From 1da88766ea52d28dbd711155639711b1c979c1ec Mon Sep 17 00:00:00 2001
From: Andrzej Rybczak <andrzej@rybczak.net>
Date: Thu, 12 Sep 2024 17:57:24 +0200
Subject: [PATCH] Properly roll back changes made to the environment with
 OnEmptyKeep

---
 effectful-core/CHANGELOG.md            |  2 ++
 effectful-core/src/Effectful/NonDet.hs | 11 ++++++++++-
 effectful/CHANGELOG.md                 |  2 ++
 effectful/tests/NonDetTests.hs         | 13 ++++++++-----
 4 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/effectful-core/CHANGELOG.md b/effectful-core/CHANGELOG.md
index 5a7257d..b3793b0 100644
--- a/effectful-core/CHANGELOG.md
+++ b/effectful-core/CHANGELOG.md
@@ -13,6 +13,8 @@
   ([#237](https://github.com/haskell-effectful/effectful/issues/237)).
 * Add `HasCallStack` constraints where appropriate for better debugging
   experience.
+* Properly roll back changes made to the environment when `OnEmptyKeep` policy
+  for the `NonDet` effect is selected.
 * **Breaking changes**:
   - `localSeqLend`, `localLend`, `localSeqBorrow` and `localBorrow` now take a
     list of effects instead of a single one.
diff --git a/effectful-core/src/Effectful/NonDet.hs b/effectful-core/src/Effectful/NonDet.hs
index febe20c..1f6cc27 100644
--- a/effectful-core/src/Effectful/NonDet.hs
+++ b/effectful-core/src/Effectful/NonDet.hs
@@ -75,7 +75,7 @@ runNonDetRollback
   :: HasCallStack
   => Eff (NonDet : es) a
   -> Eff es (Either CallStack a)
-runNonDetRollback = reinterpret (fmap noError . runError @()) $ \env -> \case
+runNonDetRollback = reinterpret setup $ \env -> \case
   Empty       -> throwError ()
   m1 :<|>: m2 -> do
     backupEnv <- cloneLocalEnv env
@@ -87,6 +87,15 @@ runNonDetRollback = reinterpret (fmap noError . runError @()) $ \env -> \case
       case mr of
         Just r  -> pure r
         Nothing -> unlift m2
+  where
+    setup action = do
+      backupEs <- unsafeEff cloneEnv
+      runError @() action >>= \case
+        Right r -> pure $ Right r
+        Left (cs, _) -> do
+          -- If the whole action failed, roll back the environment.
+          unsafeEff $ \es -> restoreEnv es backupEs
+          pure $ Left cs
 
 ----------------------------------------
 
diff --git a/effectful/CHANGELOG.md b/effectful/CHANGELOG.md
index 44628d2..20f7299 100644
--- a/effectful/CHANGELOG.md
+++ b/effectful/CHANGELOG.md
@@ -15,6 +15,8 @@
   ([#237](https://github.com/haskell-effectful/effectful/issues/237)).
 * Add `HasCallStack` constraints where appropriate for better debugging
   experience.
+* Properly roll back changes made to the environment when `OnEmptyKeep` policy
+  for the `NonDet` effect is selected.
 * **Breaking changes**:
   - `localSeqLend`, `localLend`, `localSeqBorrow` and `localBorrow` now take a
     list of effects instead of a single one.
diff --git a/effectful/tests/NonDetTests.hs b/effectful/tests/NonDetTests.hs
index 8cb558e..aae9401 100644
--- a/effectful/tests/NonDetTests.hs
+++ b/effectful/tests/NonDetTests.hs
@@ -30,11 +30,11 @@ nonDetTests = testGroup "NonDet"
 
     expectedLocalState :: OnEmptyPolicy -> Int
     expectedLocalState = \case
-      OnEmptyKeep     -> 3
-      OnEmptyRollback -> 2
+      OnEmptyKeep     -> 7
+      OnEmptyRollback -> 0
 
     expectedSharedState :: OnEmptyPolicy -> Int
-    expectedSharedState _ = 3
+    expectedSharedState _ = 7
 
 test_empty
   :: Eff [NonDet, IOE] Bool
@@ -53,12 +53,15 @@ test_state
 test_state evalState expectedState step = runEff $ do
   evalState 0 . runNonDetBoth test $ \policy result -> do
     liftIO . step $ show policy
-    U.assertEqual "result" (Just ()) (dropLeft result)
+    U.assertEqual "result" Nothing (dropLeft result)
     s <- state @Int $ \s -> (s, 0)
     U.assertEqual "state" (expectedState policy) s
   where
     test :: (NonDet :> es, State Int :> es) => Eff es ()
-    test = (modify @Int (+1) >> empty) <|> modify @Int (+2)
+    test = do
+      modify @Int (+1)
+      _<- (modify @Int (+2) >> empty) <|> (modify @Int (+4) >> empty)
+      modify @Int (+8)
 
 test_independentHandlers :: (String -> IO ()) -> Assertion
 test_independentHandlers step = runEff $ do