Day 5

1 files changed, 22 insertions, 6 deletions

diff --git a/2019/IntCode.hs b/2019/IntCode.hs
index 904e961..89b54a2 100644
--- a/2019/IntCode.hs
+++ b/2019/IntCode.hs
@@ -17,6 +17,10 @@ data IC
     | Mul Arg Arg Arg
     | Inp Arg
     | Out Arg
+    | Jnz Arg Arg
+    | Jez Arg Arg
+    | Clt Arg Arg Arg
+    | Ceq Arg Arg Arg
     | Hlt
   deriving (Show)
 
@@ -51,7 +55,9 @@ insModeSplit :: Int -> (Int, [Int])
 insModeSplit n = (n `mod` 100, map (`mod` 10) (iterate (`div` 10) (n `div` 100)))
 
 insNArgs :: Int -> Int
-insNArgs n = case n of { 1 -> 3; 2 -> 3; 3 -> 1; 4 -> 1; 99 -> 0;
+insNArgs n = case n of { 1 -> 3; 2 -> 3; 3 -> 1; 4 -> 1;
+                         5 -> 2; 6 -> 2; 7 -> 3; 8 -> 3;
+                         99 -> 0;
                          _ -> error $ "IC: Unknown instruction " ++ show n }
 
 decodeBase :: Int -> [Int] -> [Int] -> IC
@@ -59,6 +65,10 @@ decodeBase  1 (am:bm:cm:_) (a:b:c:_) = Add (insArg am a) (insArg bm b) (insArg c
 decodeBase  2 (am:bm:cm:_) (a:b:c:_) = Mul (insArg am a) (insArg bm b) (insArg cm c)
 decodeBase  3 (am:      _) (a:    _) = Inp (insArg am a)
 decodeBase  4 (am:      _) (a:    _) = Out (insArg am a)
+decodeBase  5 (am:bm:   _) (a:b:  _) = Jnz (insArg am a) (insArg bm b)
+decodeBase  6 (am:bm:   _) (a:b:  _) = Jez (insArg am a) (insArg bm b)
+decodeBase  7 (am:bm:cm:_) (a:b:c:_) = Clt (insArg am a) (insArg bm b) (insArg cm c)
+decodeBase  8 (am:bm:cm:_) (a:b:c:_) = Ceq (insArg am a) (insArg bm b) (insArg cm c)
 decodeBase 99 (         _) (      _) = Hlt
 decodeBase ic _            _         = error $ "IC: Unknown instruction " ++ show ic
 
@@ -77,14 +87,20 @@ run initMem input = runST $ do
     run' :: SA.STUArray s Int Int -> Int -> [Int] -> ST s [Int]
     run' arr ptr inp = do
         (ic, len) <- decode' arr ptr
-        let continue i' = run' arr (ptr + len) i'
+        let continueInp inp' = run' arr (ptr + len) inp'
+            continue = continueInp inp
+            jump ptr' = run' arr ptr' inp
         case ic of
-            Add a b (Addr c) -> (+) <$> resolve a <*> resolve b >>= SA.writeArray arr c >> continue inp
-            Mul a b (Addr c) -> (*) <$> resolve a <*> resolve b >>= SA.writeArray arr c >> continue inp
+            Add a b (Addr c) -> (+) <$> resolve a <*> resolve b >>= SA.writeArray arr c >> continue
+            Mul a b (Addr c) -> (*) <$> resolve a <*> resolve b >>= SA.writeArray arr c >> continue
             Inp (Addr a) -> case inp of
-                         val : rest -> SA.writeArray arr a val >> continue rest
+                         val : rest -> SA.writeArray arr a val >> continueInp rest
                          _ -> error "IC: Not enough input"
-            Out a -> resolve a >>= \val -> (val :) <$> continue inp
+            Out a -> resolve a >>= \val -> (val :) <$> continue
+            Jnz a b -> resolve a >>= \val -> if val /= 0 then resolve b >>= jump else continue
+            Jez a b -> resolve a >>= \val -> if val == 0 then resolve b >>= jump else continue
+            Clt a b (Addr c) -> ((fromEnum .) . (<)) <$> resolve a <*> resolve b >>= SA.writeArray arr c >> continue
+            Ceq a b (Addr c) -> ((fromEnum .) . (==)) <$> resolve a <*> resolve b >>= SA.writeArray arr c >> continue
             Hlt -> return []
             ins -> error $ "IC: Invalid instruction " ++ show ins
       where