Introduce a CastWeakBlobRef class for retrieveBlobs

This new class allows us to cast blob references inside the implementation of
`retrieveBlobs` differently depending on the monad it is running in. The `IO`
instance casts blob references the same way we did before the previous commit.
This reverts the slight performance impact that the previous commit introduced
by removing the `Typeable m` constraint on the `retrieveBlobs` function.

Author: jorisdral

src/Database/LSMTree/Internal/BlobRef.hs

@@ -17,6 +17,7 @@ module Database.LSMTree.Internal.BlobRef (
   , removeReferences
   , readBlob
   , readBlobIOOp
+  , CastWeakBlobRef (..)
   ) where
 
 import           Control.DeepSeq (NFData (..))
@@ -29,6 +30,7 @@ import qualified Control.RefCount as RC
 import           Data.Coerce (coerce)
 import qualified Data.Primitive.ByteArray as P (MutableByteArray,
                      newPinnedByteArray, unsafeFreezeByteArray)
+import           Data.Typeable
 import qualified Data.Vector as V
 import           Data.Word (Word32, Word64)
 import qualified Database.LSMTree.Internal.RawBytes as RB
@@ -186,3 +188,14 @@ readBlobIOOp buf bufoff
       (fromIntegral blobSpanOffset :: FS.FileOffset)
       buf (FS.BufferOffset bufoff)
       (fromIntegral blobSpanSize :: FS.ByteCount)
+
+-- | A stronger alternative to 'cast', which can be used to cast /only/ the @h@
+-- parameter to 'WeakBlobRef'.
+class CastWeakBlobRef m where
+  castWeakBlobRef ::
+       (Typeable h, Typeable h')
+    => WeakBlobRef m h
+    -> Maybe (WeakBlobRef m h')
+
+instance CastWeakBlobRef IO where
+  castWeakBlobRef = cast

src/Database/LSMTree/Normal.hs

@@ -83,6 +83,7 @@ module Database.LSMTree.Normal (
     -- ** Blobs
   , BlobRef
   , retrieveBlobs
+  , CheckBlobRefType
 
     -- * Durability (snapshots)
   , SnapshotName
@@ -115,7 +116,7 @@ import           Control.Exception (throw)
 import           Control.Monad
 import           Data.Bifunctor (Bifunctor (..))
 import           Data.Kind (Type)
-import           Data.Typeable (Proxy (..), cast)
+import           Data.Typeable (Proxy (..))
 import qualified Data.Vector as V
 import           Database.LSMTree.Common (BlobRef (BlobRef), IOLike, Range (..),
                      SerialiseKey, SerialiseValue, Session, SnapshotName,
@@ -570,6 +571,7 @@ deletes = updates . fmap (,Delete)
 retrieveBlobs ::
      ( IOLike m
      , SerialiseValue blob
+     , CheckBlobRefType m
      )
   => Session m
   -> V.Vector (BlobRef m blob)
@@ -578,10 +580,11 @@ retrieveBlobs (Internal.Session' sesh) refs =
     V.map Internal.deserialiseBlob <$>
       Internal.retrieveBlobs sesh (V.imap checkBlobRefType refs)
   where
-    checkBlobRefType _ (BlobRef (Internal.WeakBlobRef (Internal.BlobRef ref a b)))
-      | Just ref' <- cast ref = Internal.WeakBlobRef (Internal.BlobRef ref' a b)
+    checkBlobRefType _ (BlobRef ref) | Just ref' <- Internal.castWeakBlobRef ref = ref'
     checkBlobRefType i _ = throw (Internal.ErrBlobRefInvalid i)
 
+type CheckBlobRefType = Internal.CastWeakBlobRef
+
 {-------------------------------------------------------------------------------
   Snapshots
 -------------------------------------------------------------------------------}

test/Database/LSMTree/Class/Normal.hs

@@ -120,6 +120,7 @@ class (IsSession (Session h)) => IsTableHandle h where
     retrieveBlobs ::
            ( IOLike m
            , SerialiseValue blob
+           , R.CheckBlobRefType m
            )
         => proxy h
         -> Session h m

test/Test/Database/LSMTree/Class/Normal.hs

@@ -135,6 +135,7 @@ retrieveBlobsTrav ::
      , IOLike m
      , SerialiseValue blob
      , Traversable t
+     , R.CheckBlobRefType m
      )
   => proxy h
   -> Session h m
@@ -153,6 +154,7 @@ lookupsWithBlobs :: forall h m k v blob.
      , SerialiseKey k
      , SerialiseValue v
      , SerialiseValue blob
+     , R.CheckBlobRefType m
      )
   => h m k v blob
   -> Session h m
@@ -168,6 +170,7 @@ rangeLookupWithBlobs :: forall h m k v blob.
      , SerialiseKey k
      , SerialiseValue v
      , SerialiseValue blob
+     , R.CheckBlobRefType m
      )
   => h m k v blob
   -> Session h m
@@ -183,6 +186,7 @@ readCursorWithBlobs :: forall h m k v blob proxy.
      , SerialiseKey k
      , SerialiseValue v
      , SerialiseValue blob
+     , R.CheckBlobRefType m
      )
   => proxy h
   -> Session h m
