Merge pull request #330 from IntersectMBO/dcoutts/arena-specialise

Two micro-optimisations

Author: jorisdral

src/Data/Arena.hs

@@ -20,6 +20,7 @@ module Data.Arena (
 import           Control.DeepSeq (NFData (..))
 import           Control.Exception (assert)
 import           Control.Monad.Primitive
+import           Control.Monad.ST (ST)
 import           Data.Bits (complement, popCount, (.&.))
 import           Data.Primitive.ByteArray
 import           Data.Primitive.MutVar
@@ -59,19 +60,32 @@ type Alignment = Int
 blockSize :: Int
 blockSize = 0x100000
 
+{-# SPECIALIZE
+    newBlock :: ST s (Block s)
+  #-}
+{-# SPECIALIZE
+    newBlock :: IO (Block RealWorld)
+  #-}
 newBlock :: PrimMonad m => m (Block (PrimState m))
 newBlock = do
     off <- newPrimVar 0
     mba <- newAlignedPinnedByteArray blockSize 4096
     return (Block off mba)
 
+{-# INLINE withArena #-}
 withArena :: PrimMonad m => ArenaManager (PrimState m) -> (Arena (PrimState m) -> m a) -> m a
 withArena am f = do
     a <- newArena am
     x <- f a
     closeArena am a
     pure x
 
+{-# SPECIALIZE
+    newArena :: ArenaManager s -> ST s (Arena s)
+  #-}
+{-# SPECIALIZE
+    newArena :: ArenaManager RealWorld -> IO (Arena RealWorld)
+  #-}
 newArena :: PrimMonad m => ArenaManager (PrimState m) -> m (Arena (PrimState m))
 newArena (ArenaManager arenas) = do
     marena <- atomicModifyMutVar' arenas $ \case
@@ -86,6 +100,12 @@ newArena (ArenaManager arenas) = do
             full <- newMutVar []
             return Arena {..}
 
+{-# SPECIALIZE
+    closeArena :: ArenaManager s -> Arena s -> ST s ()
+  #-}
+{-# SPECIALIZE
+    closeArena :: ArenaManager RealWorld -> Arena RealWorld -> IO ()
+  #-}
 closeArena :: PrimMonad m => ArenaManager (PrimState m) -> Arena (PrimState m) -> m ()
 closeArena (ArenaManager arenas) arena = do
     scrambleArena arena
@@ -112,6 +132,12 @@ scrambleBlock (Block _ mba) = do
     setByteArray mba 0 size (0x77 :: Word8)
 #endif
 
+{-# SPECIALIZE
+    resetArena :: Arena s -> ST s ()
+  #-}
+{-# SPECIALIZE
+    resetArena :: Arena RealWorld -> IO ()
+  #-}
 -- | Reset arena, i.e. return used blocks to free list.
 resetArena :: PrimMonad m => Arena (PrimState m) -> m ()
 resetArena Arena {..} = do
@@ -135,13 +161,19 @@ withUnmanagedArena k = do
     mgr <- newArenaManager
     withArena mgr k
 
+{-# SPECIALIZE
+    allocateFromArena :: Arena s -> Size -> Alignment -> ST s (Offset, MutableByteArray s)
+  #-}
 -- | Allocate a slice of mutable byte array from the arena.
 allocateFromArena :: PrimMonad m => Arena (PrimState m)-> Size -> Alignment -> m (Offset, MutableByteArray (PrimState m))
 allocateFromArena !arena !size !alignment =
     assert (popCount alignment == 1) $ -- powers of 2
     assert (size <= blockSize) $ -- not too large allocations
     allocateFromArena' arena size alignment
 
+{-# SPECIALIZE
+    allocateFromArena' :: Arena s -> Size -> Alignment -> ST s (Offset, MutableByteArray s)
+  #-}
 -- TODO!? this is not async exception safe
 allocateFromArena' :: PrimMonad m => Arena (PrimState m)-> Size -> Alignment -> m (Offset, MutableByteArray (PrimState m))
 allocateFromArena' arena@Arena { .. } !size !alignment = do
@@ -173,6 +205,7 @@ allocateFromArena' arena@Arena { .. } !size !alignment = do
         -- * go again
         allocateFromArena' arena size alignment
 
+{-# SPECIALIZE newBlockWithFree :: MutVar s [Block s] -> ST s (Block s) #-}
 -- | Allocate new block, possibly taking it from a free list
 newBlockWithFree :: PrimMonad m => MutVar (PrimState m) [Block (PrimState m)] -> m (Block (PrimState m))
 newBlockWithFree free = do

src/Database/LSMTree/Internal/IndexCompact.hs

@@ -360,19 +360,19 @@ import           Database.LSMTree.Internal.Vector
 data IndexCompact = IndexCompact {
     -- | \(P\): Maps a page @i@ to the 64-bit slice of primary bits of its
     -- minimum key.
-    icPrimary        :: !(VU.Vector Word64)
+    icPrimary        :: {-# UNPACK #-} !(VU.Vector Word64)
     -- | \(C\): A clash on page @i@ means that the primary bits of the minimum
     -- key on that page aren't sufficient to decide whether a search for a key
     -- should continue left or right of the page.
-  , icClashes        :: !(VU.Vector Bit)
+  , icClashes        :: {-# UNPACK #-} !(VU.Vector Bit)
     -- | \(TB\): Maps a full minimum key to the page @i@ that contains it, but
     -- only if there is a clash on page @i@.
   , icTieBreaker     :: !(Map (Unsliced SerialisedKey) PageNo)
     -- | \(LTP\): Record of larger-than-page values. Given a span of pages for
     -- the larger-than-page value, the first page will map to 'False', and the
     -- remainder of the pages will be set to 'True'. Regular pages default to
     -- 'False'.
-  , icLargerThanPage :: !(VU.Vector Bit)
+  , icLargerThanPage :: {-# UNPACK #-} !(VU.Vector Bit)
   }
   deriving stock (Show, Eq)