assignRef no longer mutates the State object

1 files changed, 32 insertions, 25 deletions

diff --git a/VM.hs b/VM.hs
index c996e7b..8d7bb80 100644
--- a/VM.hs
+++ b/VM.hs
@@ -66,37 +66,46 @@ vmRunBB info state (BB _ inss term) = do
 vmRunInstr :: Info -> State -> Instruction -> IO State
 -- vmRunInstr _ _ ins | traceShow ins False = undefined
 vmRunInstr info@(Info _ _ datas) state (dest, instr) = case instr of
-    IAssign ref ->
+    IAssign ref -> do
         findRef (sTempMap state) ref >>= assignRef state dest
-    IParam i ->
+        return state
+    IParam i -> do
         let args = sArgs state
-        in if i < length args then assignRef state dest (args !! i)
-               else error $ show args ++ ", " ++ show i ++ ", param-out-of-range"
-    IClosure i ->
+        if i < length args then assignRef state dest (args !! i)
+            else error $ show args ++ ", " ++ show i ++ ", param-out-of-range"
+        return state
+    IClosure i -> do
         let closure = sCloVals state
-        in if i < length closure then assignRef state dest (closure !! i)
-               else error $ show closure ++ ", " ++ show i ++ ", closure-out-of-range"
-    IData i ->
+        if i < length closure then assignRef state dest (closure !! i)
+            else error $ show closure ++ ", " ++ show i ++ ", closure-out-of-range"
+        return state
+    IData i -> do
         if i < length datas then assignRef state dest (toRunValue (datas !! i))
             else error "data-out-of-range"
+        return state
     ICallC cl as -> do
         (rv, state') <- callClosure info state cl as
         assignRef state' dest rv
-    IAllocClo name clrefs ->
-        RVClosure name <$> mapM (findRef (sTempMap state)) clrefs >>= assignRef state dest
-    IDiscard ref ->
+        return state'
+    IAllocClo name clrefs -> do
+        clovals <- mapM (findRef (sTempMap state)) clrefs
+        assignRef state dest (RVClosure name clovals)
+        return state
+    IDiscard ref -> do
         case ref of
             RTemp _ -> assignRef state ref (RVNum 0)
-            _ -> return state
-    IPush refs ->
-        mapM (findRef (sTempMap state)) refs >>= \values ->
-            return (state { sStack = values ++ sStack state })
-    IPop refs ->
-        if length (sStack state) >= length refs
-            then let (values, newStack) = splitAt (length refs) (sStack state)
-                     state' = state { sStack = newStack }
-                 in foldM (uncurry . assignRef) state' (zip refs values)
-            else error "VM: IPop on too-small stack"
+            _ -> return ()
+        return state
+    IPush refs -> do
+        values <- mapM (findRef (sTempMap state)) refs
+        return (state { sStack = values ++ sStack state })
+    IPop refs -> do
+        when (length (sStack state) < length refs) $
+            error "VM: IPop on too-small stack"
+        let (values, newStack) = splitAt (length refs) (sStack state)
+            state' = state { sStack = newStack }
+        mapM_ (uncurry (assignRef state')) (zip refs values)
+        return state'
 
 vmRunTerm :: Info -> State -> Terminator -> IO (RunValue, State)
 vmRunTerm info@(Info bbmap _ _) state@(State { sTempMap = tmap }) term = case term of
@@ -115,10 +124,8 @@ findRef tmap (RTemp i) = A.readArray tmap i
 findRef _ (RSClo name) = return (RVClosure name [])
 findRef _ RNone = error "VM: None ref used"
 
-assignRef :: State -> Ref -> RunValue -> IO State
-assignRef state (RTemp i) rv = do
-    A.writeArray (sTempMap state) i rv
-    return state  -- TODO: now assignRef doesn't even mutate the state object anymore
+assignRef :: State -> Ref -> RunValue -> IO ()
+assignRef state (RTemp i) rv = A.writeArray (sTempMap state) i rv
 assignRef _ ref _ = error $ "VM: Unexpected ref " ++ show ref ++ " assigned"
 
 callClosure :: Info -> State -> Ref -> [Ref] -> IO (RunValue, State)