Make envRefs store Refs instead of juggling Ints

.github/workflows/haskell-ci.yml

@@ -248,11 +248,11 @@ jobs:
       - name: doctest
         run: |
           cd ${PKGDIR_effectful_core} || false
-          doctest  -XHaskell2010 -XBangPatterns -XConstraintKinds -XDataKinds -XDeriveFunctor -XDeriveGeneric -XFlexibleContexts -XFlexibleInstances -XGADTs -XGeneralizedNewtypeDeriving -XImportQualifiedPost -XLambdaCase -XMultiParamTypeClasses -XNoStarIsType -XPolyKinds -XRankNTypes -XRoleAnnotations -XScopedTypeVariables -XStandaloneDeriving -XTupleSections -XTypeApplications -XTypeFamilies -XTypeOperators src
+          doctest  -XHaskell2010 -XBangPatterns -XConstraintKinds -XDataKinds -XDeriveFunctor -XDeriveGeneric -XDerivingStrategies -XFlexibleContexts -XFlexibleInstances -XGADTs -XGeneralizedNewtypeDeriving -XImportQualifiedPost -XLambdaCase -XMultiParamTypeClasses -XNoStarIsType -XPolyKinds -XRankNTypes -XRoleAnnotations -XScopedTypeVariables -XStandaloneDeriving -XTupleSections -XTypeApplications -XTypeFamilies -XTypeOperators src
           cd ${PKGDIR_effectful_th} || false
-          doctest  -XHaskell2010 -XBangPatterns -XConstraintKinds -XDataKinds -XDeriveFunctor -XDeriveGeneric -XFlexibleContexts -XFlexibleInstances -XGADTs -XGeneralizedNewtypeDeriving -XImportQualifiedPost -XLambdaCase -XMultiParamTypeClasses -XNoStarIsType -XPolyKinds -XRankNTypes -XRecordWildCards -XRoleAnnotations -XScopedTypeVariables -XStandaloneDeriving -XTupleSections -XTypeApplications -XTypeFamilies -XTypeOperators src
+          doctest  -XHaskell2010 -XBangPatterns -XConstraintKinds -XDataKinds -XDeriveFunctor -XDeriveGeneric -XDerivingStrategies -XFlexibleContexts -XFlexibleInstances -XGADTs -XGeneralizedNewtypeDeriving -XImportQualifiedPost -XLambdaCase -XMultiParamTypeClasses -XNoStarIsType -XPolyKinds -XRankNTypes -XRecordWildCards -XRoleAnnotations -XScopedTypeVariables -XStandaloneDeriving -XTupleSections -XTypeApplications -XTypeFamilies -XTypeOperators src
           cd ${PKGDIR_effectful} || false
-          doctest  -XHaskell2010 -XBangPatterns -XConstraintKinds -XDataKinds -XDeriveFunctor -XDeriveGeneric -XFlexibleContexts -XFlexibleInstances -XGADTs -XGeneralizedNewtypeDeriving -XImportQualifiedPost -XLambdaCase -XMultiParamTypeClasses -XNoStarIsType -XPolyKinds -XRankNTypes -XRecordWildCards -XRoleAnnotations -XScopedTypeVariables -XStandaloneDeriving -XTupleSections -XTypeApplications -XTypeFamilies -XTypeOperators src
+          doctest  -XHaskell2010 -XBangPatterns -XConstraintKinds -XDataKinds -XDeriveFunctor -XDeriveGeneric -XDerivingStrategies -XFlexibleContexts -XFlexibleInstances -XGADTs -XGeneralizedNewtypeDeriving -XImportQualifiedPost -XLambdaCase -XMultiParamTypeClasses -XNoStarIsType -XPolyKinds -XRankNTypes -XRecordWildCards -XRoleAnnotations -XScopedTypeVariables -XStandaloneDeriving -XTupleSections -XTypeApplications -XTypeFamilies -XTypeOperators src
       - name: cabal check
         run: |
           cd ${PKGDIR_effectful_core} || false

doctest.sh

