Fix NonDet/OnEmptyRollback (#256)

effectful-core/CHANGELOG.md

@@ -1,6 +1,11 @@
 # effectful-core-2.5.0.0 (????-??-??)
 * Add `plusEff` (specialized version of `<|>`) to `Effectful.NonDet` and make
   `emptyEff` and `sumEff` generate better call stacks.
+* **Bugfixes**:
+  - `OnEmptyRollback` strategy of the `NonDet` effect is no longer broken.
+* **Breaking changes**:
+  - Removed `restoreEnv` function from `Effectful.Dispatch.Static.Primitive`
+    since it was broken.
 
 # effectful-core-2.4.0.0 (2024-10-08)
 * Add utility functions for handling effects that take the effect handler as the

effectful-core/src/Effectful/Dispatch/Static/Primitive.hs

@@ -31,7 +31,6 @@ module Effectful.Dispatch.Static.Primitive
     -- ** Utils
   , emptyEnv
   , cloneEnv
-  , restoreEnv
   , sizeEnv
   , tailEnv
   ) where

effectful-core/src/Effectful/Internal/Env.hs

@@ -9,12 +9,19 @@ module Effectful.Internal.Env
   , Ref(..)
   , Version
   , Storage(..)
-  , AnyEffect
-  , toAnyEffect
-  , fromAnyEffect
+
+    -- ** StorageData
+  , StorageData(..)
+  , copyStorageData
+  , restoreStorageData
+
+    -- *** Utils
   , AnyRelinker
   , toAnyRelinker
   , fromAnyRelinker
+  , AnyEffect
+  , toAnyEffect
+  , fromAnyEffect
 
     -- ** Relinker
   , Relinker(..)
@@ -29,7 +36,6 @@ module Effectful.Internal.Env
     -- * Operations
   , emptyEnv
   , cloneEnv
-  , restoreEnv
   , sizeEnv
   , tailEnv
 
@@ -134,13 +140,13 @@ newtype Version = Version Int
 
 -- | A storage of effects.
 data Storage = Storage
