Support opcode 9 (Adjust), Integer in intcode

3 files changed, 58 insertions, 43 deletions

diff --git a/2019/7.hs b/2019/7.hs
index d5910f1..40eaf96 100644
--- a/2019/7.hs
+++ b/2019/7.hs
@@ -7,13 +7,13 @@ import Input
 import IntCode
 
 
-part1 :: [Int] -> Int
+part1 :: [Integer] -> Integer
 part1 program =
     let settings = permutations [0..4]
         outcome setting = foldr (.) id [\i -> head (snd (run program [p,i])) | p <- setting] 0
     in maximum (map outcome settings)
 
-part2 :: [Int] -> Int
+part2 :: [Integer] -> Integer
 part2 program =
     let settings = permutations [5..9]
         outcome setting =
diff --git a/2019/IntCode.hs b/2019/IntCode.hs
index 699a40c..71fa3b1 100644
--- a/2019/IntCode.hs
+++ b/2019/IntCode.hs
@@ -1,6 +1,6 @@
 {-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, FlexibleContexts #-}
 module IntCode (
-    Arg(..), IC(..),
+    Arg(..), IC(..), mkAddr, mkImm, mkRel,
     parse, decode, run, unparse,
     runInterruptible, runContinue, Continuation()
 ) where
@@ -11,9 +11,14 @@ import qualified Data.Array.ST as SA
 import qualified Data.Array.IArray as IA
 
 
-data Arg = Imm Int | Addr Int
+data Arg = Addr Integer | Imm Integer | Rel Integer
   deriving (Show)
 
+mkAddr, mkImm, mkRel :: Integral i => i -> Arg
+mkAddr = Addr . fromIntegral
+mkImm = Imm . fromIntegral
+mkRel = Rel . fromIntegral
+
 data IC
     = Add Arg Arg Arg
     | Mul Arg Arg Arg
@@ -23,17 +28,18 @@ data IC
     | Jez Arg Arg
     | Clt Arg Arg Arg
     | Ceq Arg Arg Arg
+    | Adj Arg
     | Hlt
   deriving (Show)
 
-parse :: String -> [Int]
+parse :: String -> [Integer]
 parse = map read . splitOn ','
   where
     splitOn c s = case break (== c) s of
                       (pre, _ : post) -> pre : splitOn c post
                       _ -> [s]
 
-unparse :: IC -> [Int]
+unparse :: IC -> [Integer]
 unparse ic = case ic of
                  Add a b c -> go 1 [a,b,c]
                  Mul a b c -> go 2 [a,b,c]
@@ -43,16 +49,19 @@ unparse ic = case ic of
                  Jez a b   -> go 6 [a,b  ]
                  Clt a b c -> go 7 [a,b,c]
                  Ceq a b c -> go 8 [a,b,c]
+                 Adj a     -> go 9 [a    ]
                  Hlt       -> [99]
   where
     go code as = (100 * mode as + code) : map bare as
-      where bare (Imm n) = n
-            bare (Addr n) = n
+      where bare (Addr n) = n
+            bare (Imm n) = n
+            bare (Rel n) = n
     mode [] = 0
-    mode (Imm _ : as) = 1 + 10 * mode as
-    mode (Addr _ : as) = 10 * mode as
+    mode (Addr _ : as) = 0 + 10 * mode as
+    mode (Imm _ : as)  = 1 + 10 * mode as
+    mode (Rel _ : as)  = 2 + 10 * mode as
 
-decode :: [Int] -> (IC, Int)
+decode :: [Integer] -> (IC, Int)
 decode [] = error "IC: Execution fell off end of program"
 decode (ins : rest) =
     let (code, modes) = insModeSplit ins
@@ -61,10 +70,10 @@ decode (ins : rest) =
 class Monad m => ReadArray m a i e where
     readArray :: a -> i -> m e
 
-instance IA.Ix i => ReadArray (ST s) (SA.STUArray s i Int) i Int where
+instance IA.Ix i => ReadArray (ST s) (SA.STArray s i Integer) i Integer where
     readArray = SA.readArray
 
-decode' :: ReadArray m a Int Int => a -> Int -> m (IC, Int)
+decode' :: ReadArray m a Int Integer => a -> Int -> m (IC, Int)
 decode' arr ptr = do
     ins <- readArray arr ptr
     let (code, modes) = insModeSplit ins
@@ -72,16 +81,17 @@ decode' arr ptr = do
     params <- sequence [readArray arr (ptr + i) | i <- [1..nargs]]
     return (decodeBase code modes params, 1 + nargs)
 
-insModeSplit :: Int -> (Int, [Int])
+insModeSplit :: Integer -> (Integer, [Integer])
 insModeSplit n = (n `mod` 100, map (`mod` 10) (iterate (`div` 10) (n `div` 100)))
 
-insNArgs :: Int -> Int
+insNArgs :: Integer -> Int
 insNArgs n = case n of { 1 -> 3; 2 -> 3; 3 -> 1; 4 -> 1;
                          5 -> 2; 6 -> 2; 7 -> 3; 8 -> 3;
+                         9 -> 1;
                          99 -> 0;
                          _ -> error $ "IC: Unknown instruction " ++ show n }
 
-decodeBase :: Int -> [Int] -> [Int] -> IC
+decodeBase :: Integer -> [Integer] -> [Integer] -> IC
 decodeBase  1 (am:bm:cm:_) (a:b:c:_) = Add (insArg am a) (insArg bm b) (insArg cm 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)
@@ -90,61 +100,66 @@ 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  9 (am:      _) (a:    _) = Adj (insArg am a)
 decodeBase 99 (         _) (      _) = Hlt
 decodeBase ic _            _         = error $ "IC: Unknown instruction " ++ show ic
 
-insArg :: Int -> Int -> Arg
+insArg :: Integer -> Integer -> Arg
 insArg 0 n = Addr n
 insArg 1 n = Imm n
+insArg 2 n = Rel n
 insArg m n = error $ "IC: Unknown parameter mode " ++ show m ++ " for parameter " ++ show n
 
-run :: [Int] -> [Int] -> ([Int], [Int])
+run :: [Integer] -> [Integer] -> ([Integer], [Integer])
 run initMem input =
     case runInterruptible initMem input of
         Left _ -> error "IC: Not enough input"
         Right res -> res
 
-data Continuation = Continuation (A.Array Int Int) Int
+data Continuation = Continuation (A.Array Int Integer) Int Int
 
-runInterruptible :: [Int] -> [Int] -> Either (Continuation, [Int]) ([Int], [Int])
+runInterruptible :: [Integer] -> [Integer] -> Either (Continuation, [Integer]) ([Integer], [Integer])
 runInterruptible initMem input = runST $ do
     arr <- SA.newListArray (0, length initMem - 1) initMem
-    res <- runArray arr 0 input
+    res <- runArray arr 0 input 0
     case fst res of
         Left cont -> return (Left (cont, snd res))
         Right mem -> return (Right (mem, snd res))
 
-runContinue :: Continuation -> [Int] -> Either (Continuation, [Int]) ([Int], [Int])
-runContinue (Continuation frozen ptr) moreinput = runST $ do
+runContinue :: Continuation -> [Integer] -> Either (Continuation, [Integer]) ([Integer], [Integer])
+runContinue (Continuation frozen ptr offset) moreinput = runST $ do
     arr <- SA.thaw frozen
-    res <- runArray arr ptr moreinput
+    res <- runArray arr ptr moreinput offset
     case fst res of
         Left cont -> return (Left (cont, snd res))
         Right mem -> return (Right (mem, snd res))
 
-runArray :: SA.STUArray s Int Int -> Int -> [Int] -> ST s (Either Continuation [Int], [Int])
-runArray arr ptr inp = do
+runArray :: SA.STArray s Int Integer -> Int -> [Integer] -> Int -> ST s (Either Continuation [Integer], [Integer])
+runArray arr ptr inp offset = do
     (ic, len) <- decode' arr ptr
-    let continueInp inp' = runArray arr (ptr + len) inp'
+    let continueInp inp' = runArray arr (ptr + len) inp' offset
+        continueOff off' = runArray arr (ptr + len) inp (fromIntegral off')
         continue = continueInp inp
