advent20/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2023/12/25/advent-of-code-2023-day-20/
   2
   3 import Debug.Trace
   4
   5 import AoC
   6 import Data.Text (Text)
   7 import qualified Data.Text.IO as TIO
   8 import Data.Attoparsec.Text hiding (take)
   9 import Control.Applicative
  10 import Data.List
  11 import qualified Data.Map.Strict as M
  12 import Data.Map ((!))
  13 import qualified Data.Sequence as Q
  14 import Data.Sequence ((|>), (><), Seq( (:|>), (:<|) ) )
  15 import Control.Lens hiding (Level)
  16 import Control.Monad.State.Strict
  17 import Control.Monad.Reader
  18 import Control.Monad.Writer
  19 import Control.Monad.RWS.Strict
  20 import Data.Function (on)
  21
  22 type Name = String
  23
  24 data Level = Low | High deriving (Show, Eq, Ord)
  25 data Pulse = Pulse { _source :: Name, _level :: Level , _destination :: Name }
  26   deriving (Show, Eq, Ord)
  27 makeLenses ''Pulse
  28
  29 type Queue = Q.Seq Pulse
  30
  31 type Memory = M.Map Name Level
  32
  33 data Module =
  34   Broadcast
  35   | FlipFlop Bool
  36   | Conjunction Memory
  37   | Untyped
  38   deriving (Show, Eq, Ord)
  39
  40 type Network = M.Map Name [Name]
  41 type Modules = M.Map Name Module
  42
  43 data NetworkState = NetworkState { _modules :: Modules
  44                                  , _queue :: Queue
  45                                  }
  46   deriving (Show, Eq, Ord)
  47 makeLenses ''NetworkState
  48
  49 type NetworkHandler = RWS Network [Pulse] NetworkState
  50
  51 main :: IO ()
  52 main =
  53   do  dataFileName <- getDataFileName
  54       text <- TIO.readFile dataFileName
  55       let config = successfulParse text
  56       let (network, modules) = assembleNetwork config
  57       print $ part1 network modules
  58       print $ part2 network modules
  59
  60
  61
  62
  63 part1, part2 :: Network -> Modules -> Int
  64 part1 network modules = highs * lows
  65   where (_, (highs, lows)) =
  66          (!! 1000) $ iterate (pressAndEvaluate network part1Extractor) (state0, (0, 0))
  67         state0 = NetworkState modules Q.empty
  68 part2 network modules = foldl' lcm 1 cycleLengths
  69   where (_, lxPulses) =
  70          (!! 10000) $ iterate (pressAndEvaluate network part2Extractor) (state0, [(0, [])])
  71         state0 = NetworkState modules Q.empty
  72         lxHighs = filter (not . null . snd) lxPulses
  73         cycleLengths = fmap (fst . head) $
  74                           fmap sort $
  75                           groupBy ((==) `on` (_source . snd)) $
  76                           sortBy (compare `on` (\(_, p) -> p ^. source)) $
  77                           fmap ((\(n, ps) -> (n, head ps))) lxHighs
  78
  79 pressAndEvaluate :: Network -> (a -> [Pulse] -> a) -> (NetworkState, a) -> (NetworkState, a)
  80 pressAndEvaluate network resultExtractor (state0, result0) = (state1, result1)
  81   where (state1, pulses) = buttonPress network state0
  82         result1 = resultExtractor result0 pulses
  83
  84 part1Extractor :: (Int, Int) -> [Pulse] -> (Int, Int)
  85 part1Extractor (highs, lows) pulses = (highs + length hs, lows + length ls)
  86   where (hs, ls) = partition ((== High) . _level) pulses
  87
  88
  89 part2Extractor :: [(Int, [Pulse])] -> [Pulse] -> [(Int, [Pulse])]
  90 part2Extractor allRs@((i, _):rs) pulses = (i + 1, lxPulses) : allRs
  91   where lxPulses = filter catchLx pulses
  92         catchLx (Pulse _ High "lx") = True
  93         catchLx _ = False
  94
  95 buttonPress :: Network -> NetworkState -> (NetworkState, [Pulse])
  96 buttonPress network state =
  97   execRWS handlePulses network (state & queue .~ pulse0)
  98   where pulse0 = Q.singleton $ Pulse "button" Low "broadcaster"
  99
 100
 101 handlePulses :: NetworkHandler ()
 102 handlePulses =
 103   do  pulses <- gets _queue
 104       if Q.null pulses
 105         then return ()
 106         else
 107           do  let p :<| _ = pulses
 108               handlePulse p
 109               modify (\s -> s & queue %~ (Q.drop 1))
 110               handlePulses
 111
 112
 113 handlePulse :: Pulse -> NetworkHandler ()
 114 handlePulse p@(Pulse _ _ destination) =
 115   do  mdl <- gets ((! destination) . _modules)
 116       outGoings <- asks (! destination)
 117       let (mdl', maybeLevel) = processPulse p mdl
 118       modify (\s -> s & modules . ix destination .~ mdl')
 119       tell [p]
 120       case maybeLevel of
 121         Nothing -> return ()
 122         Just level' ->
 123           do let newPulses = fmap (Pulse destination level') outGoings
 124              modify (\s -> s & queue %~ (>< (Q.fromList newPulses)))
 125             --  tell newPulses
 126              return ()
 127
 128 processPulse :: Pulse -> Module -> (Module, Maybe Level)
 129 -- processPulse p m | trace ((show p) ++ " " ++ (show m) ) False = undefined
 130 processPulse (Pulse _ l _) Broadcast = (Broadcast, Just l)
 131 processPulse (Pulse _ Low _) (FlipFlop False) = (FlipFlop True, Just High)
 132 processPulse (Pulse _ Low _) (FlipFlop True) = (FlipFlop False, Just Low)
 133 processPulse (Pulse _ High _) (FlipFlop s) = (FlipFlop s, Nothing)
 134 processPulse (Pulse s l _) (Conjunction memory) =
 135   (Conjunction memory', Just outLevel)
 136   where memory' = M.insert s l memory
 137         outLevel = if all (== High) $ M.elems memory' then Low else High
 138 processPulse _ Untyped = (Untyped, Nothing)
 139
 140 -- Assemble the network
 141
 142 assembleNetwork :: [((Module, Name), [Name])] -> (Network, Modules)
 143 assembleNetwork config = (network, modules)
 144   where (network, modules0) = mkNetwork config
 145         modules1 = M.union (mkModules config) modules0
 146         modules = addConjunctionMemory network modules1
 147
 148 mkModules :: [((Module, Name), [Name])] -> Modules
 149 mkModules = M.fromList . fmap (\((m, n), _) -> (n, m))
 150
 151 mkNetwork :: [((Module, Name), [Name])] -> (Network, Modules)
 152 mkNetwork config = (net, mods)
 153   where net = M.fromList $ fmap (\((_, n), ds) -> (n, ds)) config
 154         mods = M.fromList $ concatMap (\(_, ds) -> fmap (\d -> (d, Untyped)) ds) config
 155
 156 addConjunctionMemory :: Network -> Modules -> Modules
 157 addConjunctionMemory network modules =
 158   M.foldlWithKey addMemory modules network
 159
 160 addMemory :: Modules -> Name -> [Name] -> Modules
 161 addMemory modules source connections =
 162   foldl' (addOneMemory source) modules connections
 163
 164 addOneMemory :: Name -> Modules -> Name -> Modules
 165 addOneMemory source modules destination =
 166   case modules ! destination of
 167     Conjunction memory -> M.insert destination (Conjunction $ M.insert source Low memory) modules
 168     _ -> modules
 169
 170 showDot :: Network -> String
 171 showDot network =
 172   "digraph {\n" ++ (concatMap showDotLine $ M.toList network) ++ "\n}"
 173   where showDotLine (source, destinations) =
 174           concatMap (\d -> source ++ " -> " ++ d ++ ";\n") destinations
 175
 176 -- Parse the input file
 177
 178 configLinesP :: Parser [((Module, Name), [Name])]
 179 configLineP :: Parser ((Module, Name), [Name])
 180 moduleP, broadcastP, flipFlopP, conjunctionP :: Parser (Module, Name)
 181 -- namesP :: Parser [Name]
 182 nameP :: Parser Name
 183
 184 configLinesP = configLineP `sepBy` endOfLine
 185 configLineP = (,) <$> (moduleP <* " -> ") <*> (nameP `sepBy` ", ")
 186
 187 moduleP = broadcastP <|> flipFlopP <|> conjunctionP
 188
 189 broadcastP = (Broadcast, "broadcaster") <$ "broadcaster"
 190 flipFlopP = (FlipFlop False, ) <$> ("%" *> nameP)
 191 conjunctionP = (Conjunction M.empty, ) <$> ("&" *> nameP)
 192
 193 -- namesP = nameP `sepBy` ", "
 194 nameP = many1 letter
 195
 196 successfulParse :: Text -> [((Module, Name), [Name])]
 197 successfulParse input =
 198   case parseOnly configLinesP input of
 199     Left  _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
 200     Right monkeys -> monkeys