Day 2 in Haskell with IntCode module

2 files changed, 108 insertions, 0 deletions

diff --git a/2019/2.hs b/2019/2.hs
new file mode 100644
index 0000000..51ecf12
--- /dev/null
+++ b/2019/2.hs
@@ -0,0 +1,16 @@
+module Main where
+
+import Control.Monad
+
+import Input
+import IntCode
+
+
+main :: IO ()
+main = do
+    initMem <- parse . head <$> getInput 2
+    let set12 a b mem = head mem : a : b : drop 3 mem
+        run' a b = fst (run (set12 a b initMem) []) !! 0
+    print (run' 12 2)
+    forM_ [0..99] $ \a -> forM_ [0..99] $ \b ->
+        when (run' a b == 19690720) $ print (100 * a + b)
diff --git a/2019/IntCode.hs b/2019/IntCode.hs
new file mode 100644
index 0000000..904e961
--- /dev/null
+++ b/2019/IntCode.hs
@@ -0,0 +1,92 @@
+{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, FlexibleContexts #-}
+module IntCode (
+    IC(..),
+    parse, decode, run
+) where
+
+import Control.Monad.ST
+import qualified Data.Array.ST as SA
+import qualified Data.Array.IArray as IA
+
+
+data Arg = Imm Int | Addr Int
+  deriving (Show)
+
+data IC
+    = Add Arg Arg Arg
+    | Mul Arg Arg Arg
+    | Inp Arg
+    | Out Arg
+    | Hlt
+  deriving (Show)
+
+parse :: String -> [Int]
+parse = map read . splitOn ','
+  where
+    splitOn c s = case break (== c) s of
+                      (pre, _ : post) -> pre : splitOn c post
+                      _ -> [s]
+
+decode :: [Int] -> (IC, Int)
+decode [] = error "IC: Execution fell off end of program"
+decode (ins : rest) =
+    let (code, modes) = insModeSplit ins
+    in (decodeBase code modes rest, 1 + insNArgs code)
+
+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
+    readArray = SA.readArray
+
+decode' :: ReadArray m a Int Int => a -> Int -> m (IC, Int)
+decode' arr ptr = do
+    ins <- readArray arr ptr
+    let (code, modes) = insModeSplit ins
+        nargs = insNArgs code
+    params <- sequence [readArray arr (ptr + i) | i <- [1..nargs]]
+    return (decodeBase code modes params, 1 + nargs)
+
+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;
+                         _ -> error $ "IC: Unknown instruction " ++ show n }
+
+decodeBase :: Int -> [Int] -> [Int] -> 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)
+decodeBase  4 (am:      _) (a:    _) = Out (insArg am a)
+decodeBase 99 (         _) (      _) = Hlt
+decodeBase ic _            _         = error $ "IC: Unknown instruction " ++ show ic
+
+insArg :: Int -> Int -> Arg
+insArg 0 n = Addr n
+insArg 1 n = Imm n
+insArg m n = error $ "IC: Unknown parameter mode " ++ show m ++ " for parameter " ++ show n
+
+run :: [Int] -> [Int] -> ([Int], [Int])
+run initMem input = runST $ do
+     arr <- SA.newListArray (0, length initMem - 1) initMem
+     out <- run' arr 0 input
+     mem <- SA.getElems arr
+     return (mem, out)
+  where
+    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'
+        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
+            Inp (Addr a) -> case inp of
+                         val : rest -> SA.writeArray arr a val >> continue rest
+                         _ -> error "IC: Not enough input"
+            Out a -> resolve a >>= \val -> (val :) <$> continue inp
+            Hlt -> return []
+            ins -> error $ "IC: Invalid instruction " ++ show ins
+      where
+        resolve (Imm n) = return n
+        resolve (Addr n) = SA.readArray arr n