restore base Env after an exception bubbles to top level

src/swarm-engine/Swarm/Game/CESK.hs

@@ -82,7 +82,7 @@ module Swarm.Game.CESK (
   cont,
 ) where
 
-import Control.Lens (Lens', Traversal', lens, traversal, (^.))
+import Control.Lens (Getter, Lens', Traversal', lens, to, traversal, (^.))
 import Data.Aeson (FromJSON (..), ToJSON (..), genericParseJSON, genericToJSON)
 import Data.IntMap.Strict (IntMap)
 import Data.IntMap.Strict qualified as IM
@@ -165,6 +165,9 @@ data Frame
     FProj Var
   | -- | We should suspend once we finish the current evaluation.
     FSuspend Env
+  | -- | If an exception bubbles all the way up to this frame, then
+    --   switch to Suspended mode with this saved top-level context.
+    FRestoreEnv Env
   deriving (Eq, Show, Generic)
 
 instance ToJSON Frame where
@@ -277,11 +280,12 @@ data CESK
   | -- | The machine is suspended, i.e. waiting for another term to
     --   evaluate.  This happens after we have evaluated whatever the
     --   user entered at the REPL and we are waiting for them to type
-    --   something else.  Conceptually, this is like a combination of
+    --   something else.  Conceptually, this almost like a combination of
     --   'Out' and 'In': we store a 'Value' that was just yielded by
     --   evaluation, and otherwise it is just like 'In' with a hole
-    --   for the 'Term' we are going to evaluate.
-    Suspended Value Env Store Cont
+    --   for the 'Term' we are going to evaluate.  However, we do not
+    --   store a continuation.
+    Suspended Value Env Store
   deriving (Eq, Show, Generic)
 
 instance ToJSON CESK where
