Prevent NoThunks failures using strict(er) vector operations

IntersectMBO · Dec 24, 2024 · 6e6f726 · 6e6f726
1 parent 640f48a
commit 6e6f726
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git a/src/Database/LSMTree/Internal/MergeSchedule.hs b/src/Database/LSMTree/Internal/MergeSchedule.hs
@@ -58,7 +58,7 @@ import           Database.LSMTree.Internal.RunNumber
 import           Database.LSMTree.Internal.Serialise (SerialisedBlob,
                      SerialisedKey, SerialisedValue)
 import           Database.LSMTree.Internal.UniqCounter
-import           Database.LSMTree.Internal.Vector (mapStrict)
+import           Database.LSMTree.Internal.Vector (forMStrict, mapStrict)
 import           Database.LSMTree.Internal.WriteBuffer (WriteBuffer)
 import qualified Database.LSMTree.Internal.WriteBuffer as WB
 import           Database.LSMTree.Internal.WriteBufferBlobs (WriteBufferBlobs)
@@ -200,7 +200,7 @@ mkLevelsCache reg lvls = do
       -> Levels m h
       -> m a
     foldRunAndMergeM k1 k2 ls =
-        fmap fold $ V.forM ls $ \(Level ir rs) -> do
+        fmap fold $ forMStrict ls $ \(Level ir rs) -> do
           incoming <- case ir of
             Single     r -> k1 r
             Merging _ mr -> k2 mr
@@ -251,7 +251,7 @@ duplicateLevelsCache ::
   -> LevelsCache m h
   -> m (LevelsCache m h)
 duplicateLevelsCache reg cache = do
-    rs' <- V.forM (cachedRuns cache) $ \r ->
+    rs' <- forMStrict (cachedRuns cache) $ \r ->
              withRollback reg (dupRef r) releaseRef
     return cache { cachedRuns = rs' }
 
@@ -300,9 +300,9 @@ duplicateLevels ::
   -> Levels m h
   -> m (Levels m h)
 duplicateLevels reg levels =
-    V.forM levels $ \Level {incomingRun, residentRuns} -> do
+    forMStrict levels $ \Level {incomingRun, residentRuns} -> do
       incomingRun'  <- duplicateIncomingRun reg incomingRun
-      residentRuns' <- V.forM residentRuns $ \r ->
+      residentRuns' <- forMStrict residentRuns $ \r ->
                          withRollback reg (dupRef r) releaseRef
       return $! Level {
         incomingRun  = incomingRun',

diff --git a/src/Database/LSMTree/Internal/Vector.hs b/src/Database/LSMTree/Internal/Vector.hs
@@ -9,6 +9,7 @@ module Database.LSMTree.Internal.Vector (
     mapStrict,
     mapMStrict,
     imapMStrict,
+    forMStrict,
     zipWithStrict,
     binarySearchL,
     unsafeInsertWithMStrict,
@@ -79,6 +80,11 @@ imapMStrict f v = V.imapM (\i -> f i >=> (pure $!)) v
 zipWithStrict :: forall a b c. (a -> b -> c) -> V.Vector a -> V.Vector b -> V.Vector c
 zipWithStrict f xs ys = runST (V.zipWithM (\x y -> pure $! f x y) xs ys)
 
+-- | /( O(n) /) Like 'V.forM', but strict in the produced elements of type @b@.
+{-# INLINE forMStrict #-}
+forMStrict :: Monad m => V.Vector a -> (a -> m b) -> m (V.Vector b)
+forMStrict xs f = V.forM xs (f >=> (pure $!))
+
 {-|
     Finds the lowest index in a given sorted vector at which the given element
     could be inserted while maintaining the sortedness.