Simple intcode assembler

3 files changed, 126 insertions, 2 deletions

diff --git a/2019/.gitignore b/2019/.gitignore
index ebbfb2e..ef5c4cb 100644
--- a/2019/.gitignore
+++ b/2019/.gitignore
@@ -25,3 +25,4 @@ obj/
 23
 24
 25
+ic-asm
diff --git a/2019/IntCode.hs b/2019/IntCode.hs
index 89b54a2..9e94066 100644
--- a/2019/IntCode.hs
+++ b/2019/IntCode.hs
@@ -1,7 +1,7 @@
 {-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, FlexibleContexts #-}
 module IntCode (
-    IC(..),
-    parse, decode, run
+    Arg(..), IC(..),
+    parse, decode, run, unparse
 ) where
 
 import Control.Monad.ST
@@ -31,6 +31,25 @@ parse = map read . splitOn ','
                       (pre, _ : post) -> pre : splitOn c post
                       _ -> [s]
 
+unparse :: IC -> [Int]
+unparse ic = case ic of
+                 Add a b c -> go 1 [a,b,c]
+                 Mul a b c -> go 2 [a,b,c]
+                 Inp a     -> go 3 [a    ]
+                 Out a     -> go 4 [a    ]
+                 Jnz a b   -> go 5 [a,b  ]
+                 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]
+                 Hlt       -> [99]
+  where
+    go code as = (100 * mode as + code) : map bare as
+      where bare (Imm n) = n
+            bare (Addr n) = n
+    mode [] = 0
+    mode (Imm _ : as) = 1 + 10 * mode as
+    mode (Addr _ : as) = 10 * mode as
+
 decode :: [Int] -> (IC, Int)
 decode [] = error "IC: Execution fell off end of program"
 decode (ins : rest) =
diff --git a/2019/ic-asm.hs b/2019/ic-asm.hs
new file mode 100644
index 0000000..f3f5540
--- /dev/null
+++ b/2019/ic-asm.hs
@@ -0,0 +1,104 @@
+module Main where
+
+import qualified Data.Map.Strict as Map
+import qualified Data.Set as Set
+
+import qualified IntCode as IC
+
+
+type Name = String
+
+data Ref = Imm Int | Var Name
+  deriving (Show)
+
+data Stmt
+    = Declare Name
+    | Set Name Ref
+    | Add Name Ref Ref
+    | Mul Name Ref Ref
+    | Input Name
+    | Output Ref
+    | Clt Name Ref Ref
+    | Ceq Name Ref Ref
+    | If Ref [Stmt] [Stmt]
+    | While Ref [Stmt]
+    | Halt
+  deriving (Show)
+
+assemble :: [Stmt] -> [Int]
+assemble stmts =
+    IC.unparse (IC.Jnz (IC.Imm 1) (IC.Imm (3 + nvars)))
+    ++ replicate nvars 0
+    ++ concatMap IC.unparse (snd (goL nvars stmts))
+  where
+    vars = Set.unions (map collectDeclares stmts)
+    var2idx = Map.fromList (zip (Set.toList vars) [3..])
+    nvars = Set.size vars
+
+    goL :: Int -> [Stmt] -> (Int, [IC.IC])
+    goL = (fmap concat .) . foldProduce go
+
+    go :: Int -> Stmt -> ([IC.IC], Int)
+    go off (Declare _) = ([], off)
+    go off (Set n a) = go off (Add n a (Imm 0))
+    go off (Add n a b) = ([IC.Add (ref a) (ref b) (ref (Var n))], off + 4)
+    go off (Mul n a b) = ([IC.Mul (ref a) (ref b) (ref (Var n))], off + 4)
+    go off (Input n) = ([IC.Inp (ref (Var n))], off + 2)
+    go off (Output a) = ([IC.Out (ref a)], off + 2)
+    go off (Clt n a b) = ([IC.Clt (ref a) (ref b) (ref (Var n))], off + 4)
+    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
+           , 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))]
+            ++ ics2
+            ++ [IC.Jnz (IC.Imm 1) (IC.Imm offAfterElse)]
+            ++ ics1
+           , offAfterElse)
+    go off (While a s) =
+        let (offAfterBody, ics) = goL (off + 3) s
+        in ([IC.Jez (ref a) (IC.Imm (offAfterBody + 3))]
+            ++ ics
+            ++ [IC.Jnz (ref a) (IC.Imm (off + 3))]
+           , offAfterBody + 3)
+    go off Halt = ([IC.Hlt], off + 1)
+
+    ref (Imm n) = IC.Imm n
+    ref (Var n) = case Map.lookup n var2idx of
+                      Just idx -> IC.Addr idx
+                      Nothing  -> error $ "Undeclared variable '" ++ n ++ "'"
+
+collectDeclares :: Stmt -> Set.Set Name
+collectDeclares (Declare n) = Set.singleton n
+collectDeclares (If _ a b) = Set.unions (map collectDeclares (a ++ b))
+collectDeclares (While _ a) = Set.unions (map collectDeclares a)
+collectDeclares _ = Set.empty
+
+foldProduce :: (s -> a -> (b, s)) -> s -> [a] -> (s, [b])
+foldProduce _ s [] = (s, [])
+foldProduce f s (x:xs) = let (y, s') = f s x in fmap (y :) (foldProduce f s' xs)
+
+
+main :: IO ()
+main = do
+    let prog = [Declare "a"
+               ,Declare "c"
+               ,Declare "out"
+               ,Set "out" (Imm 1)
+               ,Input "a"
+               ,Clt "c" (Imm 0) (Var "a")
+               ,While (Var "c")
+                    [Mul "out" (Var "out") (Var "a")
+                    ,Add "a" (Var "a") (Imm (-1))
+                    ,Clt "c" (Imm 0) (Var "a")
+                    ]
+               ,Output (Var "out")
+               ,Halt
+               ]
+    let res = assemble prog
+    print res
+    print (snd (IC.run res [10]))