{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, FlexibleContexts #-} module IntCode ( Arg(..), IC(..), mkAddr, mkImm, mkRel, parse, decode, run, unparse, runInterruptible, runContinue, Continuation(), initialContinuation ) 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 = Addr Integer | Imm Integer | Rel Integer deriving (Show) mkAddr, mkImm, mkRel :: Integral i => i -> Arg mkAddr = Addr . fromIntegral mkImm = Imm . fromIntegral mkRel = Rel . fromIntegral 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 | Adj Arg | Hlt deriving (Show) parse :: String -> [Integer] parse = map read . splitOn ',' where splitOn c s = case break (== c) s of (pre, _ : post) -> pre : splitOn c post _ -> [s] unparse :: IC -> [Integer] 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] Adj a -> go 9 [a ] Hlt -> [99] where go code as = (100 * mode as + code) : map bare as where bare (Addr n) = n bare (Imm n) = n bare (Rel n) = n mode [] = 0 mode (Addr _ : as) = 0 + 10 * mode as mode (Imm _ : as) = 1 + 10 * mode as mode (Rel _ : as) = 2 + 10 * mode as decode :: [Integer] -> (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.STArray s i Integer) i Integer where readArray = SA.readArray decode' :: ReadArray m a Int Integer => 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 :: Integer -> (Integer, [Integer]) insModeSplit n = (n `mod` 100, map (`mod` 10) (iterate (`div` 10) (n `div` 100))) insNArgs :: Integer -> Int insNArgs n = case n of { 1 -> 3; 2 -> 3; 3 -> 1; 4 -> 1; 5 -> 2; 6 -> 2; 7 -> 3; 8 -> 3; 9 -> 1; 99 -> 0; _ -> error $ "IC: Unknown instruction " ++ show n } decodeBase :: Integer -> [Integer] -> [Integer] -> 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 9 (am: _) (a: _) = Adj (insArg am a) decodeBase 99 ( _) ( _) = Hlt decodeBase ic _ _ = error $ "IC: Unknown instruction " ++ show ic insArg :: Integer -> Integer -> Arg insArg 0 n = Addr n insArg 1 n = Imm n insArg 2 n = Rel n insArg m n = error $ "IC: Unknown parameter mode " ++ show m ++ " for parameter " ++ show n run :: [Integer] -> [Integer] -> ([Integer], [Integer]) run initMem input = case runInterruptible initMem input of Left _ -> error "IC: Not enough input" Right res -> res data Continuation = Continuation (A.Array Int Integer) Int Int runInterruptible :: [Integer] -> [Integer] -> Either (Continuation, [Integer]) ([Integer], [Integer]) runInterruptible initMem input = runST $ do arr <- SA.newListArray (0, length initMem - 1) initMem res <- runArray arr 0 input 0 case fst res of Left cont -> return (Left (cont, snd res)) Right mem -> return (Right (mem, snd res)) runContinue :: Continuation -> [Integer] -> Either (Continuation, [Integer]) ([Integer], [Integer]) runContinue (Continuation frozen ptr offset) moreinput = runST $ do arr <- SA.thaw frozen res <- runArray arr ptr moreinput offset case fst res of Left cont -> return (Left (cont, snd res)) Right mem -> return (Right (mem, snd res)) initialContinuation :: [Integer] -> Continuation initialContinuation initMem = Continuation (A.listArray (0, length initMem - 1) initMem) 0 0 runArray :: SA.STArray s Int Integer -> Int -> [Integer] -> Int -> ST s (Either Continuation [Integer], [Integer]) runArray arr ptr inp offset = do (ic, len) <- decode' arr ptr let continue = runArray arr (ptr + len) inp offset continueArr arr' = runArray arr' (ptr + len) inp offset jump ptr' = runArray arr (fromIntegral ptr') inp offset continueOff off' = runArray arr (ptr + len) inp (fromIntegral off') continueInpArr inp' arr' = runArray arr' (ptr + len) inp' offset resize arr1 idx = do (_, right) <- SA.getBounds arr1 if idx < right then return arr1 else SA.getElems arr1 >>= \els -> SA.newListArray (0, 2 * idx) (els ++ replicate (2 * idx - right) 0) performWrite arr1 idx value = do arr' <- resize arr1 (fromIntegral idx) SA.writeArray arr' (fromIntegral idx) value return arr' case ic of Add a b c -> (+) <$> resolve a <*> resolve b >>= performWrite arr (resolveAddr c) >>= continueArr Mul a b c -> (*) <$> resolve a <*> resolve b >>= performWrite arr (resolveAddr c) >>= continueArr Inp a -> case inp of val : rest -> performWrite arr (resolveAddr a) val >>= continueInpArr rest [] -> do frozen <- SA.freeze arr return (Left (Continuation frozen ptr offset), []) 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 c -> (((fromIntegral . fromEnum) .) . (<)) <$> resolve a <*> resolve b >>= performWrite arr (resolveAddr c) >>= continueArr Ceq a b c -> (((fromIntegral . fromEnum) .) . (==)) <$> resolve a <*> resolve b >>= performWrite arr (resolveAddr c) >>= continueArr Adj a -> resolve a >>= \add -> continueOff (fromIntegral offset + add) Hlt -> do mem <- SA.getElems arr return (Right mem, []) where resolve (Addr n) = do (_, right) <- SA.getBounds arr let n' = fromIntegral n if n' > right then return 0 else SA.readArray arr n' resolve (Imm n) = return n resolve (Rel n) = resolve (Addr (fromIntegral offset + n)) resolveAddr (Addr n) = n resolveAddr (Rel n) = fromIntegral offset + n resolveAddr (Imm _) = error "IC: Using immediate value as address operand"