1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
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] -> [Integer]
assemble stmts =
IC.unparse (IC.Jnz (IC.Imm 1) (IC.mkImm (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.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.mkImm (offAfterThen + 3))]
++ ics2
++ [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.mkImm (offAfterBody + 3))]
++ ics
++ [IC.Jnz (ref a) (IC.mkImm (off + 3))]
, offAfterBody + 3)
go off Halt = ([IC.Hlt], off + 1)
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 ++ "'"
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]))
|