-  { stSize      :: !Int
-  , stVersion   :: !Version
-  , stVersions  :: !(MutablePrimArray RealWorld Version)
-  , stEffects   :: !(SmallMutableArray RealWorld AnyEffect)
-  , stRelinkers :: !(SmallMutableArray RealWorld AnyRelinker)
+  { stVersion :: !Version
+  , stData    :: {-# UNPACK #-} !StorageData
   }
 
+----------------------------------------
+-- StorageData
+
 -- | Effect in 'Storage'.
 newtype AnyEffect = AnyEffect Any
 
@@ -159,6 +165,46 @@ toAnyRelinker = AnyRelinker . toAny
 fromAnyRelinker :: AnyRelinker -> Relinker (EffectRep (DispatchOf e)) e
 fromAnyRelinker (AnyRelinker f) = fromAny f
 
+----------------------------------------
+
+data StorageData = StorageData
+  { sdSize      :: !Int
+  , sdVersions  :: !(MutablePrimArray RealWorld Version)
+  , sdEffects   :: !(SmallMutableArray RealWorld AnyEffect)
+  , sdRelinkers :: !(SmallMutableArray RealWorld AnyRelinker)
+  }
+
+-- | Make a shallow copy of the 'StorageData'.
+--
+-- @since 2.5.0.0
+copyStorageData :: HasCallStack => StorageData -> IO StorageData
+copyStorageData (StorageData storageSize vs0 es0 fs0) = do
+  vsSize <- getSizeofMutablePrimArray  vs0
+  esSize <- getSizeofSmallMutableArray es0
+  fsSize <- getSizeofSmallMutableArray fs0
+  when (vsSize /= esSize) $ do
+    error $ "vsSize (" ++ show vsSize ++ ") /= esSize (" ++ show esSize ++ ")"
+  when (esSize /= fsSize) $ do
+    error $ "esSize (" ++ show esSize ++ ") /= fsSize (" ++ show fsSize ++ ")"
+  vs <- cloneMutablePrimArray  vs0 0 vsSize
+  es <- cloneSmallMutableArray es0 0 esSize
+  fs <- cloneSmallMutableArray fs0 0 fsSize
+  pure $ StorageData storageSize vs es fs
+
+-- | Restore a shallow copy of the 'StorageData'.
+--
+-- The copy needs to be from the same 'Env' as the target.
+--
+-- @since 2.5.0.0
+restoreStorageData :: HasCallStack => StorageData -> Env es -> IO ()
+restoreStorageData newStorageData env = do
+  modifyIORef' (envStorage env) $ \(Storage version oldStorageData) ->
+    let oldSize = sdSize oldStorageData
+        newSize = sdSize newStorageData
+    in if newSize /= oldSize
+    then error $ "newSize (" ++ show newSize ++ ") /= oldSize (" ++ show oldSize ++ ")"
+    else Storage version newStorageData
+
 ----------------------------------------
 -- Relinker
 
@@ -205,18 +251,9 @@ emptyEnv = Env 0
 -- | Clone the environment to use it in a different thread.
 cloneEnv :: HasCallStack => Env es -> IO (Env es)
 cloneEnv (Env offset refs storage0) = do
-  Storage storageSize version vs0 es0 fs0 <- readIORef' storage0
-  vsSize <- getSizeofMutablePrimArray  vs0
-  esSize <- getSizeofSmallMutableArray es0
-  fsSize <- getSizeofSmallMutableArray fs0
-  when (vsSize /= esSize) $ do
-    error $ "vsSize (" ++ show vsSize ++ ") /= esSize (" ++ show esSize ++ ")"
-  when (esSize /= fsSize) $ do
-    error $ "esSize (" ++ show esSize ++ ") /= fsSize (" ++ show fsSize ++ ")"
-  vs <- cloneMutablePrimArray  vs0 0 vsSize
-  es <- cloneSmallMutableArray es0 0 esSize
-  fs <- cloneSmallMutableArray fs0 0 fsSize
-  storage <- newIORef' $ Storage storageSize version vs es fs
+  Storage version storageData0 <- readIORef' storage0
+  storageData@(StorageData storageSize _ es fs) <- copyStorageData storageData0
+  storage <- newIORef' $ Storage version storageData
   let relinkEffects = \case
         0 -> pure ()
         k -> do
@@ -230,29 +267,6 @@ cloneEnv (Env offset refs storage0) = do
   pure $ Env offset refs storage
 {-# NOINLINE cloneEnv #-}
 
--- | Restore the environment from its clone.
---
--- @since 2.2.0.0
-restoreEnv
-  :: HasCallStack
-  => Env es -- ^ Destination.
-  -> Env es -- ^ Source.
-  -> IO ()
-restoreEnv dest src = do
-  destStorage <- readIORef' (envStorage dest)
-  srcStorage  <- readIORef' (envStorage src)
-  let destStorageSize = stSize destStorage
-      srcStorageSize  = stSize srcStorage
-  when (destStorageSize /= srcStorageSize) $ do
-    error $ "destStorageSize (" ++ show destStorageSize
-         ++ ") /= srcStorageSize (" ++ show srcStorageSize ++ ")"
-  writeIORef' (envStorage dest) $ srcStorage
-    -- Decreasing the counter allows leakage of unsafeCoerce (see unsafeCoerce2
-    -- in the EnvTests module).
-    { stVersion = max (stVersion destStorage) (stVersion srcStorage)
-    }
-{-# NOINLINE restoreEnv #-}
-
 -- | Get the current size of the environment.
 sizeEnv :: Env es -> IO Int
 sizeEnv (Env offset refs _) = do
@@ -413,7 +427,7 @@ getLocation
   => Env es
   -> IO (Int, SmallMutableArray RealWorld AnyEffect)
 getLocation (Env offset refs storage) = do
-  Storage _ _ vs es _ <- readIORef' storage
+  Storage _ (StorageData _ vs es _) <- readIORef' storage
   storageVersion <- readPrimArray vs ref
   -- If version of the reference is different than version in the storage, it
   -- means that the effect in the storage is not the one that was initially
@@ -433,10 +447,12 @@ getLocation (Env offset refs storage) = do
 
 -- | Create an empty storage.
 emptyStorage :: HasCallStack => IO Storage
-emptyStorage = Storage 0 initialVersion
-  <$> newPrimArray 0
-  <*> newSmallArray 0 undefinedEffect
-  <*> newSmallArray 0 undefinedRelinker
+emptyStorage = Storage initialVersion <$> storageData
+  where
+    storageData = StorageData 0
+      <$> newPrimArray 0
+      <*> newSmallArray 0 undefinedEffect
+      <*> newSmallArray 0 undefinedRelinker
 
 -- | Insert an effect into the storage and return its reference.
 insertEffect
@@ -447,15 +463,16 @@ insertEffect
   -> Relinker (EffectRep (DispatchOf e)) e
   -> IO Ref
 insertEffect storage e f = do
-  Storage size version vs0 es0 fs0 <- readIORef' storage
+  Storage version (StorageData size vs0 es0 fs0) <- readIORef' storage
   len0 <- getSizeofSmallMutableArray es0
   case size `compare` len0 of
     GT -> error $ "size (" ++ show size ++ ") > len0 (" ++ show len0 ++ ")"
     LT -> do
       writePrimArray   vs0 size version
       writeSmallArray' es0 size (toAnyEffect e)
       writeSmallArray' fs0 size (toAnyRelinker f)
-      writeIORef' storage $ Storage (size + 1) (bumpVersion version) vs0 es0 fs0
+      writeIORef' storage $
+        Storage (bumpVersion version) (StorageData (size + 1) vs0 es0 fs0)
       pure $ Ref size version
     EQ -> do
       let len = doubleCapacity len0
@@ -468,14 +485,15 @@ insertEffect storage e f = do
       writePrimArray   vs size version
       writeSmallArray' es size (toAnyEffect e)
       writeSmallArray' fs size (toAnyRelinker f)
-      writeIORef' storage $ Storage (size + 1) (bumpVersion version) vs es fs
+      writeIORef' storage $
+        Storage (bumpVersion version) (StorageData (size + 1) vs es fs)
       pure $ Ref size version
 
 -- | Given a reference to an effect from the top of the stack, delete it from
 -- the storage.
 deleteEffect :: HasCallStack => IORef' Storage -> Ref -> IO ()
 deleteEffect storage (Ref ref version) = do
-  Storage size currentVersion vs es fs <- readIORef' storage
+  Storage currentVersion (StorageData size vs es fs) <- readIORef' storage
   when (ref /= size - 1) $ do
     error $ "ref (" ++ show ref ++ ") /= size - 1 (" ++ show (size - 1) ++ ")"
   storageVersion <- readPrimArray vs ref
@@ -485,7 +503,7 @@ deleteEffect storage (Ref ref version) = do
   writePrimArray  vs ref undefinedVersion
   writeSmallArray es ref undefinedEffect
   writeSmallArray fs ref undefinedRelinker
-  writeIORef' storage $ Storage (size - 1) currentVersion vs es fs
+  writeIORef' storage $ Storage currentVersion (StorageData (size - 1) vs es fs)
 
 -- | Relink the environment to use the new storage.
 relinkEnv :: IORef' Storage -> Env es -> IO (Env es)

effectful-core/src/Effectful/Internal/Monad.hs

@@ -553,6 +553,8 @@ send
   -> Eff es a
 send op = unsafeEff $ \es -> do
   Handler handlerEs handler <- getEnv es
+  when (envStorage es /= envStorage handlerEs) $ do
+    error "es and handlerEs point to different Storages"
   -- Prevent internal functions that rebind the effect handler from polluting
   -- its call stack by freezing it. Note that functions 'interpret',
   -- 'reinterpret', 'interpose' and 'impose' need to thaw it so that useful

effectful-core/src/Effectful/NonDet.hs

@@ -22,7 +22,7 @@ module Effectful.NonDet
   ) where
 
 import Control.Applicative
-import Data.Coerce
+import Data.IORef.Strict
 import GHC.Generics
 import GHC.Stack
 
@@ -31,13 +31,14 @@ import Effectful.Dispatch.Dynamic
 import Effectful.Dispatch.Static
 import Effectful.Dispatch.Static.Primitive
 import Effectful.Error.Static
-import Effectful.Internal.Monad (LocalEnv(..), NonDet(..))
+import Effectful.Internal.Env qualified as I
+import Effectful.Internal.Monad (NonDet(..))
 
 -- | Policy of dealing with modifications to __thread local__ state in the
 -- environment in branches that end up calling the 'Empty' operation.
 --
--- /Note:/ 'OnEmptyKeep' is significantly faster as there is no need to back up
--- the environment on each call to ':<|>:'.
+-- /Note:/ 'OnEmptyKeep' is faster as there is no need to back up the
+-- environment on each call to ':<|>:'.
 --
 -- @since 2.2.0.0
 data OnEmptyPolicy
@@ -84,23 +85,23 @@ runNonDetRollback
 runNonDetRollback = reinterpret setup $ \env -> \case
   Empty       -> throwError ErrorEmpty
   m1 :<|>: m2 -> do
-    backupEnv <- cloneLocalEnv env
+    backupData <- unsafeEff backupStorageData
     localSeqUnlift env $ \unlift -> do
       mr <- (Just <$> unlift m1) `catchError` \_ ErrorEmpty -> do
-        -- If m1 failed, roll back the environment.
-        restoreLocalEnv env backupEnv
+        -- If m1 failed, restore the data.
+        unsafeEff $ I.restoreStorageData backupData
         pure Nothing
       case mr of
         Just r  -> pure r
         Nothing -> unlift m2
   where
     setup action = do
-      backupEs <- unsafeEff cloneEnv
+      backupData <- unsafeEff backupStorageData
       runError @ErrorEmpty action >>= \case
         Right r -> pure $ Right r
         Left (cs, _) -> do
-          -- If the whole action failed, roll back the environment.
-          unsafeEff $ \es -> restoreEnv es backupEs
+          -- If the whole action failed, restore the data.
+          unsafeEff $ I.restoreStorageData backupData
           pure $ Left cs
 
 ----------------------------------------
@@ -139,15 +140,5 @@ instance Show ErrorEmpty where
 noError :: Either (cs, e) a -> Either cs a
 noError = either (Left . fst) Right
 
-cloneLocalEnv
-  :: HasCallStack
-  => LocalEnv localEs handlerEs
-  -> Eff es (LocalEnv localEs handlerEs)
-cloneLocalEnv = coerce . unsafeEff_ . cloneEnv . coerce
-
-restoreLocalEnv
-  :: HasCallStack
-  => LocalEnv localEs handlerEs
-  -> LocalEnv localEs handlerEs
-  -> Eff es ()
-restoreLocalEnv dest src = unsafeEff_ $ restoreEnv (coerce dest) (coerce src)
+backupStorageData :: HasCallStack => Env es -> IO I.StorageData
+backupStorageData env = I.copyStorageData . I.stData =<< readIORef' (I.envStorage env)

effectful/CHANGELOG.md

@@ -1,6 +1,11 @@
 # effectful-2.5.0.0 (????-??-??)
 * Add `plusEff` (specialized version of `<|>`) to `Effectful.NonDet` and make
   `emptyEff` and `sumEff` generate better call stacks.
+* **Bugfixes**:
+  - `OnEmptyRollback` strategy of the `NonDet` effect is no longer broken.
+* **Breaking changes**:
+  - Removed `restoreEnv` function from `Effectful.Dispatch.Static.Primitive`
+    since it was broken.
 
 # effectful-2.4.0.0 (2024-10-08)
 * Add utility functions for handling effects that take the effect handler as the

effectful/tests/EnvTests.hs

@@ -104,8 +104,6 @@ test_noUnsafeCoerce :: Assertion
 test_noUnsafeCoerce = do
   r1 <- try @ErrorCall . evaluate $ unsafeCoerce1 @Int 'a'
   assertBool "unsafeCoerce1" (isLeft r1)
-  r2 <- try @ErrorCall . evaluate $ unsafeCoerce2 @Int 'a'
-  assertBool "unsafeCoerce2" (isLeft r2)
 
 unsafeCoerce1 :: forall b a. a -> b
 unsafeCoerce1 a = runPureEff $ do
@@ -122,26 +120,6 @@ unsafeCoerce1 a = runPureEff $ do
   Box b <- runReader (Box a) $ raise oops
   pure b
 
-unsafeCoerce2 :: forall b a. a -> b
-unsafeCoerce2 a = runPureEff $ do
-  backupEs <- unsafeEff cloneEnv
-  -- 'oops' gains access to the effect stack with Reader (Box b) via the
-  -- unlifting function that escaped its scope. The problem here is that this
-  -- effect is no longer in scope.
-  oops <- runReader @(Box b) (Box undefined) $ do
-    raiseWith SeqUnlift $ \unlift -> do
-      pure . unlift $ ask @(Box b)
-  -- If restoreEnv messes up versioning (i.e. restores the version counter from
-  -- before the above Reader was used), below code will succeed because the
-  -- Reader (Box a) will have the same version as Reader (Box b).
-  unsafeEff $ \es -> restoreEnv es backupEs
-  -- Put Reader (Box a) where the Reader (Box b) was before and attempt to
-  -- retrieve 'a' coerced to 'b'. It should fail because 'getLocation' in
-  -- 'Effectful.Internal.Env' checks that version of the reference is the same
-  -- as version of the effect.
-  Box b <- runReader (Box a) $ raise oops
-  pure b
-
 data Box a = Box a
 
 ----------------------------------------

effectful/tests/StateTests.hs

@@ -51,7 +51,7 @@ test_stateM = runEff $ do
   U.assertEqual "correct a" "hi"    a
   U.assertEqual "correct b" "hi!!!" b
   where
-    getEffectReps = unsafeEff $ \es -> stEffects <$> readIORef' (envStorage es)
+    getEffectReps = unsafeEff $ \es -> sdEffects . stData <$> readIORef' (envStorage es)
 
 test_deepStack :: Assertion
 test_deepStack = runEff $ do