Skip to content

Commit

Permalink
Merge pull request #494 from IntersectMBO/wenkokke/rc-fundep
Browse files Browse the repository at this point in the history
Use functional dependencies rather than type families to enable SPECIALISE
  • Loading branch information
dcoutts authored Dec 12, 2024
2 parents 21f4bf7 + 2eeabd6 commit 9b4d561
Show file tree
Hide file tree
Showing 6 changed files with 57 additions and 35 deletions.
74 changes: 51 additions & 23 deletions src-control/Control/RefCount.hs
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
{-# LANGUAGE CPP #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE TypeFamilies #-}

module Control.RefCount (
-- * Using references
Expand Down Expand Up @@ -30,15 +31,12 @@ module Control.RefCount (
, checkForgottenRefs
) where

import Data.Kind (Type)
import Data.Primitive.PrimVar

import Control.DeepSeq
import Control.Exception (assert)
import Control.Monad (when)
import Control.Monad.Class.MonadThrow
import Control.Monad.Primitive

import Data.Primitive.PrimVar
import GHC.Stack (CallStack, prettyCallStack)

#ifdef NO_IGNORE_ASSERTS
Expand Down Expand Up @@ -195,30 +193,36 @@ instance NFData obj => NFData (Ref obj) where
-- For objects in this class the guarantee is that (when the 'Ref' rules are
-- followed) the object's finaliser is called exactly once.
--
class RefCounted obj where
type FinaliserM obj :: Type -> Type
getRefCounter :: obj -> RefCounter (FinaliserM obj)
class RefCounted m obj | obj -> m where
getRefCounter :: obj -> RefCounter m

#ifdef NO_IGNORE_ASSERTS
#define HasCallStackIfDebug HasCallStack
#else
#define HasCallStackIfDebug ()
#endif

-- GHC says specialising is too complicated! But it's ok, each of these can
-- inline to calling a few other specialised helpers.
{-# INLINE newRef #-}
{-# INLINE releaseRef #-}
{-# INLINE dupRef #-}
{-# INLINE deRefWeak #-}

{-# SPECIALISE
newRef ::
RefCounted IO obj
=> IO ()
-> (RefCounter IO -> obj)
-> IO (Ref obj)
#-}
-- | Make a new reference.
--
-- The given finaliser is run when the last reference is released. The
-- finaliser is run with async exceptions masked.
--
{-# SPECIALISE
newRef ::
RefCounted IO obj
=> IO ()
-> (RefCounter IO -> obj)
-> IO (Ref obj)
#-}
newRef ::
(RefCounted obj, FinaliserM obj ~ m, PrimMonad m)
(RefCounted m obj, PrimMonad m)
=> HasCallStackIfDebug
=> m ()
-> (RefCounter m -> obj)
Expand All @@ -232,8 +236,14 @@ newRef finaliser mkObject = do
-- | Release a reference to an object that will no longer be used (via this
-- reference).
--
{-# SPECIALISE
releaseRef ::
RefCounted IO obj
=> Ref obj
-> IO ()
#-}
releaseRef ::
(RefCounted obj, FinaliserM obj ~ m, PrimMonad m, MonadMask m)
(RefCounted m obj, PrimMonad m, MonadMask m)
=> HasCallStackIfDebug
=> Ref obj
-> m ()
Expand Down Expand Up @@ -261,6 +271,12 @@ deRef ref@Ref{refobj} =
`seq` refobj
#endif

{-# SPECIALISE
withRef ::
Ref obj
-> (obj -> IO a)
-> IO a
#-}
{-# INLINE withRef #-}
-- | Use the object in a 'Ref'. Do not retain the object after the scope of
-- the body. If you cannot use scoped \"with\" style, use pattern 'DeRef'.
Expand All @@ -276,10 +292,16 @@ withRef ref@Ref{refobj} f = do
assertNoUseAfterRelease ref
f refobj

{-# SPECIALISE
dupRef ::
RefCounted IO obj
=> Ref obj
-> IO (Ref obj)
#-}
-- | Duplicate an existing reference, to produce a new reference.
--
dupRef ::
(RefCounted obj, FinaliserM obj ~ m, PrimMonad m)
(RefCounted m obj, PrimMonad m)
=> HasCallStackIfDebug
=> Ref obj
-> m (Ref obj)
Expand Down Expand Up @@ -308,11 +330,17 @@ mkWeakRef Ref {refobj} = WeakRef refobj
mkWeakRefFromRaw :: obj -> WeakRef obj
mkWeakRefFromRaw obj = WeakRef obj

{-# SPECIALISE
deRefWeak ::
RefCounted IO obj
=> WeakRef obj
-> IO (Maybe (Ref obj))
#-}
-- | If the object is still alive, obtain a /new/ normal reference. The normal
-- rules for 'Ref' apply, including the need to eventually call 'releaseRef'.
--
deRefWeak ::
(RefCounted obj, FinaliserM obj ~ m, PrimMonad m)
(RefCounted m obj, PrimMonad m)
=> HasCallStackIfDebug
=> WeakRef obj
-> m (Maybe (Ref obj))
Expand Down
3 changes: 1 addition & 2 deletions src/Database/LSMTree/Internal/BlobFile.hs
Original file line number Diff line number Diff line change
Expand Up @@ -32,8 +32,7 @@ data BlobFile m h = BlobFile {
}
deriving stock (Show)

instance RefCounted (BlobFile m h) where
type FinaliserM (BlobFile m h) = m
instance RefCounted m (BlobFile m h) where
getRefCounter = blobFileRefCounter

instance NFData h => NFData (BlobFile m h) where
Expand Down
3 changes: 1 addition & 2 deletions src/Database/LSMTree/Internal/MergeSchedule.hs
Original file line number Diff line number Diff line change
Expand Up @@ -325,8 +325,7 @@ data MergingRun m h = MergingRun {
, mergeRefCounter :: !(RefCounter m)
}

instance RefCounted (MergingRun m h) where
type FinaliserM (MergingRun m h) = m
instance RefCounted m (MergingRun m h) where
getRefCounter = mergeRefCounter

{-# SPECIALISE newMergingRun ::
Expand Down
3 changes: 1 addition & 2 deletions src/Database/LSMTree/Internal/Run.hs
Original file line number Diff line number Diff line change
Expand Up @@ -101,8 +101,7 @@ instance NFData h => NFData (Run m h) where
rnf a `seq` rwhnf b `seq` rnf c `seq` rnf d `seq` rnf e `seq`
rnf f `seq` rnf g `seq` rnf h `seq` rwhnf i `seq` rwhnf j

instance RefCounted (Run m h) where
type FinaliserM (Run m h) = m
instance RefCounted m (Run m h) where
getRefCounter = runRefCounter

size :: Ref (Run m h) -> NumEntries
Expand Down
3 changes: 1 addition & 2 deletions src/Database/LSMTree/Internal/WriteBufferBlobs.hs
Original file line number Diff line number Diff line change
Expand Up @@ -112,8 +112,7 @@ data WriteBufferBlobs m h =
instance NFData h => NFData (WriteBufferBlobs m h) where
rnf (WriteBufferBlobs a b c) = rnf a `seq` rnf b `seq` rnf c

instance RefCounted (WriteBufferBlobs m h) where
type FinaliserM (WriteBufferBlobs m h) = m
instance RefCounted m (WriteBufferBlobs m h) where
getRefCounter = writeBufRefCounter

{-# SPECIALISE new :: HasFS IO h -> FS.FsPath -> IO (Ref (WriteBufferBlobs IO h)) #-}
Expand Down
6 changes: 2 additions & 4 deletions test-control/Test/Control/RefCount.hs
Original file line number Diff line number Diff line change
Expand Up @@ -108,14 +108,12 @@ readRefCount (RefCounter countVar _) = readPrimVar countVar
#ifdef NO_IGNORE_ASSERTS
data TestObject = TestObject !(RefCounter IO)

instance RefCounted TestObject where
type FinaliserM TestObject = IO
instance RefCounted IO TestObject where
getRefCounter (TestObject rc) = rc

data TestObject2 = TestObject2 (Ref TestObject)

instance RefCounted TestObject2 where
type FinaliserM TestObject2 = IO
instance RefCounted IO TestObject2 where
getRefCounter (TestObject2 (DeRef to1)) = getRefCounter to1

prop_ref_double_free :: Property
Expand Down

0 comments on commit 9b4d561

Please sign in to comment.