@@ -295,7 +299,7 @@ instance FromJSON CESK where
 finalValue :: CESK -> Maybe Value
 {-# INLINE finalValue #-}
 finalValue (Out v _ []) = Just v
-finalValue (Suspended v _ _ _) = Just v
+finalValue (Suspended v _ _) = Just v
 finalValue _ = Nothing
 
 -- | Extract the environment from a suspended CESK machine (/e.g./ to
@@ -304,7 +308,7 @@ suspendedEnv :: Traversal' CESK Env
 suspendedEnv = traversal go
  where
   go :: Applicative f => (Env -> f Env) -> CESK -> f CESK
-  go f (Suspended v e s k) = Suspended v <$> f e <*> pure s <*> pure k
+  go f (Suspended v e s) = Suspended v <$> f e <*> pure s
   go _ cesk = pure cesk
 
 -- | Lens focusing on the store of a CESK machine.
@@ -316,30 +320,24 @@ store = lens get set
     Out _ s _ -> s
     Up _ s _ -> s
     Waiting _ c -> get c
-    Suspended _ _ s _ -> s
+    Suspended _ _ s -> s
   set cesk s = case cesk of
     In t e _ k -> In t e s k
     Out v _ k -> Out v s k
     Up x _ k -> Up x s k
     Waiting t c -> Waiting t (set c s)
-    Suspended v e _ k -> Suspended v e s k
+    Suspended v e _ -> Suspended v e s
 
 -- | Lens focusing on the continuation of a CESK machine.
-cont :: Lens' CESK Cont
-cont = lens get set
+cont :: Getter CESK Cont
+cont = to get
  where
   get = \case
     In _ _ _ k -> k
     Out _ _ k -> k
     Up _ _ k -> k
     Waiting _ c -> get c
-    Suspended _ _ _ k -> k
-  set cesk k = case cesk of
-    In t e s _ -> In t e s k
-    Out v s _ -> Out v s k
-    Up x s _ -> Up x s k
-    Waiting t c -> Waiting t (set c k)
-    Suspended v e s _ -> Suspended v e s k
+    Suspended {} -> []
 
 -- | Create a brand new CESK machine, with empty environment and
 --   store, to evaluate a given term.  We always initialize the
@@ -356,7 +354,7 @@ initMachine t = In (prepareTerm mempty t) mempty emptyStore [FExec]
 --   term is suitable for execution by the base (REPL) robot.
 continue :: TSyntax -> CESK -> CESK
 continue t = \case
-  Suspended _ e s k -> In (insertSuspend $ prepareTerm e t) e s (FExec : k)
+  Suspended _ e s -> In (insertSuspend $ prepareTerm e t) e s [FExec, FRestoreEnv e]
   cesk -> In (insertSuspend $ prepareTerm mempty t) mempty (cesk ^. store) [FExec]
 
 -- | Prepare a term for evaluation by a CESK machine in the given
@@ -397,7 +395,7 @@ instance PrettyPrec CESK where
     Out v _ k -> prettyCont k (11, "◀" <> ppr (valueToTerm v) <> "▶")
     Up e _ k -> prettyCont k (11, "!" <> (pretty (formatExn mempty e) <> "!"))
     Waiting t cesk -> "🕑" <> pretty t <> "(" <> ppr cesk <> ")"
-    Suspended v _ _ k -> prettyCont k (11, "◀" <> ppr (valueToTerm v) <> "...▶")
+    Suspended v _ _ -> prettyCont [] (11, "◀" <> ppr (valueToTerm v) <> "...▶")
 
 -- | Take a continuation, and the pretty-printed expression which is
 --   the focus of the continuation (i.e. the expression whose value
@@ -443,6 +441,7 @@ prettyFrame f (p, inner) = case f of
     pprEq (x, Just t) = pretty x <+> "=" <+> ppr t
   FProj x -> (11, pparens (p < 11) inner <> "." <> pretty x)
   FSuspend _ -> (10, "suspend" <+> pparens (p < 11) inner)
+  FRestoreEnv _ -> (10, "restore" <+> pparens (p < 11) inner)
 
 -- | Pretty-print a special "prefix application" frame, i.e. a frame
 --   formatted like @X· inner@.  Unlike typical applications, these

src/swarm-engine/Swarm/Game/Step.hs

@@ -752,7 +752,7 @@ stepCESK cesk = case cesk of
   -- suspend.
   Out (VSuspend t e) s (FExec : k) -> return $ In t e s (FSuspend e : k)
   -- Once we've finished, enter the Suspended state.
-  Out v s (FSuspend e : k) -> return $ Suspended v e s k
+  Out v s (FSuspend e : _) -> return $ Suspended v e s
   -- Any other type of value wiwth an FExec frame is an error (should
   -- never happen).
   Out _ s (FExec : _) -> badMachineState s "FExec frame with non-executable value"
@@ -763,9 +763,42 @@ stepCESK cesk = case cesk of
   -- First, if we were running a try block but evaluation completed normally,
   -- just ignore the try block and continue.
   Out v s (FTry {} : k) -> return $ Out v s k
-  Up exn s [] -> do
-    -- Here, an exception has risen all the way to the top level without being
-    -- handled.
+  -- Also ignore restore frames when returning normally.
+  Out v s (FRestoreEnv {} : k) -> return $ Out v s k
+  -- If raising an exception up the stack and we reach the top, handle
+  -- it appropriately.
+  Up exn s [] -> handleException exn s Nothing
+  Up exn s (FRestoreEnv e : _) -> handleException exn s (Just e)
+  -- If we are raising a catchable exception up the continuation
+  -- stack and come to a Try frame, force and then execute the associated catch
+  -- block.
+  Up exn s (FTry c : k)
+    | isCatchable exn -> return $ Out c s (FApp (VCApp Force []) : FExec : k)
+  -- If we are raising an exception up the stack and we see an FRestoreEnv frame,
+  -- switch into a suspended state.
+
+  -- Otherwise, keep popping from the continuation stack.
+  Up exn s (_ : k) -> return $ Up exn s k
+  -- Finally, if we're done evaluating and the continuation stack is
+  -- empty, OR if we've hit a suspend, return the machine unchanged.
+  done@(Out _ _ []) -> return done
+  suspended@(Suspended {}) -> return suspended
+ where
+  badMachineState s msg =
+    let msg' =
+          T.unlines
+            [ T.append "Bad machine state in stepRobot: " msg
+            , prettyText cesk
+            ]
+     in return $ Up (Fatal msg') s []
+
+  isCatchable = \case
+    Fatal {} -> False
+    Incapable {} -> False
+    InfiniteLoop {} -> False
+    _ -> True
+
+  handleException exn s menv = do
     case exn of
       CmdFailed _ _ (Just a) -> do
         grantAchievement a
@@ -783,31 +816,9 @@ stepCESK cesk = case cesk of
     h <- hasCapability CLog
     em <- use $ landscape . terrainAndEntities . entityMap
     when h $ void $ traceLog RobotError Error (formatExn em exn)
-    return $ Out VExc s' []
-
-  -- Fatal errors, capability errors, and infinite loop errors can't
-  -- be caught; just throw away the continuation stack.
-  Up exn@Fatal {} s _ -> return $ Up exn s []
-  Up exn@Incapable {} s _ -> return $ Up exn s []
-  Up exn@InfiniteLoop {} s _ -> return $ Up exn s []
-  -- Otherwise, if we are raising an exception up the continuation
-  -- stack and come to a Try frame, force and then execute the associated catch
-  -- block.
-  Up _ s (FTry c : k) -> return $ Out c s (FApp (VCApp Force []) : FExec : k)
-  -- Otherwise, keep popping from the continuation stack.
-  Up exn s (_ : k) -> return $ Up exn s k
-  -- Finally, if we're done evaluating and the continuation stack is
-  -- empty, OR if we've hit a suspend, return the machine unchanged.
-  done@(Out _ _ []) -> return done
-  suspended@(Suspended {}) -> return suspended
- where
-  badMachineState s msg =
-    let msg' =
-          T.unlines
-            [ T.append "Bad machine state in stepRobot: " msg
-            , prettyText cesk
-            ]
-     in return $ Up (Fatal msg') s []
+    return $ case menv of
+      Nothing -> Out VExc s' []
+      Just env -> Suspended VExc env s'
 
 -- | Execute the given program *hypothetically*: i.e. in a fresh
 -- CESK machine, using *copies* of the current store, robot