Use clone to trim arrays

`cloneSmallMutableArray#` looks like a much better fit for
implementing `trim` than does `copySmallMutableArray#`.

Author: treeowl

Data/HashMap/Array.hs

@@ -29,6 +29,7 @@ module Data.HashMap.Array
     , insertM
     , delete
     , sameArray1
+    , trim
 
     , unsafeFreeze
     , unsafeThaw
@@ -70,13 +71,13 @@ import Prelude hiding (filter, foldr, length, map, read)
 import GHC.Exts (SmallArray#, newSmallArray#, readSmallArray#, writeSmallArray#,
                  indexSmallArray#, unsafeFreezeSmallArray#, unsafeThawSmallArray#,
                  SmallMutableArray#, sizeofSmallArray#, copySmallArray#, thawSmallArray#,
-                 sizeofSmallMutableArray#, copySmallMutableArray#)
+                 sizeofSmallMutableArray#, copySmallMutableArray#, cloneSmallMutableArray#)
 
 #else
 import GHC.Exts (Array#, newArray#, readArray#, writeArray#,
                  indexArray#, unsafeFreezeArray#, unsafeThawArray#,
                  MutableArray#, sizeofArray#, copyArray#, thawArray#,
-                 sizeofMutableArray#, copyMutableArray#)
+                 sizeofMutableArray#, copyMutableArray#, cloneMutableArray#)
 #endif
 
 #if defined(ASSERTS)
@@ -125,6 +126,13 @@ copyArray# :: SmallArray# a
            -> State# d
 copyArray# = copySmallArray#
 
+cloneMutableArray# :: SmallMutableArray# s a
+                   -> Int#
+                   -> Int#
+                   -> State# s
+                   -> (# State# s, SmallMutableArray# s a #)
+cloneMutableArray# = cloneSmallMutableArray#
+
 thawArray# :: SmallArray# a
            -> Int#
            -> Int#
@@ -329,6 +337,19 @@ copyM !src !_sidx@(I# sidx#) !dst !_didx@(I# didx#) _n@(I# n#) =
     case copyMutableArray# (unMArray src) sidx# (unMArray dst) didx# n# s# of
         s2 -> (# s2, () #)
 
+cloneM :: MArray s a -> Int -> Int -> ST s (MArray s a)
+cloneM _mary@(MArray mary#) _off@(I# off#) _len@(I# len#) =
+    CHECK_BOUNDS("cloneM_off", lengthM _mary, _off - 1)
+    CHECK_BOUNDS("cloneM_end", lengthM _mary, _off + _len - 1)
+    ST $ \ s ->
+    case cloneMutableArray# 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

Data/HashMap/Base.hs

@@ -1564,14 +1564,6 @@ foldrWithKey f = go
 ------------------------------------------------------------------------
 -- * Filter
 
--- | Create a new array of the @n@ first elements of @mary@.
-trim :: A.MArray s a -> Int -> ST s (A.Array a)
-trim mary n = do
-    mary2 <- A.new_ n
-    A.copyM mary 0 mary2 0 n
-    A.unsafeFreeze mary2
-{-# INLINE trim #-}
-
 -- | /O(n)/ Transform this map by applying a function to every value
 --   and retaining only some of them.
 mapMaybeWithKey :: (k -> v1 -> Maybe v2) -> HashMap k v1 -> HashMap k v2
@@ -1634,9 +1626,9 @@ filterMapAux onLeaf onColl = go
                     ch <- A.read mary 0
                     case ch of
                       t | isLeafOrCollision t -> return t
-                      _                       -> BitmapIndexed b <$> trim mary 1
+                      _                       -> BitmapIndexed b <$> A.trim mary 1
                 _ -> do
-                    ary2 <- trim mary j
+                    ary2 <- A.trim mary j
                     return $! if j == maxChildren
                               then Full ary2
                               else BitmapIndexed b ary2
@@ -1663,7 +1655,7 @@ filterMapAux onLeaf onColl = go
                         return $! Leaf h l
                 _ | i == j -> do ary2 <- A.unsafeFreeze mary
                                  return $! Collision h ary2
-                  | otherwise -> do ary2 <- trim mary j
+                  | otherwise -> do ary2 <- A.trim 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