Switch to using unboxed type proxies

IntersectMBO · Dec 6, 2024 · 0b36ff3 · 0b36ff3
1 parent 58dba23
commit 0b36ff3
Show file tree

Hide file tree

Showing 6 changed files with 27 additions and 17 deletions.
diff --git a/src/Database/LSMTree/Internal/ChecksumHandle.hs b/src/Database/LSMTree/Internal/ChecksumHandle.hs
@@ -1,3 +1,5 @@
+{-# LANGUAGE MagicHash #-}
+
 module Database.LSMTree.Internal.ChecksumHandle
   (
     -- * Checksum handles
@@ -25,7 +27,7 @@ import           Control.Monad.Primitive
 import           Data.BloomFilter (Bloom)
 import qualified Data.ByteString.Lazy as BSL
 import           Data.Primitive.PrimVar
-import           Data.Proxy (Proxy)
+import           GHC.Exts (Proxy#)
 import           Data.Word (Word64)
 import           Database.LSMTree.Internal.BlobRef (BlobSpan (..), RawBlobRef)
 import qualified Database.LSMTree.Internal.BlobRef as BlobRef
@@ -205,13 +207,13 @@ writeFilter hfs filterHandle bf =
 {-# SPECIALISE writeIndexHeader ::
      HasFS IO h
   -> ForIndex (ChecksumHandle RealWorld h)
-  -> Proxy IndexCompact
+  -> Proxy# IndexCompact
   -> IO () #-}
 writeIndexHeader ::
      (MonadSTM m, PrimMonad m)
   => HasFS m h
   -> ForIndex (ChecksumHandle (PrimState m) h)
-  -> Proxy IndexCompact
+  -> Proxy# IndexCompact
   -> m ()
 writeIndexHeader hfs indexHandle indexTypeProxy =
     writeToHandle hfs (unForIndex indexHandle) $

diff --git a/src/Database/LSMTree/Internal/Index.hs b/src/Database/LSMTree/Internal/Index.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE MagicHash    #-}
 {-# LANGUAGE TypeFamilies #-}
 
 {-|
@@ -14,7 +15,7 @@ where
 import           Control.Monad.ST.Strict (ST)
 import           Data.ByteString.Lazy (LazyByteString)
 import           Data.ByteString.Short (ShortByteString)
-import           Data.Proxy (Proxy)
+import           GHC.Exts (Proxy#)
 import           Data.Word (Word32)
 import           Database.LSMTree.Internal.Chunk (Chunk)
 import           Database.LSMTree.Internal.Entry (NumEntries)
@@ -54,7 +55,7 @@ class Index i where
         See the documentation of the 'Index' class for how to generate a
         complete serialised index.
     -}
-    headerLBS :: Proxy i -> LazyByteString
+    headerLBS :: Proxy# i -> LazyByteString
 
     {-|
         Yields the footer of the serialised form of an index.

diff --git a/src/Database/LSMTree/Internal/Index/Compact.hs b/src/Database/LSMTree/Internal/Index/Compact.hs
@@ -1,3 +1,5 @@
+{-# LANGUAGE MagicHash #-}
+
 -- | A compact fence-pointer index for uniformly distributed keys.
 --
 -- TODO: add utility functions for clash probability calculations
@@ -36,7 +38,7 @@ import           Data.Maybe (fromMaybe)
 import           Data.Primitive.ByteArray (ByteArray (..), indexByteArray,
                      sizeofByteArray)
 import           Data.Primitive.Types (sizeOf)
-import           Data.Proxy (Proxy (Proxy))
+import           GHC.Exts (Proxy#, proxy#)
 import qualified Data.Vector.Algorithms.Search as VA
 import qualified Data.Vector.Generic as VG
 import qualified Data.Vector.Primitive as VP
@@ -460,7 +462,7 @@ sizeInPages = NumPages . toEnum . VU.length . icPrimary
 -- | Serialises a compact index in one go.
 toLBS :: NumEntries -> IndexCompact -> LBS.ByteString
 toLBS numEntries index =
-     headerLBS (Proxy @IndexCompact)
+     headerLBS (proxy# @IndexCompact)
   <> LBS.fromStrict (Chunk.toByteString (word64VectorToChunk (icPrimary index)))
   <> finalLBS numEntries index
 
@@ -478,7 +480,7 @@ supportedTypeAndVersion = 0x0001
     For a specification of this operation, see the documentation of [its
     polymorphic version]('Index.headerLBS').
 -}
-headerLBS :: Proxy IndexCompact -> LBS.ByteString
+headerLBS :: Proxy# IndexCompact -> LBS.ByteString
 headerLBS _ =
     -- create a single 4 byte chunk
     BB.toLazyByteStringWith (BB.safeStrategy 4 BB.smallChunkSize) mempty $
@@ -686,7 +688,7 @@ instance Index IndexCompact where
     sizeInPages :: IndexCompact -> NumPages
     sizeInPages = sizeInPages
 
-    headerLBS :: Proxy IndexCompact -> LBS.ByteString
+    headerLBS :: Proxy# IndexCompact -> LBS.ByteString
     headerLBS = headerLBS
 
     finalLBS :: NumEntries -> IndexCompact -> LBS.ByteString

diff --git a/src/Database/LSMTree/Internal/Index/Ordinary.hs b/src/Database/LSMTree/Internal/Index/Ordinary.hs
@@ -1,3 +1,5 @@
+{-# LANGUAGE MagicHash #-}
+
 {- HLINT ignore "Avoid restricted alias" -}
 
 -- | A general-purpose fence pointer index.
@@ -22,7 +24,7 @@ import           Data.ByteString.Short (ShortByteString (SBS))
 import qualified Data.ByteString.Short as ShortByteString (length)
 import           Data.Primitive.ByteArray (ByteArray (ByteArray),
                      indexByteArray)
-import           Data.Proxy (Proxy)
+import           GHC.Exts (Proxy#)
 import           Data.Vector (Vector, drop, findIndex, findIndexR, fromList,
                      last, length, (!))
 import qualified Data.Vector.Primitive as Primitive (Vector (Vector), drop,
@@ -104,7 +106,7 @@ instance Index IndexOrdinary where
     sizeInPages (IndexOrdinary lastKeys)
         = NumPages $ fromIntegral (length lastKeys)
 
-    headerLBS :: Proxy IndexOrdinary -> LazyByteString
+    headerLBS :: Proxy# IndexOrdinary -> LazyByteString
     headerLBS _ = toLazyByteString        $
                   word32Host              $
                   supportedTypeAndVersion

diff --git a/src/Database/LSMTree/Internal/RunBuilder.hs b/src/Database/LSMTree/Internal/RunBuilder.hs
@@ -1,3 +1,5 @@
+{-# LANGUAGE MagicHash #-}
+
 -- | A mutable run ('RunBuilder') that is under construction.
 --
 module Database.LSMTree.Internal.RunBuilder (
@@ -18,7 +20,7 @@ import           Control.Monad.Primitive
 import           Data.BloomFilter (Bloom)
 import           Data.Foldable (for_, traverse_)
 import           Data.Primitive.PrimVar
-import           Data.Proxy (Proxy (Proxy))
+import           GHC.Exts (proxy#)
 import           Data.Word (Word64)
 import           Database.LSMTree.Internal.BlobRef (RawBlobRef)
 import           Database.LSMTree.Internal.ChecksumHandle
@@ -92,7 +94,7 @@ new hfs hbio runBuilderFsPaths numEntries alloc = do
     runBuilderHandles <- traverse (makeHandle hfs) (pathsForRunFiles runBuilderFsPaths)
 
     let builder = RunBuilder { runBuilderHasFS = hfs, runBuilderHasBlockIO = hbio, .. }
-    writeIndexHeader hfs (forRunIndex runBuilderHandles) (Proxy @IndexCompact)
+    writeIndexHeader hfs (forRunIndex runBuilderHandles) (proxy# @IndexCompact)
     return builder
 
 {-# SPECIALISE addKeyOp ::

diff --git a/test/Test/Database/LSMTree/Internal/Index/Compact.hs b/test/Test/Database/LSMTree/Internal/Index/Compact.hs
@@ -1,5 +1,6 @@
 {-# LANGUAGE LambdaCase      #-}
 {-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE MagicHash       #-}
 {-# OPTIONS_GHC -Wno-orphans #-}
 {- HLINT ignore "Eta reduce" -}
 
@@ -22,7 +23,7 @@ import           Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
 import           Data.Primitive.ByteArray (ByteArray (..), byteArrayFromList,
                      sizeofByteArray)
-import           Data.Proxy (Proxy (Proxy))
+import           GHC.Exts (proxy#)
 import qualified Data.Vector.Primitive as VP
 import qualified Data.Vector.Unboxed as VU
 import qualified Data.Vector.Unboxed.Base as VU
@@ -100,7 +101,7 @@ tests = testGroup "Test.Database.LSMTree.Internal.Index.Compact" [
                 , 7, 0
                 ]
 
-          let header = LBS.unpack (headerLBS (Proxy @IndexCompact))
+          let header = LBS.unpack (headerLBS (proxy# @IndexCompact))
           let primary = LBS.unpack $
                         LBS.fromChunks (map Chunk.toByteString chunks)
           let rest = LBS.unpack (finalLBS (NumEntries 7) index)
@@ -232,7 +233,7 @@ prop_roundtrip_chunks (Chunks chunks index) numEntries =
     counterexample ("rest:\n" <> showBS bsRest) $
       Right (numEntries, index) === fromSBS sbs
   where
-    bsVersion = headerLBS (Proxy @IndexCompact)
+    bsVersion = headerLBS (proxy# @IndexCompact)
     bsPrimary = LBS.fromChunks $
                 map (Chunk.toByteString . word64VectorToChunk) chunks
     bsRest = finalLBS numEntries index
@@ -294,7 +295,7 @@ writeIndexCompact numEntries (ChunkSize csize) ps = runST $ do
     cs <- mapM (`append` ica) (toAppends ps)
     (c, index) <- unsafeEnd ica
     return
-      ( headerLBS (Proxy @IndexCompact)
+      ( headerLBS (proxy# @IndexCompact)
       , LBS.fromChunks $
         foldMap (map Chunk.toByteString) $ cs <> pure (toList c)
       , finalLBS numEntries index