Skip to content

Commit

Permalink
Fix NonDet/OnEmptyRollback (#256)
Browse files Browse the repository at this point in the history
  • Loading branch information
arybczak authored Oct 10, 2024
1 parent df06ee3 commit 5bda849
Show file tree
Hide file tree
Showing 8 changed files with 98 additions and 100 deletions.
5 changes: 5 additions & 0 deletions effectful-core/CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -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
Expand Down
1 change: 0 additions & 1 deletion effectful-core/src/Effectful/Dispatch/Static/Primitive.hs
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,6 @@ module Effectful.Dispatch.Static.Primitive
-- ** Utils
, emptyEnv
, cloneEnv
, restoreEnv
, sizeEnv
, tailEnv
) where
Expand Down
126 changes: 72 additions & 54 deletions effectful-core/src/Effectful/Internal/Env.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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(..)
Expand All @@ -29,7 +36,6 @@ module Effectful.Internal.Env
-- * Operations
, emptyEnv
, cloneEnv
, restoreEnv
, sizeEnv
, tailEnv

Expand Down Expand Up @@ -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

Expand All @@ -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

Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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)
Expand Down
2 changes: 2 additions & 0 deletions effectful-core/src/Effectful/Internal/Monad.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
35 changes: 13 additions & 22 deletions effectful-core/src/Effectful/NonDet.hs
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ module Effectful.NonDet
) where

import Control.Applicative
import Data.Coerce
import Data.IORef.Strict
import GHC.Generics
import GHC.Stack

Expand All @@ -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
Expand Down Expand Up @@ -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

----------------------------------------
Expand Down Expand Up @@ -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)
5 changes: 5 additions & 0 deletions effectful/CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -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
Expand Down
22 changes: 0 additions & 22 deletions effectful/tests/EnvTests.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

----------------------------------------
Expand Down
2 changes: 1 addition & 1 deletion effectful/tests/StateTests.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 5bda849

Please sign in to comment.