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) . sort) $
  74                           groupBy ((==) `on` (_source . snd)) $
  75                           sortBy (compare `on` (\(_, p) -> p ^. source)) $
  76                           fmap (\(n, ps) -> (n, head ps)) lxHighs
  77
  78 pressAndEvaluate :: Network -> (a -> [Pulse] -> a) -> (NetworkState, a) -> (NetworkState, a)
  79 pressAndEvaluate network resultExtractor (state0, result0) = (state1, result1)
  80   where (state1, pulses) = buttonPress network state0
  81         result1 = resultExtractor result0 pulses
  82
  83 part1Extractor :: (Int, Int) -> [Pulse] -> (Int, Int)
  84 part1Extractor (highs, lows) pulses = (highs + length hs, lows + length ls)
  85   where (hs, ls) = partition ((== High) . _level) pulses
  86
  87 part2Extractor :: [(Int, [Pulse])] -> [Pulse] -> [(Int, [Pulse])]
  88 part2Extractor allRs@((i, _):rs) pulses = (i + 1, lxPulses) : allRs
  89   where lxPulses = filter catchLx pulses
  90         catchLx (Pulse _ High "lx") = True
  91         catchLx _ = False
  92
  93 buttonPress :: Network -> NetworkState -> (NetworkState, [Pulse])
  94 buttonPress network state =
  95   execRWS handlePulses network (state & queue .~ pulse0)
  96   where pulse0 = Q.singleton $ Pulse "button" Low "broadcaster"
  97
  98
  99 handlePulses :: NetworkHandler ()
 100 handlePulses =
 101   do  pulses <- gets _queue
 102       case pulses of
 103         Q.Empty -> return ()
 104         (p :<| ps) ->
 105           do modify (\s -> s & queue .~ ps)
 106              handlePulse p
 107              handlePulses
 108
 109 handlePulse :: Pulse -> NetworkHandler ()
 110 handlePulse p@(Pulse _ _ destination) =
 111   do  mdl <- gets ((! destination) . _modules)
 112       outGoings <- asks (! destination)
 113       let (mdl', maybeLevel) = processPulse p mdl
 114       modify (\s -> s & modules . ix destination .~ mdl')
 115       tell [p]
 116       case maybeLevel of
 117         Nothing -> return ()
 118         Just level' ->
 119           do let newPulses = fmap (Pulse destination level') outGoings
 120              modify (\s -> s & queue %~ (>< (Q.fromList newPulses)))
 121
 122 processPulse :: Pulse -> Module -> (Module, Maybe Level)
 123 -- processPulse p m | trace ((show p) ++ " " ++ (show m) ) False = undefined
 124 processPulse (Pulse _ l _) Broadcast = (Broadcast, Just l)
 125 processPulse (Pulse _ Low _) (FlipFlop False) = (FlipFlop True, Just High)
 126 processPulse (Pulse _ Low _) (FlipFlop True) = (FlipFlop False, Just Low)
 127 processPulse (Pulse _ High _) (FlipFlop s) = (FlipFlop s, Nothing)
 128 processPulse (Pulse s l _) (Conjunction memory) =
 129   (Conjunction memory', Just outLevel)
 130   where memory' = M.insert s l memory
 131         outLevel = if all (== High) $ M.elems memory' then Low else High
 132 processPulse _ Untyped = (Untyped, Nothing)
 133
 134 -- Assemble the network
 135
 136 assembleNetwork :: [((Module, Name), [Name])] -> (Network, Modules)
 137 assembleNetwork config = (network, modules)
 138   where (network, modules0) = mkNetwork config
 139         modules1 = M.union (mkModules config) modules0
 140         modules = addConjunctionMemory network modules1
 141
 142 mkNetwork :: [((Module, Name), [Name])] -> (Network, Modules)
 143 mkNetwork config = (net, mods)
 144   where net = M.fromList $ fmap (\((_, n), ds) -> (n, ds)) config
 145         mods = M.fromList $ concatMap (\(_, ds) -> fmap (\d -> (d, Untyped)) ds) config
 146
 147 mkModules :: [((Module, Name), [Name])] -> Modules
 148 mkModules = M.fromList . fmap (\((m, n), _) -> (n, m))
 149
 150 addConjunctionMemory :: Network -> Modules -> Modules
 151 addConjunctionMemory network modules =
 152   M.foldlWithKey addMemory modules network
 153
 154 addMemory :: Modules -> Name -> [Name] -> Modules
 155 addMemory modules source connections =
 156   foldl' (addOneMemory source) modules connections
 157
 158 addOneMemory :: Name -> Modules -> Name -> Modules
 159 addOneMemory source modules destination =
 160   case modules ! destination of
 161     Conjunction memory ->
 162       M.insert destination
 163                (Conjunction $ M.insert source Low memory)
 164                modules
 165     _ -> modules
 166
 167 showDot :: Network -> String
 168 showDot network =
 169   "digraph {\n" ++ (concatMap showDotLine $ M.toList network) ++ "\n}"
 170   where showDotLine (source, destinations) =
 171           concatMap (\d -> source ++ " -> " ++ d ++ ";\n") destinations
 172
 173 -- Parse the input file
 174
 175 configLinesP :: Parser [((Module, Name), [Name])]
 176 configLineP :: Parser ((Module, Name), [Name])
 177 moduleP, broadcastP, flipFlopP, conjunctionP :: Parser (Module, Name)
 178 -- namesP :: Parser [Name]
 179 nameP :: Parser Name
 180
 181 configLinesP = configLineP `sepBy` endOfLine
 182 configLineP = (,) <$> (moduleP <* " -> ") <*> (nameP `sepBy` ", ")
 183
 184 moduleP = broadcastP <|> flipFlopP <|> conjunctionP
 185
 186 broadcastP = (Broadcast, "broadcaster") <$ "broadcaster"
 187 flipFlopP = (FlipFlop False, ) <$> ("%" *> nameP)
 188 conjunctionP = (Conjunction M.empty, ) <$> ("&" *> nameP)
 189
 190 -- namesP = nameP `sepBy` ", "
 191 nameP = many1 letter
 192
 193 successfulParse :: Text -> [((Module, Name), [Name])]
 194 successfulParse input =
 195   case parseOnly configLinesP input of
 196     Left  _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
 197     Right monkeys -> monkeys