Add missing IO and ST specializations

src/Database/LSMTree/Internal/Arena.hs

@@ -34,6 +34,12 @@ import           Data.Word (Word8)
 
 data ArenaManager s = ArenaManager (MutVar s [Arena s])
 
+{-# SPECIALISE
+    newArenaManager :: ST s (ArenaManager s)
+  #-}
+{-# SPECIALISE
+    newArenaManager :: IO (ArenaManager RealWorld)
+  #-}
 newArenaManager :: PrimMonad m => m (ArenaManager (PrimState m))
 newArenaManager = do
     m <- newMutVar []
@@ -116,8 +122,12 @@ closeArena (ArenaManager arenas) arena = do
 
     atomicModifyMutVar' arenas $ \xs -> (arena : xs, ())
 
-
-
+{-# SPECIALISE
+    scrambleArena :: Arena s -> ST s ()
+  #-}
+{-# SPECIALISE
+    scrambleArena :: Arena RealWorld -> IO ()
+  #-}
 scrambleArena :: PrimMonad m => Arena (PrimState m) -> m ()
 #ifndef NO_IGNORE_ASSERTS
 scrambleArena _ = return ()
@@ -127,6 +137,12 @@ scrambleArena Arena {..} = do
     readMutVar full >>= mapM_ scrambleBlock
     readMutVar free >>= mapM_ scrambleBlock
 
+{-# SPECIALISE
+    scrambleBlock :: Block s -> ST s ()
+  #-}
+{-# SPECIALISE
+    scrambleBlock :: Block RealWorld -> IO ()
+  #-}
 scrambleBlock :: PrimMonad m => Block (PrimState m) -> m ()
 scrambleBlock (Block _ mba) = do
     size <- getSizeofMutableByteArray mba
@@ -157,6 +173,12 @@ resetArena Arena {..} = do
 -- | Create unmanaged arena.
 --
 -- Never use this in non-tests code.
+{-# SPECIALISE
+    withUnmanagedArena :: (Arena s -> ST s a) -> ST s a
+  #-}
+{-# SPECIALISE
+    withUnmanagedArena :: (Arena RealWorld -> IO a) -> IO a
+  #-}
 withUnmanagedArena :: PrimMonad m => (Arena (PrimState m) -> m a) -> m a
 withUnmanagedArena k = do
     mgr <- newArenaManager
@@ -165,6 +187,9 @@ withUnmanagedArena k = do
 {-# SPECIALISE
     allocateFromArena :: Arena s -> Size -> Alignment -> ST s (Offset, MutableByteArray s)
   #-}
+{-# SPECIALISE
+    allocateFromArena :: Arena RealWorld -> Size -> Alignment -> IO (Offset, MutableByteArray RealWorld)
+  #-}
 -- | 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 =
@@ -175,6 +200,9 @@ allocateFromArena !arena !size !alignment =
 {-# SPECIALISE
     allocateFromArena' :: Arena s -> Size -> Alignment -> ST s (Offset, MutableByteArray s)
   #-}
+{-# SPECIALISE
+    allocateFromArena' :: Arena RealWorld -> Size -> Alignment -> IO (Offset, MutableByteArray RealWorld)
+  #-}
 -- 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
@@ -206,7 +234,12 @@ allocateFromArena' arena@Arena { .. } !size !alignment = do
         -- * go again
         allocateFromArena' arena size alignment
 
-{-# SPECIALISE newBlockWithFree :: MutVar s [Block s] -> ST s (Block s) #-}
+{-# SPECIALISE
+    newBlockWithFree :: MutVar s [Block s] -> ST s (Block s)
+  #-}
+{-# SPECIALISE
+    newBlockWithFree :: MutVar RealWorld [Block RealWorld] -> IO (Block RealWorld)
+  #-}
 -- | 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/CRC32C.hs

@@ -466,11 +466,10 @@ data FileCorruptedError
 
 {-# SPECIALISE
   expectValidFile ::
-      (MonadThrow m)
-    => FsPath
+       FsPath
     -> FileFormat
     -> Either String a
-    -> m a
+    -> IO a
   #-}
 expectValidFile ::
      (MonadThrow m)

src/Database/LSMTree/Internal/IncomingRun.hs

@@ -246,6 +246,11 @@ creditThresholdForLevel conf (LevelNo _i) =
 -- This is /not/ itself thread safe. All 'TableContent' update operations are
 -- expected to be serialised by the caller. See concurrency comments for
 -- 'TableContent' for detail.
+{-# SPECIALISE depositNominalCredits ::
+     NominalDebt
+  -> PrimVar RealWorld NominalCredits
+  -> NominalCredits
+  -> IO (NominalCredits, NominalCredits) #-}
 depositNominalCredits ::
      PrimMonad m
   => NominalDebt

src/Database/LSMTree/Internal/MergeSchedule.hs

@@ -402,6 +402,10 @@ newtype UnionCache m h = UnionCache {
     cachedTree :: MT.LookupTree (V.Vector (Ref (Run m h)))
   }
 
+{-# SPECIALISE mkUnionCache ::
+     ActionRegistry IO
+  -> Ref (MergingTree IO h)
+  -> IO (UnionCache IO h) #-}
 mkUnionCache ::
      (PrimMonad m, MonadMVar m, MonadMask m)
   => ActionRegistry m
@@ -410,6 +414,10 @@ mkUnionCache ::
 mkUnionCache reg mt =
     UnionCache <$> MT.buildLookupTree reg mt
 
+{-# SPECIALISE duplicateUnionCache ::
+     ActionRegistry IO
+  -> UnionCache IO h
+  -> IO (UnionCache IO h) #-}
 duplicateUnionCache ::
      (PrimMonad m, MonadMask m)
   => ActionRegistry m
@@ -419,6 +427,10 @@ duplicateUnionCache reg (UnionCache mt) =
     UnionCache <$>
       MT.mapMStrict (mapMStrict (\r -> withRollback reg (dupRef r) releaseRef)) mt
 
+{-# SPECIALISE releaseUnionCache ::
+     ActionRegistry IO
+  -> UnionCache IO h
+  -> IO () #-}
 releaseUnionCache ::
      (PrimMonad m, MonadMask m)
   => ActionRegistry m

src/Database/LSMTree/Internal/MergingTree.hs

@@ -532,6 +532,9 @@ supplyCredits hfs hbio resolve runParams threshold root uc = \mt0 c0 -> do
 
 -- | This does /not/ release the reference, but allocates a new reference for
 -- the returned run, which must be released at some point.
+{-# SPECIALISE expectCompleted ::
+     Ref (MergingTree IO h)
+  -> IO (Ref (Run IO h)) #-}
 expectCompleted ::
      (MonadMVar m, MonadSTM m, MonadST m, MonadMask m)
   => Ref (MergingTree m h) -> m (Ref (Run m h))

src/Database/LSMTree/Internal/MergingTree/Lookup.hs

@@ -33,6 +33,10 @@ data LookupTree a =
 
 -- | Deriving 'Traversable' leads to functions that are not strict in the
 -- elements of the vector of children. This function avoids that issue.
+{-# SPECIALISE mapMStrict ::
+     (a -> IO b)
+  -> LookupTree a
+  -> IO (LookupTree b) #-}
 mapMStrict :: Monad m => (a -> m b) -> LookupTree a -> m (LookupTree b)
 mapMStrict f = \case
   LookupBatch a  -> LookupBatch <$!> f a