@@ -199,6 +203,7 @@ readCursorAllWithBlobs :: forall h m k v blob proxy.
      , SerialiseKey k
      , SerialiseValue v
      , SerialiseValue blob
+     , R.CheckBlobRefType m
      )
   => proxy h
   -> Session h m

test/Test/Database/LSMTree/Normal/StateMachine.hs

@@ -72,6 +72,7 @@ import           Database.LSMTree.Extras (showPowersOf)
 import           Database.LSMTree.Extras.Generators (KeyForIndexCompact)
 import           Database.LSMTree.Extras.NoThunks (assertNoThunks)
 import           Database.LSMTree.Internal (LSMTreeError (..))
+import qualified Database.LSMTree.Internal.BlobRef as Internal
 import qualified Database.LSMTree.Model.Normal.Session as Model
 import qualified Database.LSMTree.ModelIO.Normal as M
 import qualified Database.LSMTree.Normal as R
@@ -128,6 +129,13 @@ tests = testGroup "Normal.StateMachine" [
 labelledExamples :: IO ()
 labelledExamples = QC.labelledExamples $ Lockstep.Run.tagActions (Proxy @(ModelState R.TableHandle))
 
+instance Arbitrary M.TableConfig where
+  arbitrary :: Gen M.TableConfig
+  arbitrary = pure M.TableConfig
+
+deriving via AllowThunk (M.Session IO)
+    instance NoThunks (M.Session IO)
+
 propLockstep_ModelIOImpl ::
      Actions (Lockstep (ModelState M.TableHandle))
   -> QC.Property
@@ -181,8 +189,17 @@ propLockstep_ModelIOImpl =
           | otherwise
           = Nothing
 
-deriving via AllowThunk (M.Session IO)
-    instance NoThunks (M.Session IO)
+instance Arbitrary R.TableConfig where
+  arbitrary :: Gen R.TableConfig
+  arbitrary = pure $ R.TableConfig {
+        R.confMergePolicy       = R.MergePolicyLazyLevelling
+      , R.confSizeRatio         = R.Four
+      , R.confWriteBufferAlloc  = R.AllocNumEntries (R.NumEntries 30)
+      , R.confBloomFilterAlloc  = R.AllocFixed 10
+      , R.confFencePointerIndex = R.CompactIndex
+      , R.confDiskCachePolicy   = R.DiskCacheNone
+      , R.confMergeSchedule     = R.OneShot
+      }
 
 propLockstep_RealImpl_RealFS_IO ::
      Actions (Lockstep (ModelState R.TableHandle))
@@ -230,6 +247,10 @@ propLockstep_RealImpl_MockFS_IO =
       mockfs <- atomically $ readTMVar fsVar
       assert (MockFS.numOpenHandles mockfs == 0) $ pure ()
 
+instance Internal.CastWeakBlobRef (IOSim s) where
+  castWeakBlobRef (Internal.WeakBlobRef (Internal.BlobRef ref a b)) =
+      (\ref' -> Internal.WeakBlobRef (Internal.BlobRef ref' a b)) <$> cast ref
+
 propLockstep_RealImpl_MockFS_IOSim ::
      Actions (Lockstep (ModelState R.TableHandle))
   -> QC.Property
@@ -281,22 +302,6 @@ createSystemTempDirectory prefix = do
     hasBlockIO <- ioHasBlockIO hasFS defaultIOCtxParams
     pure (tempDir, hasFS, hasBlockIO)
 
-instance Arbitrary M.TableConfig where
-  arbitrary :: Gen M.TableConfig
-  arbitrary = pure M.TableConfig
-
-instance Arbitrary R.TableConfig where
-  arbitrary :: Gen R.TableConfig
-  arbitrary = pure $ R.TableConfig {
-        R.confMergePolicy       = R.MergePolicyLazyLevelling
-      , R.confSizeRatio         = R.Four
-      , R.confWriteBufferAlloc  = R.AllocNumEntries (R.NumEntries 30)
-      , R.confBloomFilterAlloc  = R.AllocFixed 10
-      , R.confFencePointerIndex = R.CompactIndex
-      , R.confDiskCachePolicy   = R.DiskCacheNone
-      , R.confMergeSchedule     = R.OneShot
-      }
-
 {-------------------------------------------------------------------------------
   Key and value types
 -------------------------------------------------------------------------------}