1 files changed, 51 insertions, 28 deletions
diff --git a/2019/IntCode.hs b/2019/IntCode.hs
index 9e94066..699a40c 100644
--- a/2019/IntCode.hs
+++ b/2019/IntCode.hs
@@ -1,10 +1,12 @@
 {-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, FlexibleContexts #-}
 module IntCode (
     Arg(..), IC(..),
-    parse, decode, run, unparse
+    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
 
@@ -97,31 +99,52 @@ 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)
+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
-    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
+    resolve (Imm n) = return n
+    resolve (Addr n) = SA.readArray arr n