From 7d2a4e861e6f61413f8bf216bdb0ea8aca8593c9 Mon Sep 17 00:00:00 2001
From: Brent Yorgey <byorgey@gmail.com>
Date: Sat, 15 Jun 2024 08:47:50 -0500
Subject: [PATCH 1/3] properly handle suspend with more continuation remaining

---
 src/swarm-engine/Swarm/Game/CESK.hs       |  2 +-
 src/swarm-engine/Swarm/Game/Step.hs       | 19 +++++++++++++++++--
 src/swarm-engine/Swarm/Game/Step/Const.hs |  9 +++------
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/src/swarm-engine/Swarm/Game/CESK.hs b/src/swarm-engine/Swarm/Game/CESK.hs
index 94a9a5ffe9..32f544c6b9 100644
--- a/src/swarm-engine/Swarm/Game/CESK.hs
+++ b/src/swarm-engine/Swarm/Game/CESK.hs
@@ -298,7 +298,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
diff --git a/src/swarm-engine/Swarm/Game/Step.hs b/src/swarm-engine/Swarm/Game/Step.hs
index 80106b33ab..e084c46a4c 100644
--- a/src/swarm-engine/Swarm/Game/Step.hs
+++ b/src/swarm-engine/Swarm/Game/Step.hs
@@ -757,6 +757,22 @@ stepCESK cesk = case cesk of
   -- never happen).
   Out _ s (FExec : _) -> badMachineState s "FExec frame with non-executable value"
   ------------------------------------------------------------
+  -- Suspension
+  ------------------------------------------------------------
+
+  -- If we're suspended and see the env restore frame, we can discard
+  -- it: it was only there in case an exception was thrown.
+  Suspended v e s (FRestoreEnv _ : k) -> return $ Suspended v e s k
+  -- If we're suspended but we were on the LHS of a bind, switch to
+  -- evaluating that, except with the environment from the suspension
+  -- instead of the environment stored in the FBind frame, as if the
+  -- RHS of the bind had been grafted in right where the suspend was.
+  Suspended _ e s (FBind _ _ t2 _ : k) -> return $ In t2 e s (FExec : k)
+  -- Otherwise, if we're suspended with nothing else left to do,
+  -- return the machine unchanged (but throw away the rest of the
+  -- continuation).
+  Suspended v e s _ -> return $ Suspended v e s []
+  ------------------------------------------------------------
   -- Exception handling
   ------------------------------------------------------------
 
@@ -780,9 +796,8 @@ stepCESK cesk = case cesk of
   -- 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.
+  -- empty, return the machine unchanged.
   done@(Out _ _ []) -> return done
-  suspended@(Suspended {}) -> return suspended
  where
   badMachineState s msg =
     let msg' =
diff --git a/src/swarm-engine/Swarm/Game/Step/Const.hs b/src/swarm-engine/Swarm/Game/Step/Const.hs
index 143ea7b911..c4ba5545e6 100644
--- a/src/swarm-engine/Swarm/Game/Step/Const.hs
+++ b/src/swarm-engine/Swarm/Game/Step/Const.hs
@@ -1226,13 +1226,10 @@ execConst runChildProg c vs s k = do
 
         case mt of
           Nothing -> return $ mkReturn ()
-          Just pt -> do
+          Just t -> do
             void $ traceLog CmdStatus Info "run: OK."
-            case k of
-              [] -> return $ In (insertSuspend (prepareTerm mempty pt)) mempty s [FExec]
-              _ -> return $ In (prepareTerm mempty pt) mempty s (FExec : k)
-      -- mempty above means we will wipe out all our current definitions every time
-      -- we execute 'run'...
+            cesk <- use machine
+            return $ continue t cesk
       _ -> badConst
     Not -> case vs of
       [VBool b] -> return $ Out (VBool (not b)) s k

From 03d6951f17c1d78af51134e344e707aa861b2132 Mon Sep 17 00:00:00 2001
From: Brent Yorgey <byorgey@gmail.com>
Date: Sat, 15 Jun 2024 08:47:50 -0500
Subject: [PATCH 2/3] Remove last XXX comment

---
 src/swarm-lang/Swarm/Language/Elaborate.hs | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/src/swarm-lang/Swarm/Language/Elaborate.hs b/src/swarm-lang/Swarm/Language/Elaborate.hs
index ea2ba40cd0..6e46d63544 100644
--- a/src/swarm-lang/Swarm/Language/Elaborate.hs
+++ b/src/swarm-lang/Swarm/Language/Elaborate.hs
@@ -24,9 +24,6 @@ import Swarm.Language.Types
 --   over.
 elaborate :: TSyntax -> TSyntax
 elaborate =
-  -- XXX Maybe pass the current environment into elaborate, so we can tell
-  -- _which_ free variables need to be wrapped in Force?
-
   -- Wrap all *free* variables in 'Force'.  Free variables must be
   -- referring to a previous definition, which are all wrapped in
   -- 'TDelay'.

From 6de8a179c0d7a404d48b1694fda4911baf62c743 Mon Sep 17 00:00:00 2001
From: "Restyled.io" <commits@restyled.io>
Date: Sat, 15 Jun 2024 16:39:07 +0000
Subject: [PATCH 3/3] Restyled by fourmolu

---
 test/unit/TestEval.hs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/unit/TestEval.hs b/test/unit/TestEval.hs
index 99542e1180..4161d0657f 100644
--- a/test/unit/TestEval.hs
+++ b/test/unit/TestEval.hs
@@ -335,7 +335,7 @@ testEval g =
       Left err ->
         p err
           @? "Expected predicate did not hold on error message "
-          ++ from @Text @String err
+            ++ from @Text @String err
 
   evaluatesTo :: Text -> Value -> Assertion
   evaluatesTo tm val = do