summaryrefslogtreecommitdiff
path: root/2019/IntCode.hs
blob: 699a40ca0b10d6799606b64d604aed21ca9d71e9 (plain)
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
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, FlexibleContexts #-}
module IntCode (
    Arg(..), IC(..),
    parse, decode, run, unparse,
    runInterruptible, runContinue, Continuation()
) where

import Control.Monad.ST
import qualified Data.Array as A
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
    | Jnz Arg Arg
    | Jez Arg Arg
    | Clt Arg Arg Arg
    | Ceq Arg Arg 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]

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) =
    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;
                         5 -> 2; 6 -> 2; 7 -> 3; 8 -> 3;
                         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  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

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 =
    case runInterruptible initMem input of
        Left _ -> error "IC: Not enough input"
        Right res -> res

data Continuation = Continuation (A.Array Int Int) Int

runInterruptible :: [Int] -> [Int] -> Either (Continuation, [Int]) ([Int], [Int])
runInterruptible initMem input = runST $ do
    arr <- SA.newListArray (0, length initMem - 1) initMem
    res <- runArray arr 0 input
    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
    arr <- SA.thaw frozen
    res <- runArray arr ptr moreinput
    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
    (ic, len) <- decode' arr ptr
    let continueInp inp' = runArray arr (ptr + len) inp'
        continue = continueInp inp
        jump ptr' = runArray arr ptr' inp
    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
        Inp (Addr a) -> case inp of
                     val : rest -> SA.writeArray arr a val >> continueInp rest
                     [] -> do
                         frozen <- SA.freeze arr
                         return (Left (Continuation frozen ptr), [])
        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
        Hlt -> do
            mem <- SA.getElems arr
            return (Right mem, [])
        ins -> error $ "IC: Invalid instruction " ++ show ins
  where
    resolve (Imm n) = return n
    resolve (Addr n) = SA.readArray arr n