summaryrefslogtreecommitdiff
path: root/2019/IntCode.hs
blob: 89b54a247cdbd124da4e1804fe2349e4726cfbda (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
{-# 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
    | 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]

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 = 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 continueInp inp' = run' arr (ptr + len) inp'
            continue = continueInp inp
            jump ptr' = run' 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
                         _ -> error "IC: Not enough input"
            Out a -> resolve a >>= \val -> (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 -> return []
            ins -> error $ "IC: Invalid instruction " ++ show ins
      where
        resolve (Imm n) = return n
        resolve (Addr n) = SA.readArray arr n