Resolving TODO: PrimVar for UniqCounter

src-extras/Database/LSMTree/Extras/NoThunks.hs

@@ -158,7 +158,7 @@ deriving via AllowThunksIn ["cursorSession", "cursorSessionEnv"] (CursorEnv m h)
 -------------------------------------------------------------------------------}
 
 deriving stock instance Generic (UniqCounter m)
-deriving anyclass instance NoThunks (StrictMVar m Word64)
+deriving anyclass instance (NoThunks (PrimVar (PrimState m) Int))
                         => NoThunks (UniqCounter m)
 
 {-------------------------------------------------------------------------------

src/Database/LSMTree/Internal/Snapshot.hs

@@ -365,7 +365,7 @@ snapshotRuns reg hbio0 (NamedSnapshotDir targetDir) levels = do
 --
 -- The result must ultimately be released using 'releaseRuns'.
 openRuns ::
-     (MonadMask m, MonadSTM m, MonadST m, MonadMVar m)
+     (MonadMask m, MonadSTM m, MonadST m)
   => ActionRegistry m
   -> HasFS m h
   -> HasBlockIO m h

src/Database/LSMTree/Internal/UniqCounter.hs

@@ -7,32 +7,32 @@ module Database.LSMTree.Internal.UniqCounter (
   uniqueToRunNumber,
 ) where
 
-import           Control.Concurrent.Class.MonadMVar.Strict
-import           Data.Coerce (coerce)
+import           Control.Monad.Primitive (PrimMonad, PrimState)
+import           Data.Primitive.PrimVar as P
 import           Data.Word (Word64)
 import           Database.LSMTree.Internal.RunNumber
 
--- | A newtype wrapper around 'Word64'.
-newtype Unique = Unique Word64
+-- | A unique value derived from a 'UniqCounter'.
+newtype Unique = Unique Int
 
 -- | Avoid this function, use specialised versions like 'uniqueToRunNumber' if possible.
 uniqueToWord64 :: Unique -> Word64
-uniqueToWord64 = coerce
+uniqueToWord64 (Unique n) = fromIntegral n
 
 uniqueToRunNumber :: Unique -> RunNumber
-uniqueToRunNumber = coerce
+uniqueToRunNumber (Unique n) = RunNumber (fromIntegral n)
 
 -- | An atomic counter for producing 'Unique' values.
 --
--- TODO: could we use a PrimVar here?
-newtype UniqCounter m = UniqCounter (StrictMVar m Word64)
+newtype UniqCounter m = UniqCounter (PrimVar (PrimState m) Int)
 
 {-# INLINE newUniqCounter #-}
-newUniqCounter :: MonadMVar m => Word64 -> m (UniqCounter m)
-newUniqCounter x = UniqCounter <$> newMVar x
+newUniqCounter :: PrimMonad m => Int -> m (UniqCounter m)
+newUniqCounter = fmap UniqCounter . P.newPrimVar
 
 {-# INLINE incrUniqCounter #-}
--- | Return the current state of the atomic counter, and then increment the
+-- | Atomically, return the current state of the counter, and increment the
 -- counter.
-incrUniqCounter :: MonadMVar m => UniqCounter m -> m Unique
-incrUniqCounter (UniqCounter uniqVar) = modifyMVar uniqVar (\x -> pure ((x+1), Unique x))
+incrUniqCounter :: PrimMonad m => UniqCounter m -> m Unique
+incrUniqCounter (UniqCounter uniqVar) =
+  Unique <$> P.fetchAddInt uniqVar 1