Use IOArray for TempMap in VM

1 files changed, 36 insertions, 29 deletions

diff --git a/VM.hs b/VM.hs
index 554a310..c996e7b 100644
--- a/VM.hs
+++ b/VM.hs
@@ -1,5 +1,7 @@
+{-# LANGUAGE TupleSections, LambdaCase #-}
 module VM(vmRun) where
 
+import qualified Data.Array.IO as A
 import Control.Monad
 import Data.Char
 import Data.List
@@ -19,7 +21,7 @@ data Info =
          (Map.Map Name GlobFuncDef)  -- global functions
          [Value]  -- data table
 
-type TempMap = IMap.IntMap RunValue
+type TempMap = A.IOArray Int RunValue
 
 data State = State
     { sTempMap :: TempMap
@@ -44,11 +46,13 @@ kErrorExit :: String
 kErrorExit = "VM:exit"
 
 vmRun :: IRProgram -> IO ()
-vmRun (IRProgram bbs gfds datas) =
+vmRun irprogram@(IRProgram bbs gfds datas) = do
+    let alltemps = onlyTemporaries (allRefs irprogram)
+    tmap <- A.newArray (minimum alltemps, maximum alltemps) (RVNum 0)
     let bbmap = Map.fromList [(bidOf bb, bb) | bb <- bbs]
         info = Info bbmap gfds datas
-        state = State IMap.empty [] [] [] (IMap.fromList [(-1, stdin), (-2, stdout), (-3, stderr)]) 0
-    in IO.catchIOError (void $ vmRunBB info state (head bbs)) vmErrorHandler
+        state = State tmap [] [] [] (IMap.fromList [(-1, stdin), (-2, stdout), (-3, stderr)]) 0
+    IO.catchIOError (void $ vmRunBB info state (head bbs)) vmErrorHandler
 
 vmErrorHandler :: IOError -> IO ()
 vmErrorHandler e =
@@ -63,69 +67,72 @@ vmRunInstr :: Info -> State -> Instruction -> IO State
 -- vmRunInstr _ _ ins | traceShow ins False = undefined
 vmRunInstr info@(Info _ _ datas) state (dest, instr) = case instr of
     IAssign ref ->
-        return (assignRef state dest (findRef (sTempMap state) ref))
+        findRef (sTempMap state) ref >>= assignRef state dest
     IParam i ->
         let args = sArgs state
-        in if i < length args then return (assignRef state dest (args !! i))
+        in if i < length args then assignRef state dest (args !! i)
                else error $ show args ++ ", " ++ show i ++ ", param-out-of-range"
     IClosure i ->
         let closure = sCloVals state
-        in if i < length closure then return (assignRef state dest (closure !! i))
+        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 datas then return (assignRef state dest (toRunValue (datas !! i)))
+        if i < length datas then assignRef state dest (toRunValue (datas !! i))
             else error "data-out-of-range"
     ICallC cl as -> do
         (rv, state') <- callClosure info state cl as
-        return (assignRef state' dest rv)
+        assignRef state' dest rv
     IAllocClo name clrefs ->
-        let cloVals = (map (findRef (sTempMap state)) clrefs)
-        in return (assignRef state dest (RVClosure name cloVals))
+        RVClosure name <$> mapM (findRef (sTempMap state)) clrefs >>= assignRef state dest
     IDiscard ref ->
         case ref of
-            RTemp i -> return (state { sTempMap = IMap.delete i (sTempMap state) })
+            RTemp _ -> assignRef state ref (RVNum 0)
             _ -> return state
     IPush refs ->
-        return (state { sStack = map (findRef (sTempMap state)) refs ++ sStack state })
+        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 return (foldl (uncurry . assignRef) state' (zip refs values))
+                 in foldM (uncurry . assignRef) state' (zip refs values)
             else error "VM: IPop on too-small stack"
 
 vmRunTerm :: Info -> State -> Terminator -> IO (RunValue, State)
 vmRunTerm info@(Info bbmap _ _) state@(State { sTempMap = tmap }) term = case term of
-    IBr ref b1 b2 -> vmRunBB info state . (bbmap Map.!) $ if truthy (findRef tmap ref) then b1 else b2
+    IBr ref b1 b2 -> do
+        val <- findRef tmap ref
+        vmRunBB info state . (bbmap Map.!) $ if truthy val then b1 else b2
     IJmp b -> vmRunBB info state (bbmap Map.! b)
-    IRet ref -> return (findRef tmap ref, state)
+    IRet ref -> (,state) <$> findRef tmap ref
     ITailC cl as -> callClosure info state cl as
     IExit -> IO.ioError (IO.userError kErrorExit)
     IUnknown -> undefined
 
-findRef :: TempMap -> Ref -> RunValue
-findRef _ (RConst n) = RVNum n
-findRef tmap (RTemp i) = case IMap.lookup i tmap of
-                             Nothing -> error "Use of declared but uninitialised variable"
-                             Just v -> v
-findRef _ (RSClo name) = RVClosure name []
+findRef :: TempMap -> Ref -> IO RunValue
+findRef _ (RConst n) = return (RVNum n)
+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 -> State
-assignRef state (RTemp i) rv = state { sTempMap = IMap.insert i rv (sTempMap state) }
+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 _ ref _ = error $ "VM: Unexpected ref " ++ show ref ++ " assigned"
 
 callClosure :: Info -> State -> Ref -> [Ref] -> IO (RunValue, State)
 callClosure info@(Info bbmap gfds _) state@(State { sTempMap = tmap }) cl as =
-    case findRef tmap cl of
+    findRef tmap cl >>= \case
         RVClosure clname clvals -> case Map.lookup clname gfds of
             Just (GlobFuncDef b _ _) -> do
-                (rv, state') <- vmRunBB info (state { sArgs = map (findRef tmap) as, sCloVals = clvals })
-                                        (bbmap Map.! b)
+                args <- mapM (findRef tmap) as
+                (rv, state') <- vmRunBB info (state { sArgs = args, sCloVals = clvals }) (bbmap Map.! b)
                 return (rv, state' { sArgs = sArgs state, sCloVals = sCloVals state })
-            Nothing ->
+            Nothing -> do
                 -- Take 'tail as' to skip the first self-link argument
-                vmRunBuiltin state clname (map (findRef tmap) (tail as))
+                args <- mapM (findRef tmap) (tail as)
+                vmRunBuiltin state clname args
         obj -> error $ "VM: Cannot call non-closure object: " ++ show obj
 
 vmRunBuiltin :: State -> Name -> [RunValue] -> IO (RunValue, State)