Use shrink in filterMapAux (#433)

This results in a ~8% speedup in the filterWithKey benchmark. Context: #362
haskell-unordered-containers · Apr 25, 2022 · 0a474c7 · 0a474c7
1 parent 8f7ebc5
commit 0a474c7
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 9 deletions.
diff --git a/Data/HashMap/Internal.hs b/Data/HashMap/Internal.hs
@@ -2101,9 +2101,9 @@ filterMapAux onLeaf onColl = go
                     ch <- A.read mary 0
                     case ch of
                       t | isLeafOrCollision t -> return t
-                      _                       -> BitmapIndexed b <$> A.trim mary 1
+                      _                       -> BitmapIndexed b <$> (A.unsafeFreeze =<< A.shrink mary 1)
                 _ -> do
-                    ary2 <- A.trim mary j
+                    ary2 <- A.unsafeFreeze =<< A.shrink mary j
                     return $! if j == maxChildren
                               then Full ary2
                               else BitmapIndexed b ary2
@@ -2130,7 +2130,7 @@ filterMapAux onLeaf onColl = go
                         return $! Leaf h l
                 _ | i == j -> do ary2 <- A.unsafeFreeze mary
                                  return $! Collision h ary2
-                  | otherwise -> do ary2 <- A.trim mary j
+                  | otherwise -> do ary2 <- A.unsafeFreeze =<< A.shrink mary j
                                     return $! Collision h ary2
             | Just el <- onColl $! A.index ary i
                 = A.write mary j el >> step ary mary (i+1) (j+1) n

diff --git a/Data/HashMap/Internal/Array.hs b/Data/HashMap/Internal/Array.hs
@@ -52,14 +52,14 @@ module Data.HashMap.Internal.Array
     , insertM
     , delete
     , sameArray1
-    , trim
 
     , unsafeFreeze
     , unsafeThaw
     , unsafeSameArray
     , run
     , copy
     , copyM
+    , cloneM
 
       -- * Folds
     , foldl
@@ -318,11 +318,6 @@ cloneM _mary@(MArray mary#) _off@(I# off#) _len@(I# len#) =
     case cloneSmallMutableArray# mary# off# len# s of
       (# s', mary'# #) -> (# s', MArray mary'# #)
 
--- | Create a new array of the @n@ first elements of @mary@.
-trim :: MArray s a -> Int -> ST s (Array a)
-trim mary n = cloneM mary 0 n >>= unsafeFreeze
-{-# INLINE trim #-}
-
 -- | \(O(n)\) Insert an element at the given position in this array,
 -- increasing its size by one.
 insert :: Array e -> Int -> e -> Array e