-        jump ptr' = runArray arr ptr' inp
+        jump ptr' = runArray arr (fromIntegral ptr') inp offset
     case ic of
-        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
+        Add a b (Addr c) -> (+) <$> resolve a <*> resolve b >>= SA.writeArray arr (fromIntegral c) >> continue
+        Mul a b (Addr c) -> (*) <$> resolve a <*> resolve b >>= SA.writeArray arr (fromIntegral c) >> continue
         Inp (Addr a) -> case inp of
-                     val : rest -> SA.writeArray arr a val >> continueInp rest
+                     val : rest -> SA.writeArray arr (fromIntegral a) val >> continueInp rest
                      [] -> do
                          frozen <- SA.freeze arr
-                         return (Left (Continuation frozen ptr), [])
+                         return (Left (Continuation frozen ptr offset), [])
         Out a -> resolve a >>= \val -> fmap (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
+        Clt a b (Addr c) -> (((fromIntegral . fromEnum) .) . (<)) <$> resolve a <*> resolve b >>= SA.writeArray arr (fromIntegral c) >> continue
+        Ceq a b (Addr c) -> (((fromIntegral . fromEnum) .) . (==)) <$> resolve a <*> resolve b >>= SA.writeArray arr (fromIntegral c) >> continue
+        Adj a -> resolve a >>= \add -> continueOff (fromIntegral offset + add)
         Hlt -> do
             mem <- SA.getElems arr
             return (Right mem, [])
         ins -> error $ "IC: Invalid instruction " ++ show ins
   where
+    resolve (Addr n) = SA.readArray arr (fromIntegral n)
     resolve (Imm n) = return n
-    resolve (Addr n) = SA.readArray arr n
+    resolve (Rel n) = return (fromIntegral offset + n)
diff --git a/2019/ic-asm.hs b/2019/ic-asm.hs
index f3f5540..aec6e51 100644
--- a/2019/ic-asm.hs
+++ b/2019/ic-asm.hs
@@ -25,9 +25,9 @@ data Stmt
     | Halt
   deriving (Show)
 
-assemble :: [Stmt] -> [Int]
+assemble :: [Stmt] -> [Integer]
 assemble stmts =
-    IC.unparse (IC.Jnz (IC.Imm 1) (IC.Imm (3 + nvars)))
+    IC.unparse (IC.Jnz (IC.Imm 1) (IC.mkImm (3 + nvars)))
     ++ replicate nvars 0
     ++ concatMap IC.unparse (snd (goL nvars stmts))
   where
@@ -49,25 +49,25 @@ assemble stmts =
     go off (Ceq n a b) = ([IC.Clt (ref a) (ref b) (ref (Var n))], off + 4)
     go off (If a [] s2) =
         let (off', ics2) = goL (off + 3) s2
-        in (IC.Jnz (ref a) (IC.Imm off') : ics2
+        in (IC.Jnz (ref a) (IC.mkImm off') : ics2
            , off')
     go off (If a s1 s2) =
         let (offAfterThen, ics2) = goL (off + 3) s2
             (offAfterElse, ics1) = goL (offAfterThen + 3) s1
-        in ([IC.Jez (ref a) (IC.Imm (offAfterThen + 3))]
+        in ([IC.Jez (ref a) (IC.mkImm (offAfterThen + 3))]
             ++ ics2
-            ++ [IC.Jnz (IC.Imm 1) (IC.Imm offAfterElse)]
+            ++ [IC.Jnz (IC.Imm 1) (IC.mkImm offAfterElse)]
             ++ ics1
            , offAfterElse)
     go off (While a s) =
         let (offAfterBody, ics) = goL (off + 3) s
-        in ([IC.Jez (ref a) (IC.Imm (offAfterBody + 3))]
+        in ([IC.Jez (ref a) (IC.mkImm (offAfterBody + 3))]
             ++ ics
-            ++ [IC.Jnz (ref a) (IC.Imm (off + 3))]
+            ++ [IC.Jnz (ref a) (IC.mkImm (off + 3))]
            , offAfterBody + 3)
     go off Halt = ([IC.Hlt], off + 1)
 
-    ref (Imm n) = IC.Imm n
+    ref (Imm n) = IC.mkImm n
     ref (Var n) = case Map.lookup n var2idx of
                       Just idx -> IC.Addr idx
                       Nothing  -> error $ "Undeclared variable '" ++ n ++ "'"