@@ -19,6 +19,7 @@ run_doctest() {
     -XDataKinds \
     -XDeriveFunctor \
     -XDeriveGeneric \
+    -XDerivingStrategies \
     -XFlexibleContexts \
     -XFlexibleInstances \
     -XGADTs \

effectful-core/effectful-core.cabal

@@ -39,6 +39,7 @@ common language
                         DataKinds
                         DeriveFunctor
                         DeriveGeneric
+                        DerivingStrategies
                         FlexibleContexts
                         FlexibleInstances
                         GADTs

effectful-core/src/Effectful/Dispatch/Dynamic.hs

@@ -1096,16 +1096,14 @@ copyRefs src@(Env soffset srefs _) dest@(Env doffset drefs storage) = do
   requireMatchingStorages src dest
   let size = sizeofPrimArray drefs - doffset
       es = reifyIndices @es @srcEs
-      esSize = 2 * length es
+      esSize = length es
   mrefs <- newPrimArray (esSize + size)
   copyPrimArray mrefs esSize drefs doffset size
   let writeRefs i = \case
         [] -> pure ()
         (x : xs) -> do
-          let ix = soffset + 2 * x
+          writePrimArray mrefs i $ indexPrimArray srefs (soffset + x)
-          writePrimArray mrefs  i      $ indexPrimArray srefs  ix
+          writeRefs (i + 1) xs
-          writePrimArray mrefs (i + 1) $ indexPrimArray srefs (ix + 1)
-          writeRefs (i + 2) xs
   writeRefs 0 es
   refs <- unsafeFreezePrimArray mrefs
   pure $ Env 0 refs storage

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

@@ -1,9 +1,13 @@
 {-# LANGUAGE AllowAmbiguousTypes #-}
 {-# LANGUAGE CPP #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE UnboxedTuples #-}
 {-# OPTIONS_HADDOCK not-home #-}
 module Effectful.Internal.Env
   ( -- * The environment
     Env(..)
+  , Ref(..)
+  , Version
   , Storage(..)
   , AnyEffect
   , toAnyEffect
@@ -49,6 +53,8 @@ import Control.Monad.Primitive
 import Data.IORef.Strict
 import Data.Primitive.PrimArray
 import Data.Primitive.SmallArray
+import Data.Primitive.Types
+import GHC.Exts ((*#), (+#))
 import GHC.Stack
 
 import Effectful.Internal.Effect
@@ -81,15 +87,50 @@ type role Env nominal
 --
 data Env (es :: [Effect]) = Env
   { envOffset  :: !Int
-  , envRefs    :: !(PrimArray Int)
+  , envRefs    :: !(PrimArray Ref)
   , envStorage :: !(IORef' Storage)
   }
 
+-- | Reference to the effect in 'Storage'.
+data Ref = Ref !Int !Version
+
+instance Prim Ref where
+  sizeOf# _ = 2# *# sizeOf# (undefined :: Int)
+  alignment# _ = alignment# (undefined :: Int)
+  indexByteArray# arr# i# =
+    let ref = indexByteArray# arr# (2# *# i#)
+        version = indexByteArray# arr# (2# *# i# +# 1#)
+    in Ref ref version
+  readByteArray# arr# i# =
+    \s0 -> case readByteArray# arr# (2# *# i#) s0 of
+       (# s1#, ref #) -> case readByteArray# arr# (2# *# i# +# 1#) s1# of
+          (# s2#, version #) -> (# s2#, Ref ref version #)
+  writeByteArray# arr# i# (Ref ref version) =
+    \s0 -> case writeByteArray# arr# (2# *# i#) ref s0 of
+       s1 -> case writeByteArray# arr# (2# *# i# +# 1#) version s1 of
+         s2 -> s2
+  indexOffAddr# addr# i# =
+    let ref = indexOffAddr# addr# (2# *# i#)
+        version = indexOffAddr# addr# (2# *# i# +# 1#)
+    in Ref ref version
+  readOffAddr# addr# i# =
+    \s0 -> case readOffAddr# addr# (2# *# i#) s0 of
+       (# s1, ref #) -> case readOffAddr# addr# (2# *# i# +# 1#) s1 of
+         (# s2, version #) -> (# s2, Ref ref version #)
+  writeOffAddr# addr# i# (Ref ref version) =
+    \s0 -> case writeOffAddr# addr# (2# *# i#) ref s0 of
+       s1 -> case writeOffAddr# addr# (2# *# i# +# 1#) version s1 of
+         s2 -> s2
+
+-- | Version of the effect.
+newtype Version = Version Int
+  deriving newtype (Eq, Ord, Prim, Show)
+
 -- | A storage of effects.
 data Storage = Storage
   { stSize      :: !Int
-  , stVersion   :: !Int
+  , stVersion   :: !Version
-  , stVersions  :: !(MutablePrimArray RealWorld Int)
+  , stVersions  :: !(MutablePrimArray RealWorld Version)
   , stEffects   :: !(SmallMutableArray RealWorld AnyEffect)
   , stRelinkers :: !(SmallMutableArray RealWorld AnyRelinker)
   }
@@ -209,12 +250,12 @@ restoreEnv dest src = do
 -- | Get the current size of the environment.
 sizeEnv :: Env es -> IO Int
 sizeEnv (Env offset refs _) = do
-  pure $ (sizeofPrimArray refs - offset) `div` 2
+  pure $ sizeofPrimArray refs - offset
 
 -- | Access the tail of the environment.
 tailEnv :: Env (e : es) -> IO (Env es)
 tailEnv (Env offset refs storage) = do
-  pure $ Env (offset + 2) refs storage
+  pure $ Env (offset + 1) refs storage
 
 ----------------------------------------
 -- Extending and shrinking
@@ -229,11 +270,10 @@ consEnv
   -> IO (Env (e : es))
 consEnv e f (Env offset refs0 storage) = do
   let size = sizeofPrimArray refs0 - offset
-  mrefs <- newPrimArray (size + 2)
+  mrefs <- newPrimArray (size + 1)
-  copyPrimArray mrefs 2 refs0 offset size
+  copyPrimArray mrefs 1 refs0 offset size
-  (ref, version) <- insertEffect storage e f
+  ref <- insertEffect storage e f
   writePrimArray mrefs 0 ref
-  writePrimArray mrefs 1 version
   refs <- unsafeFreezePrimArray mrefs
   pure $ Env 0 refs storage
 {-# NOINLINE consEnv #-}
@@ -264,10 +304,8 @@ replaceEnv e f (Env offset refs0 storage) = do
   let size = sizeofPrimArray refs0 - offset
   mrefs <- newPrimArray size
   copyPrimArray mrefs 0 refs0 offset size
-  (ref, version) <- insertEffect storage e f
+  ref <- insertEffect storage e f
-  let i = 2 * reifyIndex @e @es
+  writePrimArray mrefs (reifyIndex @e @es) ref
-  writePrimArray mrefs  i      ref
-  writePrimArray mrefs (i + 1) version
   refs <- unsafeFreezePrimArray mrefs
   pure $ Env 0 refs storage
 {-# NOINLINE replaceEnv #-}
@@ -278,7 +316,7 @@ replaceEnv e f (Env offset refs0 storage) = do
 -- usable.
 unreplaceEnv :: forall e es. (HasCallStack, e :> es) => Env es -> IO ()
 unreplaceEnv (Env offset refs storage) = do
-  deleteEffect storage $ indexPrimArray refs (offset + 2 * reifyIndex @e @es)
+  deleteEffect storage $ indexPrimArray refs (offset + reifyIndex @e @es)
 {-# NOINLINE unreplaceEnv #-}
 
 ----------------------------------------
@@ -287,11 +325,9 @@ unreplaceEnv (Env offset refs storage) = do
 subsumeEnv :: forall e es. e :> es => Env es -> IO (Env (e : es))
 subsumeEnv (Env offset refs0 storage) = do
   let size = sizeofPrimArray refs0 - offset
-  mrefs <- newPrimArray (size + 2)
+  mrefs <- newPrimArray (size + 1)
-  copyPrimArray mrefs 2 refs0 offset size
+  copyPrimArray mrefs 1 refs0 offset size
-  let ix = offset + 2 * reifyIndex @e @es
+  writePrimArray mrefs 0 $ indexPrimArray refs0 (offset + reifyIndex @e @es)
-  writePrimArray mrefs 0 $ indexPrimArray refs0  ix
-  writePrimArray mrefs 1 $ indexPrimArray refs0 (ix + 1)
   refs <- unsafeFreezePrimArray mrefs
   pure $ Env 0 refs storage
 {-# NOINLINE subsumeEnv #-}
@@ -303,8 +339,8 @@ subsumeEnv (Env offset refs0 storage) = do
 injectEnv :: forall xs es. Subset xs es => Env es -> IO (Env xs)
 injectEnv (Env offset refs0 storage) = do
   let xs         = reifyIndices @xs @es
-      permSize   = 2 * length xs
+      permSize   = length xs
-      prefixSize = 2 * prefixLength @es
+      prefixSize = prefixLength @es
       suffixSize = if subsetFullyKnown @xs @es
                    then 0
                    else sizeofPrimArray refs0 - offset - prefixSize
@@ -313,10 +349,8 @@ injectEnv (Env offset refs0 storage) = do
   let writePermRefs i = \case
         []       -> pure ()
         (e : es) -> do
-          let ix = offset + 2 * e
+          writePrimArray mrefs i $ indexPrimArray refs0 (offset + e)
-          writePrimArray mrefs  i      $ indexPrimArray refs0  ix
+          writePermRefs (i + 1) es
-          writePrimArray mrefs (i + 1) $ indexPrimArray refs0 (ix + 1)
-          writePermRefs (i + 2) es
   writePermRefs 0 xs
   refs <- unsafeFreezePrimArray mrefs
   pure $ Env 0 refs storage
@@ -373,9 +407,6 @@ getLocation
   => Env es
   -> IO (Int, SmallMutableArray RealWorld AnyEffect)
 getLocation (Env offset refs storage) = do
-  let i       = offset + 2 * reifyIndex @e @es
-      ref     = indexPrimArray refs  i
-      version = indexPrimArray refs (i + 1)
   Storage _ _ vs es _ <- readIORef' storage
   storageVersion <- readPrimArray vs ref
   -- If version of the reference is different than version in the storage, it
@@ -388,13 +419,15 @@ getLocation (Env offset refs storage) = do
          ++ "of the scope of effects it captures, have a look at "
          ++ "UnliftingStrategy (SeqForkUnlift)."
   pure (ref, es)
+  where
+    Ref ref version = indexPrimArray refs (offset + reifyIndex @e @es)
 
 ----------------------------------------
 -- Internal helpers
 
 -- | Create an empty storage.
 emptyStorage :: HasCallStack => IO Storage
-emptyStorage = Storage 0 (noVersion + 1)
+emptyStorage = Storage 0 initialVersion
   <$> newPrimArray 0
   <*> newSmallArray 0 undefinedEffect
   <*> newSmallArray 0 undefinedRelinker
@@ -406,7 +439,7 @@ insertEffect
   -> EffectRep (DispatchOf e) e
   -- ^ The representation of the effect.
   -> Relinker (EffectRep (DispatchOf e)) e
-  -> IO (Int, Int)
+  -> IO Ref
 insertEffect storage e f = do
   Storage size version vs0 es0 fs0 <- readIORef' storage
   len0 <- getSizeofSmallMutableArray es0
@@ -416,8 +449,8 @@ insertEffect storage e f = do
       writePrimArray   vs0 size version
       writeSmallArray' es0 size (toAnyEffect e)
       writeSmallArray' fs0 size (toAnyRelinker f)
-      writeIORef' storage $ Storage (size + 1) (version + 1) vs0 es0 fs0
+      writeIORef' storage $ Storage (size + 1) (bumpVersion version) vs0 es0 fs0
-      pure (size, version)
+      pure $ Ref size version
     EQ -> do
       let len = doubleCapacity len0
       vs <- newPrimArray len
@@ -429,20 +462,24 @@ insertEffect storage e f = do
       writePrimArray   vs size version
       writeSmallArray' es size (toAnyEffect e)
       writeSmallArray' fs size (toAnyRelinker f)
-      writeIORef' storage $ Storage (size + 1) (version + 1) vs es fs
+      writeIORef' storage $ Storage (size + 1) (bumpVersion version) vs es fs
-      pure (size, version)
+      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 -> Int -> IO ()
+deleteEffect :: HasCallStack => IORef' Storage -> Ref -> IO ()
-deleteEffect storage ref = do
+deleteEffect storage (Ref ref version) = do
-  Storage size version vs es fs <- readIORef' storage
+  Storage size currentVersion vs es fs <- readIORef' storage
   when (ref /= size - 1) $ do
     error $ "ref (" ++ show ref ++ ") /= size - 1 (" ++ show (size - 1) ++ ")"
-  writePrimArray  vs ref noVersion
+  storageVersion <- readPrimArray vs ref
+  when (version /= storageVersion) $ do
+    error $ "version (" ++ show version ++ ") /= storageVersion ("
+         ++ show storageVersion ++ ")\n"
+  writePrimArray  vs ref undefinedVersion
   writeSmallArray es ref undefinedEffect
   writeSmallArray fs ref undefinedRelinker
-  writeIORef' storage $ Storage (size - 1) version vs es fs
+  writeIORef' storage $ Storage (size - 1) currentVersion vs es fs
 
 -- | Relink the environment to use the new storage.
 relinkEnv :: IORef' Storage -> Env es -> IO (Env es)
@@ -452,8 +489,14 @@ relinkEnv storage (Env offset refs _) = pure $ Env offset refs storage
 doubleCapacity :: Int -> Int
 doubleCapacity n = max 1 n * 2
 
-noVersion :: Int
+undefinedVersion :: Version
-noVersion = 0
+undefinedVersion = Version 0
+
+initialVersion :: Version
+initialVersion = Version 1
+
+bumpVersion :: Version -> Version
+bumpVersion (Version n) = Version (n + 1)
 
 undefinedEffect :: HasCallStack => AnyEffect
 undefinedEffect = toAnyEffect . errorWithoutStackTrace $ unlines

effectful-core/src/Effectful/Provider.hs

@@ -204,8 +204,8 @@ copyRef (Env hoffset hrefs hstorage) (Env offset refs0 storage) = do
   when (hstorage /= storage) $ do
     error "storages do not match"
   let size = sizeofPrimArray refs0 - offset
-  mrefs <- newPrimArray (size + 2)
+  mrefs <- newPrimArray (size + 1)
-  copyPrimArray mrefs 0 hrefs hoffset 2
+  copyPrimArray mrefs 0 hrefs hoffset 1
-  copyPrimArray mrefs 2 refs0 offset size
+  copyPrimArray mrefs 1 refs0 offset size
   refs <- unsafeFreezePrimArray mrefs
   pure $ Env 0 refs storage

effectful-core/src/Effectful/Provider/List.hs

@@ -148,7 +148,7 @@ copyRefs (Env hoffset hrefs hstorage) (Env offset refs0 storage) = do
   when (hstorage /= storage) $ do
     error "storages do not match"
   let size = sizeofPrimArray refs0 - offset
-      listSize = 2 * knownEffectsLength @list
+      listSize = knownEffectsLength @list
   mrefs <- newPrimArray (size + listSize)
   copyPrimArray mrefs 0 hrefs hoffset listSize
   copyPrimArray mrefs listSize refs0 offset size

effectful-plugin/effectful-plugin.cabal

@@ -35,6 +35,7 @@ common language
                         DataKinds
                         DeriveFunctor
                         DeriveGeneric
+                        DerivingStrategies
                         FlexibleContexts
                         FlexibleInstances
                         GADTs

effectful-th/effectful-th.cabal

@@ -35,6 +35,7 @@ common language
                         DataKinds
                         DeriveFunctor
                         DeriveGeneric
+                        DerivingStrategies
                         FlexibleContexts
                         FlexibleInstances
                         GADTs

effectful/effectful.cabal

@@ -44,6 +44,7 @@ common language
                         DataKinds
                         DeriveFunctor
                         DeriveGeneric
+                        DerivingStrategies
                         FlexibleContexts
                         FlexibleInstances
                         GADTs