From 50040e9ef533a005d267a3527782922113a6f541 Mon Sep 17 00:00:00 2001
From: Eric Mertens <emertens@galois.com>
Date: Tue, 17 Dec 2024 09:52:22 -0800
Subject: [PATCH] generalize

---
 solutions/src/2024/17.hs | 40 +++++++++++++++++++++++++++++-----------
 1 file changed, 29 insertions(+), 11 deletions(-)

diff --git a/solutions/src/2024/17.hs b/solutions/src/2024/17.hs
index 12e9783..c48ec39 100644
--- a/solutions/src/2024/17.hs
+++ b/solutions/src/2024/17.hs
@@ -16,7 +16,7 @@ import Data.List (intercalate)
 import Data.SBV
     (Word64, SWord64, SBool,
      (.==), (./=), (.&.), (.&&), sShiftRight, shiftR, xor,
-     optLexicographic, free, minimize, constrain, getModelValue)
+     optLexicographic, free, minimize, constrain, getModelValue, sFalse, sTrue, (.||))
 
 -- | >>> :main
 -- 2,7,4,7,2,1,7,5,1
@@ -31,11 +31,11 @@ main =
       Program: %u&,%n|]
 
     putStrLn (intercalate "," (map show (run (Machine a b c) program)))
-
+    
     res <- optLexicographic
       do a2 <- free "a"
          minimize "smallest" a2
-         constrain (part2 a2 program)
+         constrain (run2 (SMachine program a2 0 0) program)
     case getModelValue "a" res of
       Just x -> print (x :: Word64)
       Nothing -> fail "no solution"
@@ -48,12 +48,12 @@ run m0 pgm = go m0 pgm
     go m = \case
       0 : x : ip' -> go m{ rA = rA m `shiftR` combo x } ip'
       1 : x : ip' -> go m{ rB = rB m `xor`          x } ip'
-      2 : x : ip' -> go m{ rB = combo x .&. 7         } ip'
-      3 : x : ip' -> go m (if rA m == 0 then ip' else drop x pgm)
+      2 : x : ip' -> go m{ rB = 7    .&.      combo x } ip'
       4 : _ : ip' -> go m{ rB = rB m `xor`    rC m    } ip'
-      5 : x : ip' -> combo x .&. 7 : go m ip'
       6 : x : ip' -> go m{ rB = rA m `shiftR` combo x } ip'
       7 : x : ip' -> go m{ rC = rA m `shiftR` combo x } ip'
+      3 : x : ip' -> go m (if rA m == 0 then ip' else drop x pgm)
+      5 : x : ip' -> combo x .&. 7 : go m ip'
       _           -> []
       where
         combo = \case
@@ -61,9 +61,27 @@ run m0 pgm = go m0 pgm
           4 -> rA m; 5 -> rB m; 6 -> rC m
           _ -> error "invalid combo operand"
 
-part2 :: SWord64 -> [Int] -> SBool
-part2 a []     = a .== 0
-part2 a (o:os) = a ./= 0 .&& b2 .== fromIntegral o .&& part2 (a `shiftR` 3) os
+data SMachine = SMachine { outs :: [Int], sA, sB, sC :: SWord64 }
+
+run2 :: SMachine -> [Int] -> SBool
+run2 m0 pgm = go m0 pgm
   where
-    b1 = a .&. 7 `xor` 2
-    b2 = (a `sShiftRight` b1 `xor` b1 `xor` 3) .&. 7
+    go m = \case
+      0 : x : ip' -> go m{ sA = sA m `sShiftRight` combo x } ip'
+      1 : x : ip' -> go m{ sB = sB m `xor`  fromIntegral x } ip'
+      2 : x : ip' -> go m{ sB = 7    .&.      combo x } ip'
+      4 : _ : ip' -> go m{ sB = sB m `xor`    sC m    } ip'
+      6 : x : ip' -> go m{ sB = sA m `sShiftRight` combo x } ip'
+      7 : x : ip' -> go m{ sC = sA m `sShiftRight` combo x } ip'
+      3 : x : ip' -> (sA m .== 0) .&& go m ip' .||
+                     (sA m ./= 0) .&& go m (drop x pgm)
+      5 : x : ip' ->
+        case outs m of
+          [] -> sFalse
+          o:os -> combo x .&. 7 .== fromIntegral o .&& go m{ outs = os} ip'
+      _           -> if null (outs m) then sTrue else sFalse
+      where
+        combo = \case
+          0 -> 0; 1 -> 1; 2 -> 2; 3 -> 3
+          4 -> sA m; 5 -> sB m; 6 -> sC m
+          _ -> error "invalid combo operand"