cleanup

solutions/src/2023/20.hs

@@ -32,6 +32,7 @@ data Node
   = Broadcast                      [String] -- ^ broadcast node
   | FlipFlop !Bool                 [String] -- ^ flip-flop
   | Conjunction !Int !(Set String) [String] -- ^ conjunction gate
+  deriving (Show)
 
 data Stream a = a :| Stream a
 
@@ -48,57 +49,68 @@ main =
     let incoming = Map.fromListWith (++) [(k, [v]) | (_, v, ks) <- input, k <- ks]
     let nodes = Map.fromList [(name, node incoming name kind conns) | (kind, name, conns) <- input]
 
-    print (part1 0 0 0 (sim nodes))
+    print (part1 (sim nodes))
     print (part2 incoming (sim nodes))
 
 node :: Map String [String] -> String -> K -> [String] -> Node
 node incoming name = \case
   K           -> Broadcast
-  K_AMPERSAND -> Conjunction (length (incoming Map.! name)) Set.empty
+  K_AMPERSAND -> Conjunction (length (Map.findWithDefault [] name incoming)) Set.empty
   K_PERCENT   -> FlipFlop False
 
-part1 :: Int -> Int -> Int -> Stream (String, a, Bool) -> Int
-part1 n l h ((src,_,sig) :| xs)
-  | n == 1000, src == "button" = l * h
-  | otherwise =
-    part1 (if src == "button" then n+1 else n)
-          (if sig then l else l+1)
-          (if sig then h+1 else h) xs
+part1 :: Stream (String, a, Bool) -> Int
+part1 = go 0 0 0
+  where
+    go n l h ((src,_,sig) :| xs)
+      | n == 1000, src == "button" = l * h
+      | otherwise =
+        go (if src == "button" then n+1 else n)
+              (if sig then l else l+1)
+              (if sig then h+1 else h) xs
 
 part2 :: Map String [String] -> Stream (String, String, Bool) -> Int
-part2 incoming msgs = foldl1 lcm [buttonsFor 0 gate msgs | gate <- incoming Map.! specialConj]
+part2 incoming = foldl lcm 1 . go 0 (Set.fromList (Map.findWithDefault [] conj incoming))
   where
-    [specialConj] = incoming Map.! "rx"
+    [conj] = incoming Map.! "rx"
 
-    buttonsFor n gate ((src, dst, msg) :| xs)
-      | "button" == src = buttonsFor (n+1) gate xs
-      | msg, src == gate, dst == specialConj = n
-    buttonsFor n gate (_ :| xs) = buttonsFor n gate xs
+    go n gates ((src, dst, msg) :| xs)
+      | Set.null gates = []
+      | "button" == src = go (n + 1) gates xs
+      | msg, dst == conj, Set.member src gates = n : go n (Set.delete src gates) xs
+      | otherwise = go n gates xs
 
+-- | Generate a stream of messages generated by this network.
+-- The button is automatically pressed any time the network
+-- stalls.
 sim :: Map String Node -> Stream (String, String, Bool)
-sim fwd = go fwd Queue.Empty
+sim = go Queue.Empty
   where
-    go st (x Queue.:<| q') = dispatch st x q'
-    go st q = dispatch st ("button", "broadcaster", False) q
-
-    dispatch st (src, dst, msg) q' =
+    go ((src, dst, msg) Queue.:<| q) = dispatch src dst msg q
+    go q                             = dispatch "button" "broadcaster" False q
+
+    -- Dispatch msg from src to dst
+    dispatch src dst msg q st =
       (src, dst, msg) :|
       case Map.lookup dst st of
-        Just (Broadcast next) -> continue st msg next -- forward message
+
+        -- broadcast: just forward message to nexts
+        Just (Broadcast next) -> send msg next st
+
+        -- flipflop: on low, toggle state and send to nexts
         Just (FlipFlop mode next)
-          | not msg -> continue st' out next -- was on sends low
+          | not msg -> send out next st' -- was on sends low
           where
             st' = Map.insert dst (FlipFlop out next) st
             out = not mode
-        Just (Conjunction sz inc next) -> continue st' out next
+
+        -- conjunction: remember incoming value, transmit nand
+        Just (Conjunction sz inc next) -> send out next st'
           where
-            inc' = mark src msg inc
-            st' = Map.insert dst (Conjunction sz inc' next) st
-            out = sz /= length inc'
-        _ -> go st q' -- ignored
-      where
-        continue st' msg' next = go st' (Queue.appendList q' [(dst, t, msg') | t <- next])
+            inc' = (if msg then Set.insert else Set.delete) src inc
+            st'  = Map.insert dst (Conjunction sz inc' next) st
+            out  = sz /= length inc'
 
-mark :: Ord a => a -> Bool -> Set a -> Set a
-mark key True = Set.insert key
-mark key False = Set.delete key
+        -- output node or flipflop on high: ignored
+        _ -> go q st
+      where
+        send msg' next = go (Queue.appendList q [(dst, t, msg') | t <- next])