WIP: sketch out a reference API

src-control/Control/RefCount.hs

@@ -1,9 +1,35 @@
-{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE CPP               #-}
+{-# LANGUAGE DefaultSignatures #-}
+{-# LANGUAGE MagicHash         #-}
+{-# LANGUAGE TypeFamilies      #-}
 
 {- HLINT ignore "Evaluate" -}
 
 module Control.RefCount (
-    RefCounter (..)
+    -- * Using references
+    Ref
+  , releaseRef
+  , withRef
+  , unsafeDeRef
+    -- ** Shared references
+  , dupRef
+    -- ** Weak references
+  , WeakRef
+  , mkWeakRef
+  , deRefWeak
+    -- * Implementing objects with finalisers
+  , HasFinaliser (..)
+  , HasSharedFinaliser (..)
+  , newRef
+  , newSharedRef
+    -- ** Low level reference counts
+  , RawRefCounter (..)
+  , newRawRefCounter
+  , incrementRefCounter
+  , decrementRefCounter
+  , tryIncrementRefCounter
+
+    -- * Old API
   , RefCount (..)
   , unsafeMkRefCounterN
   , mkRefCounterN
@@ -14,15 +40,23 @@
   , readRefCount
   ) where
 
+import           Data.Kind (Type)
+import           Data.Maybe
+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.Maybe
-import           Data.Primitive.PrimVar
+
 import           GHC.Stack
 
+#ifdef NO_IGNORE_ASSERTS
+import           Data.IORef
+import           System.IO.Unsafe (unsafeDupablePerformIO)
+#endif
+
 -- | A reference counter with an optional finaliser action. Once the reference
 -- count reaches @0@, the finaliser will be run.
 data RefCounter m = RefCounter {
@@ -131,3 +165,231 @@
 -- compiling with @-O@ or @-fignore-asserts@.
 assertWithCallStack :: HasCallStack => Bool -> a -> a
 assertWithCallStack b = assert (const b callStack)
+
+-------------------------------------------------------------------------------
+-- New API
+--
+
+newtype RawRefCounter m = RawRefCounter (PrimVar (PrimState m) Int)
+  deriving stock Eq
+
+newRawRefCounter :: PrimMonad m => m (RawRefCounter m)
+newRawRefCounter = do
+    countVar <- newPrimVar 1
+    return $! RawRefCounter countVar
+
+incrementRefCounter :: PrimMonad m => RawRefCounter m -> m ()
+incrementRefCounter (RawRefCounter countVar) = do
+    prevCount <- fetchAddInt countVar 1
+    assert (prevCount > 0) $ pure ()
+
+decrementRefCounter :: PrimMonad m => RawRefCounter m -> m Bool
+decrementRefCounter (RawRefCounter countVar) = do
+    prevCount <- fetchSubInt countVar 1
+    assert (prevCount > 0) $ pure $! prevCount == 1
+
+tryIncrementRefCounter :: PrimMonad m => RawRefCounter m -> m Bool
+tryIncrementRefCounter (RawRefCounter countVar) = do
+    prevCount <- atomicReadInt countVar
+    casLoop prevCount
+  where
+    -- A classic lock-free CAS loop.
+    -- Check the value before is non-zero, return failure or continue.
+    -- Atomically write the new (incremented) value if the old value is
+    -- unchanged, and return the old value (either way).
+    -- If no other thread changed the old value, we succeed.
+    -- Otherwise we go round the loop again.
+    casLoop prevCount
+      | prevCount <= 0 = return False
+      | otherwise      = do
+          prevCount' <- casInt countVar prevCount (prevCount+1)
+          if prevCount' == prevCount
+            then return True
+            else casLoop prevCount'
+
+
+-- | A reference to an object of type @a@. Use references to support prompt
+-- finalisation of object resources.
+--
+-- The reference API works relatively uniformly, both for objects that only
+-- expect to have a single reference to them, and also objects supporting
+-- reference counting that may have many references.
+--
+-- Rules of use:
+--
+-- * Each 'Ref' must eventually be released /exactly/ once with 'releaseRef'.
+-- * Use 'withRef', or 'unsafeDeRef' to (temporarily) obtain the underlying
+--   object.
+-- * After calling 'releaseRef', the operations 'withRef' and 'unsafeDeRef'
+--   must /not/ be used.
+-- * After calling 'releaseRef', any value obtained previously from
+--   'unsafeDeRef' must /not/ be used. For this reason, it is advisable to use
+--   'withRef' where possible, and be careful with use of 'unsafeDeRef'.
+--
+-- Additionally, for objects in the 'HasSharedFinaliser' class:
+--
+-- * A 'Ref' may be duplicated using 'dupRef' to produce an independent
+--   reference (which must itself be released with 'releaseRef').
+--
+-- Provided that all these rules are followed, this guarantees that the
+-- object's finaliser will be run exactly once, promptly. For objects in the
+-- 'HasSharedFinaliser' class, this is when the final reference is released.
+--
+-- In debug mode (when using CPP define @NO_IGNORE_ASSERTS@), adherence to
+-- these rules is checked dynamically.
+#ifndef NO_IGNORE_ASSERTS
+newtype Ref obj = Ref obj
+#else
+data Ref obj = Ref !obj !(IORef Bool)
+#endif
+
+#ifdef NO_IGNORE_ASSERTS
+assertNotReleased :: PrimMonad m => IORef Bool -> m ()
+assertNotReleased released = do
+    r <- unsafeIOToPrim $ readIORef released
+    assert (not r) $ return ()
+#endif
+
+-- | Class of objects which support 'Ref'.
+--
+-- For objects in this class but not 'HasSharedFinaliser', the guarantee is
+-- that (when the 'Ref' rules are followed) the 'unsharedFinaliser' is called
+-- exactly once.
+--
+-- For objects that are also in 'HasSharedFinaliser', the 'unsharedFinaliser'
+-- is repurposed for reference counting, but the guarantee instead is that the
+-- object's 'sharedFinaliser' will be called exactly once.
+--
+class HasFinaliser obj where
+  type FinaliserM obj :: Type -> Type
+  unsharedFinaliser :: FinaliserM obj ~ m => obj -> m ()
+
+  default unsharedFinaliser :: (HasSharedFinaliser obj, FinaliserM obj ~ m,
+                                PrimMonad m, MonadMask m)
+                            => obj -> m ()
+  unsharedFinaliser obj =
+    mask_ $ do
+      done <- decrementRefCounter (getRefCounter obj)
+      when done (sharedFinaliser obj)
+
+-- | Class of objects that support 'Ref' including 'dupRef', with the guarantee
+-- that the object's 'sharedFinaliser' will be called exactly once.
+--
+-- Do not define 'unsharedFinaliser' for types in this class. The default
+-- implementation is used for the reference counting.
+--
+class HasSharedFinaliser obj where
+  getRefCounter   :: obj -> RawRefCounter (FinaliserM obj)
+
+  sharedFinaliser :: FinaliserM obj ~ m => obj -> m ()
+
+newRef :: PrimMonad m => obj -> m (Ref obj)
+#ifndef NO_IGNORE_ASSERTS
+newRef obj = return (Ref obj)
+#else
+newRef obj = do
+    released <- unsafeIOToPrim $ newIORef False
+    return (Ref obj released)
+#endif
+
+-- | Convenience helper to make an object containing a reference count.
+newSharedRef :: (HasSharedFinaliser obj, FinaliserM obj ~ m, PrimMonad m)
+             => (RawRefCounter m -> obj) -> m (Ref obj)
+newSharedRef mkObject = do
+    rc  <- newRawRefCounter
+    let !obj = mkObject rc
+    assert (getRefCounter obj == rc) $
+      newRef obj
+
+-- | Release a reference to an object that will no longer be used (via this
+-- reference).
+--
+releaseRef :: (HasFinaliser obj, FinaliserM obj ~ m, PrimMonad m)
+           => Ref obj -> m ()
+#ifndef NO_IGNORE_ASSERTS
+releaseRef (Ref obj) = unsharedFinaliser obj
+#else
+releaseRef (Ref obj released) = do
+    assertNotReleased released
+    unsharedFinaliser obj
+    unsafeIOToPrim $ writeIORef released True
+#endif
+
+-- | 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 'unsafeDeRef'.
+--
+#ifndef NO_IGNORE_ASSERTS
+withRef :: Ref obj -> (obj -> m a) -> m a
+withRef (Ref obj) f = f obj
+#else
+withRef :: PrimMonad m => Ref obj -> (obj -> m a) -> m a
+withRef (Ref obj released) f = assertNotReleased released
+                            >> f obj
+#endif
+
+-- | Get the object in a 'Ref'. Be careful with retaining the object for too
+-- long, since the object must not be used after 'releaseRef' is called.
+--
+unsafeDeRef :: Ref obj -> obj
+#ifndef NO_IGNORE_ASSERTS
+unsafeDeRef (Ref obj) = obj
+#else
+unsafeDeRef (Ref obj released) =
+          unsafeDupablePerformIO (assertNotReleased released)
+    `seq` obj
+#endif
+
+-- | Duplicate an existing reference, to produce a new reference.
+--
+dupRef :: (HasSharedFinaliser obj, FinaliserM obj ~ m, PrimMonad m)
+       => Ref obj -> m (Ref obj)
+#ifndef NO_IGNORE_ASSERTS
+dupRef (Ref obj) = do
+    incrementRefCounter (getRefCounter obj)
+    return (Ref obj)
+#else
+dupRef (Ref obj released) = do
+    assertNotReleased released
+    incrementRefCounter (getRefCounter obj)
+    released' <- unsafeIOToPrim $ newIORef False
+    return (Ref obj released')
+#endif
+
+
+-- | A \"weak\" reference to an object: that is, a reference that does not
+-- guarantee to keep the object alive. If however the object is still alive
+-- (due to other normal references still existing) then it can be converted
+-- back into a normal reference with 'deRefWeak'.
+--
+-- Weak references do not themselves need to be released.
+--
+newtype WeakRef a = WeakRef a
+
+-- | Given an existing normal reference, create a new weak reference.
+--
+mkWeakRef :: Ref obj -> WeakRef obj
+#ifndef NO_IGNORE_ASSERTS
+mkWeakRef (Ref obj) = WeakRef obj
+#else
+mkWeakRef (Ref obj _) = WeakRef obj
+#endif
+
+-- | If the object is still alive, obtain a /new/ normal reference. The normal
+-- rules for 'Ref' apply, including the need to eventually calling 'releaseRef'.
+--
+deRefWeak :: (HasSharedFinaliser obj, FinaliserM obj ~ m, PrimMonad m)
+          => WeakRef obj -> m (Maybe (Ref obj))
+#ifndef NO_IGNORE_ASSERTS
+deRefWeak (WeakRef obj) = do
+    success <- tryIncrementRefCounter (getRefCounter obj)
+    if success then return (Just (Ref obj))
+               else return Nothing
+#else
+deRefWeak (WeakRef obj) = do
+    success <- tryIncrementRefCounter (getRefCounter obj)
+    if success
+      then do released <- unsafeIOToPrim $ newIORef False
+              return (Just (Ref obj released))
+      else return Nothing
+#endif
+