Properly roll back changes made to the environment with OnEmptyKeep

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.

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
 
 ----------------------------------------
 

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.

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