{-# Language ImportQualifiedPost, BlockArguments #-}
{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}
module Advent.Prelude where
import Control.Applicative (Alternative(empty))
import Control.Monad.Trans.State (StateT(StateT, runStateT))
import Data.Array.Unboxed qualified as A
import Data.Array.Base qualified as AB
import GHC.Arr qualified as GA
import Data.Coerce (coerce)
import Data.Foldable (toList)
import Data.IntMap (IntMap)
import Data.IntMap qualified as IntMap
import Data.List (foldl', scanl', inits, sortBy, tails, mapAccumL)
import Data.Map (Map)
import Data.Map.Strict qualified as SMap
import Data.Ord (comparing)
import Data.Set (Set)
import Data.Set qualified as Set
import GHC.Stack (HasCallStack)
import Language.Haskell.TH (DecsQ, stringE)
import Language.Haskell.TH.Quote (QuasiQuoter(..))
count :: (Foldable f, Eq a) => a -> f a -> Int
count :: forall (f :: * -> *) a. (Foldable f, Eq a) => a -> f a -> Int
count = (a -> Bool) -> f a -> Int
forall (f :: * -> *) a. Foldable f => (a -> Bool) -> f a -> Int
countBy ((a -> Bool) -> f a -> Int) -> (a -> a -> Bool) -> a -> f a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(==)
countBy :: Foldable f => (a -> Bool) -> f a -> Int
countBy :: forall (f :: * -> *) a. Foldable f => (a -> Bool) -> f a -> Int
countBy a -> Bool
p = (Int -> a -> Int) -> Int -> f a -> Int
forall b a. (b -> a -> b) -> b -> f a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Int
acc a
x -> if a -> Bool
p a
x then Int
accInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1 else Int
acc) Int
0
same :: Foldable t => Eq a => t a -> Bool
same :: forall (t :: * -> *) a. (Foldable t, Eq a) => t a -> Bool
same t a
x =
case t a -> [a]
forall a. t a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList t a
x of
[] -> Bool
True
a
y:[a]
ys -> (a -> Bool) -> [a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (a
y a -> a -> Bool
forall a. Eq a => a -> a -> Bool
==) [a]
ys
pickOne :: [a] -> [(a, [a])]
pickOne :: forall a. [a] -> [(a, [a])]
pickOne [a]
xs = [ (a
x, [a]
l[a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++[a]
r) | ([a]
l,a
x:[a]
r) <- [[a]] -> [[a]] -> [([a], [a])]
forall a b. [a] -> [b] -> [(a, b)]
zip ([a] -> [[a]]
forall a. [a] -> [[a]]
inits [a]
xs) ([a] -> [[a]]
forall a. [a] -> [[a]]
tails [a]
xs) ]
ordNub :: Ord a => [a] -> [a]
ordNub :: forall a. Ord a => [a] -> [a]
ordNub [a]
xs = (a -> (Set a -> [a]) -> Set a -> [a])
-> (Set a -> [a]) -> [a] -> Set a -> [a]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> (Set a -> [a]) -> Set a -> [a]
forall {a}. Ord a => a -> (Set a -> [a]) -> Set a -> [a]
f ([a] -> Set a -> [a]
forall a b. a -> b -> a
const []) [a]
xs Set a
forall a. Set a
Set.empty
where
f :: a -> (Set a -> [a]) -> Set a -> [a]
f a
x Set a -> [a]
rec Set a
seen =
case Set a -> [a]
rec (Set a -> [a]) -> (Bool, Set a) -> (Bool, [a])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Bool -> (Bool, Bool)) -> a -> Set a -> (Bool, Set a)
forall a (f :: * -> *).
(Ord a, Functor f) =>
(Bool -> f Bool) -> a -> Set a -> f (Set a)
Set.alterF (\Bool
old -> (Bool
old, Bool
True)) a
x Set a
seen of
(Bool
True, [a]
ys) -> [a]
ys
(Bool
False, [a]
ys) -> a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
ys
minimumMaybe :: Ord a => [a] -> Maybe a
minimumMaybe :: forall a. Ord a => [a] -> Maybe a
minimumMaybe [a]
xs
| [a] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
xs = Maybe a
forall a. Maybe a
Nothing
| Bool
otherwise = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$! [a] -> a
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [a]
xs
maximumMaybe :: (Foldable f, Ord a) => f a -> Maybe a
maximumMaybe :: forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Maybe a
maximumMaybe f a
xs
| f a -> Bool
forall a. f a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null f a
xs = Maybe a
forall a. Maybe a
Nothing
| Bool
otherwise = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$! f a -> a
forall a. Ord a => f a -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum f a
xs
counts :: (Foldable f, Ord a) => f a -> Map a Int
counts :: forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Map a Int
counts f a
xs = (Int -> Int -> Int) -> [(a, Int)] -> Map a Int
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
SMap.fromListWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) [(a
x,Int
1) | a
x <- f a -> [a]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList f a
xs]
compose :: [a -> a] -> a -> a
compose :: forall a. [a -> a] -> a -> a
compose = ((a -> a) -> (a -> a) -> a -> a) -> (a -> a) -> [a -> a] -> a -> a
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (a -> a) -> (a -> a) -> a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
(.) a -> a
forall a. a -> a
id
chunks :: Int -> [a] -> [[a]]
chunks :: forall a. Int -> [a] -> [[a]]
chunks Int
_ [] = []
chunks Int
n [a]
xs =
case Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n [a]
xs of
([a]
a,[a]
b) -> [a]
a [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
chunks Int
n [a]
b
löb :: Functor f => f (f a -> a) -> f a
löb :: forall (f :: * -> *) a. Functor f => f (f a -> a) -> f a
löb = (((f a -> a) -> a) -> f (f a -> a) -> f a) -> f (f a -> a) -> f a
forall a b c. (((a -> b) -> b) -> c -> a) -> c -> a
möb ((f a -> a) -> a) -> f (f a -> a) -> f a
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
möb :: (((a -> b) -> b) -> c -> a) -> c -> a
möb :: forall a b c. (((a -> b) -> b) -> c -> a) -> c -> a
möb ((a -> b) -> b) -> c -> a
f = \c
x -> let go :: a
go = ((a -> b) -> b) -> c -> a
f ((a -> b) -> a -> b
forall a b. (a -> b) -> a -> b
$ a
go) c
x in a
go
arrIx :: (A.IArray a e, A.Ix i, Alternative f) => a i e -> i -> f e
arrIx :: forall (a :: * -> * -> *) e i (f :: * -> *).
(IArray a e, Ix i, Alternative f) =>
a i e -> i -> f e
arrIx a i e
a i
i
| (i, i) -> i -> Bool
forall a. Ix a => (a, a) -> a -> Bool
A.inRange (i, i)
b i
i = e -> f e
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (e -> f e) -> e -> f e
forall a b. (a -> b) -> a -> b
$! a i e -> Int -> e
forall i. Ix i => a i e -> Int -> e
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> Int -> e
AB.unsafeAt a i e
a ((i, i) -> i -> Int
forall a. Ix a => (a, a) -> a -> Int
GA.unsafeIndex (i, i)
b i
i)
| Bool
otherwise = f e
forall a. f a
forall (f :: * -> *) a. Alternative f => f a
empty
where b :: (i, i)
b = a i e -> (i, i)
forall i. Ix i => a i e -> (i, i)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
A.bounds a i e
a
{-# Inline arrIx #-}
times :: Int -> (a -> a) -> a -> a
times :: forall a. Int -> (a -> a) -> a -> a
times Int
n a -> a
f a
x
| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = a
x
| Bool
otherwise = Int -> (a -> a) -> a -> a
forall a. Int -> (a -> a) -> a -> a
times (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) a -> a
f (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$! a -> a
f a
x
{-# INLINE times #-}
timesM :: Monad m => Int -> (a -> m a) -> a -> m a
timesM :: forall (m :: * -> *) a. Monad m => Int -> (a -> m a) -> a -> m a
timesM Int
n a -> m a
f a
x
| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x
| Bool
otherwise = Int -> (a -> m a) -> a -> m a
forall (m :: * -> *) a. Monad m => Int -> (a -> m a) -> a -> m a
timesM (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) a -> m a
f (a -> m a) -> m a -> m a
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< a -> m a
f a
x
{-# INLINE timesM #-}
uniqueAssignment ::
(Traversable t, Ord a) =>
t (Set a) ->
[t a]
uniqueAssignment :: forall (t :: * -> *) a.
(Traversable t, Ord a) =>
t (Set a) -> [t a]
uniqueAssignment t (Set a)
m =
[ ([a], t a) -> t a
forall a b. (a, b) -> b
snd (([a] -> Set a -> ([a], a)) -> [a] -> t (Set a) -> ([a], t a)
forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL (\(a
x:[a]
xs) Set a
_ -> ([a]
xs,a
x)) (IntMap a -> [a]
forall a. IntMap a -> [a]
IntMap.elems IntMap a
a) t (Set a)
m)
| IntMap a
a <- IntMap a -> [(Int, Set a)] -> [IntMap a]
forall a. Ord a => IntMap a -> [(Int, Set a)] -> [IntMap a]
go IntMap a
forall a. IntMap a
IntMap.empty ([Int] -> [Set a] -> [(Int, Set a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] (t (Set a) -> [Set a]
forall a. t a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList t (Set a)
m))]
where
go :: Ord a => IntMap a -> [(Int, Set a)] -> [IntMap a]
go :: forall a. Ord a => IntMap a -> [(Int, Set a)] -> [IntMap a]
go IntMap a
a [(Int, Set a)]
xs =
case ((Int, Set a) -> (Int, Set a) -> Ordering)
-> [(Int, Set a)] -> [(Int, Set a)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (((Int, Set a) -> Int) -> (Int, Set a) -> (Int, Set a) -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (Set a -> Int
forall a. Set a -> Int
Set.size (Set a -> Int) -> ((Int, Set a) -> Set a) -> (Int, Set a) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, Set a) -> Set a
forall a b. (a, b) -> b
snd)) [(Int, Set a)]
xs of
[] -> [IntMap a
a]
(Int
k,Set a
vs):[(Int, Set a)]
rest ->
do a
v <- Set a -> [a]
forall a. Set a -> [a]
Set.toList Set a
vs
IntMap a -> [(Int, Set a)] -> [IntMap a]
forall a. Ord a => IntMap a -> [(Int, Set a)] -> [IntMap a]
go (Int -> a -> IntMap a -> IntMap a
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
k a
v IntMap a
a) ((Set a -> Set a) -> (Int, Set a) -> (Int, Set a)
forall a b. (a -> b) -> (Int, a) -> (Int, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.delete a
v) ((Int, Set a) -> (Int, Set a)) -> [(Int, Set a)] -> [(Int, Set a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Int, Set a)]
rest)
fromDigits :: HasCallStack => Integral a => a -> [a] -> a
fromDigits :: forall a. (HasCallStack, Integral a) => a -> [a] -> a
fromDigits a
base
| a
base a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
2 = [Char] -> [a] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"fromDigits: bad base"
| Bool
otherwise = (a -> a -> a) -> a -> [a] -> a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\a
acc a
x -> a
acc a -> a -> a
forall a. Num a => a -> a -> a
* a
base a -> a -> a
forall a. Num a => a -> a -> a
+ a
x) a
0
toDigits :: HasCallStack => Integral a => a -> a -> [a]
toDigits :: forall a. (HasCallStack, Integral a) => a -> a -> [a]
toDigits a
base a
x
| a
base a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
2 = [Char] -> [a]
forall a. HasCallStack => [Char] -> a
error [Char]
"toDigits: bad base"
| a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0 = [Char] -> [a]
forall a. HasCallStack => [Char] -> a
error [Char]
"toDigits: negative number"
| Bool
otherwise = [a] -> a -> [a]
go [] a
x
where
go :: [a] -> a -> [a]
go [a]
xs a
0 = [a]
xs
go [a]
xs a
n = case a -> a -> (a, a)
forall a. Integral a => a -> a -> (a, a)
quotRem a
n a
base of
(a
n', a
digit) -> [a] -> a -> [a]
go (a
digita -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
xs) a
n'
power :: HasCallStack => (a -> a -> a) -> a -> Integer -> a
power :: forall a. HasCallStack => (a -> a -> a) -> a -> Integer -> a
power a -> a -> a
(#) a
one Integer
n
| Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
1 = [Char] -> a
forall a. HasCallStack => [Char] -> a
error ([Char]
"power: bad argument " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Integer -> [Char]
forall a. Show a => a -> [Char]
show Integer
n)
| Bool
otherwise = Integer -> a
forall {t}. Integral t => t -> a
go Integer
n
where
double :: a -> a
double a
x = a
x a -> a -> a
# a
x
go :: t -> a
go t
i
| t
1 t -> t -> Bool
forall a. Eq a => a -> a -> Bool
== t
i = a
one
| t -> Bool
forall a. Integral a => a -> Bool
even t
i = a -> a
double (t -> a
go (t
i t -> t -> t
forall a. Integral a => a -> a -> a
`quot` t
2))
| Bool
otherwise = a -> a
double (t -> a
go (t
i t -> t -> t
forall a. Integral a => a -> a -> a
`quot` t
2)) a -> a -> a
# a
one
scanlM :: (Traversable t, Monad m) => (b -> a -> m (c, a)) -> a -> t b -> m (t c, a)
scanlM :: forall (t :: * -> *) (m :: * -> *) b a c.
(Traversable t, Monad m) =>
(b -> a -> m (c, a)) -> a -> t b -> m (t c, a)
scanlM b -> a -> m (c, a)
f a
z t b
t = StateT a m (t c) -> a -> m (t c, a)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT ((b -> StateT a m c) -> t b -> StateT a m (t c)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> t a -> f (t b)
traverse ((b -> a -> m (c, a)) -> b -> StateT a m c
forall a b. Coercible a b => a -> b
coerce b -> a -> m (c, a)
f) t b
t) a
z
stageTH :: DecsQ
stageTH :: DecsQ
stageTH = [Dec] -> DecsQ
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
multiline :: QuasiQuoter
multiline :: QuasiQuoter
multiline = QuasiQuoter {
quoteExp :: [Char] -> Q Exp
quoteExp = [Char] -> Q Exp
forall {m :: * -> *}. (MonadFail m, Quote m) => [Char] -> m Exp
trimLeadingWhitespace,
quoteType :: [Char] -> Q Type
quoteType = Q Type -> [Char] -> Q Type
forall a b. a -> b -> a
const ([Char] -> Q Type
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
"multiline doesn't do types"),
quoteDec :: [Char] -> DecsQ
quoteDec = DecsQ -> [Char] -> DecsQ
forall a b. a -> b -> a
const ([Char] -> DecsQ
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
"multiline doesn't do decs"),
quotePat :: [Char] -> Q Pat
quotePat = Q Pat -> [Char] -> Q Pat
forall a b. a -> b -> a
const ([Char] -> Q Pat
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
"multiline doesn't do patterns")
}
where
trimLeadingWhitespace :: [Char] -> m Exp
trimLeadingWhitespace [Char]
str =
case [Char] -> [[Char]]
lines [Char]
str of
[] -> [Char] -> m Exp
forall a. [Char] -> m a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
"empty multiline string??"
[Char]
x:[[Char]]
xs
| Bool -> Bool
not ([Char] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Char]
x) -> [Char] -> m Exp
forall a. [Char] -> m a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
"stuff on first line??"
| [[Char]] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[Char]]
xs -> [Char] -> m Exp
forall a. [Char] -> m a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
"empty multiline string??"
| Bool
otherwise -> [Char] -> m Exp
forall (m :: * -> *). Quote m => [Char] -> m Exp
stringE ([[Char]] -> [Char]
unlines (([Char] -> [Char]) -> [[Char]] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> [Char] -> [Char]
forall a. Int -> [a] -> [a]
drop Int
n) [[Char]]
xs))
where
n :: Int
n = [Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum (([Char] -> Int) -> [[Char]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ([Char] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Char] -> Int) -> ([Char] -> [Char]) -> [Char] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> [Char] -> [Char]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Char
' 'Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
==)) [[Char]]
xs)
partialSums :: Num a => [a] -> [a]
partialSums :: forall a. Num a => [a] -> [a]
partialSums = (a -> a -> a) -> a -> [a] -> [a]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl' a -> a -> a
forall a. Num a => a -> a -> a
(+) a
0
{-# Inline partialSums #-}
binSearchLargest ::
(Int -> Bool) ->
Int ->
Int ->
Int
binSearchLargest :: (Int -> Bool) -> Int -> Int -> Int
binSearchLargest Int -> Bool
p Int
lo Int
hi
| Int
lo Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
hi = Int
lo
| Int -> Bool
p Int
mid = (Int -> Bool) -> Int -> Int -> Int
binSearchLargest Int -> Bool
p Int
mid Int
hi
| Bool
otherwise = (Int -> Bool) -> Int -> Int -> Int
binSearchLargest Int -> Bool
p Int
lo Int
mid
where
mid :: Int
mid = Int
lo Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
hi Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
lo) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2