Add size-aware 'filter, mapMaybe' functions

Author: rockbmb

Data/HashMap/Base.hs

@@ -1600,12 +1600,35 @@ mapMaybeWithKey f = filterMapAux onLeaf onColl
                        | otherwise = Nothing
 {-# INLINE mapMaybeWithKey #-}
 
+-- | /O(n)/ Transform this map by applying a function to every value
+--   and retaining only some of them.
+-- Returns a tuple with the size of the result hashmap and the the hashmap
+-- itself.
+mapMaybeWithKeyInternal
+    :: (k -> v1 -> Maybe v2)
+    -> HashMap k v1
+    -> (Int, HashMap k v2)
+mapMaybeWithKeyInternal f = filterMapAuxInternal onLeaf onColl
+  where onLeaf (Leaf h (L k v)) | Just v' <- f k v = Just (Leaf h (L k v'))
+        onLeaf _ = Nothing
+
+        onColl (L k v) | Just v' <- f k v = Just (L k v')
+                       | otherwise = Nothing
+{-# INLINE mapMaybeWithKeyInternal #-}
+
 -- | /O(n)/ Transform this map by applying a function to every value
 --   and retaining only some of them.
 mapMaybe :: (v1 -> Maybe v2) -> HashMap k v1 -> HashMap k v2
 mapMaybe f = mapMaybeWithKey (const f)
 {-# INLINE mapMaybe #-}
 
+-- | /O(n)/ Transform this map by applying a function to every value
+--   and retaining only some of them.
+-- Returns a tuple with the result hashmap's size and the hasmap itself.
+mapMaybeInternal :: (v1 -> Maybe v2) -> HashMap k v1 -> (Int, HashMap k v2)
+mapMaybeInternal f = mapMaybeWithKeyInternal (const f)
+{-# INLINE mapMaybeInternal #-}
+
 -- | /O(n)/ Filter this map by retaining only elements satisfying a
 -- predicate.
 filterWithKey :: forall k v. (k -> v -> Bool) -> HashMap k v -> HashMap k v
@@ -1617,6 +1640,17 @@ filterWithKey pred = filterMapAux onLeaf onColl
         onColl _ = Nothing
 {-# INLINE filterWithKey #-}
 
+-- | /O(n)/ Filter this map by retaining only elements satisfying a
+-- predicate.
+-- Returns a tuple with the result hashmap's size and the hashmap itself.
+filterWithKeyInternal :: forall k v. (k -> v -> Bool) -> HashMap k v -> (Int, HashMap k v)
+filterWithKeyInternal pred = filterMapAuxInternal onLeaf onColl
+  where onLeaf t@(Leaf _ (L k v)) | pred k v = Just t
+        onLeaf _ = Nothing
+
+        onColl el@(L k v) | pred k v = Just el
+        onColl _ = Nothing
+{-# INLINE filterWithKeyInternal #-}
 
 -- | Common implementation for 'filterWithKey' and 'mapMaybeWithKey',
 --   allowing the former to former to reuse terms.
@@ -1687,12 +1721,91 @@ filterMapAux onLeaf onColl = go
             | otherwise = step ary mary (i+1) j n
 {-# INLINE filterMapAux #-}
 
+-- | Common implementation for 'filterWithKey' and 'mapMaybeWithKey',
+--   allowing the former to former to reuse terms.
+-- Returns the change in the hashmap's size, and the hashmap itself.
+filterMapAuxInternal :: forall k v1 v2
+              . (HashMap k v1 -> Maybe (HashMap k v2))
+             -> (Leaf k v1 -> Maybe (Leaf k v2))
+             -> HashMap k v1
+             -> (Int, HashMap k v2)
+filterMapAuxInternal onLeaf onColl = go 0
+  where
+    go !sz Empty = (sz, Empty)
+    go !sz t@Leaf{}
+        | Just t' <- onLeaf t = (sz + 1, t')
+        | otherwise = (sz, Empty)
+    go !sz (BitmapIndexed b ary) = filterA sz ary b
+    go !sz (Full ary) = filterA sz ary fullNodeMask
+    go !sz (Collision h ary) = filterC sz ary h
+
+    filterA size ary0 b0 =
+        let !n = A.length ary0
+        in runST $ do
+            mary <- A.new_ n
+            step ary0 mary b0 0 0 1 n size
+      where
+        step :: A.Array (HashMap k v1) -> A.MArray s (HashMap k v2)
+             -> Bitmap -> Int -> Int -> Bitmap -> Int -> Int
+             -> ST s (Int, HashMap k v2)
+        step !ary !mary !b i !j !bi n !sz
+            | i >= n = case j of
+                0 -> return (sz, Empty)
+                1 -> do
+                    ch <- A.read mary 0
+                    case ch of
+                      t | isLeafOrCollision t -> return (sz, t)
+                      _                       -> (sz,) . BitmapIndexed b <$> trim mary 1
+                _ -> do
+                    ary2 <- trim mary j
+                    return $! (sz,) (if j == maxChildren
+                                     then Full ary2
+                                     else BitmapIndexed b ary2)
+            | bi .&. b == 0 = step ary mary b i j (bi `unsafeShiftL` 1) n sz
+            | otherwise = case go sz (A.index ary i) of
+                (dsz, Empty) -> step ary mary (b .&. complement bi) (i+1) j
+                                (bi `unsafeShiftL` 1) n dsz
+                (dsz, t)     -> do A.write mary j t
+                                   step ary mary b (i+1) (j+1)
+                                       (bi `unsafeShiftL` 1) n dsz
+
+    filterC size ary0 h =
+        let !n = A.length ary0
+        in runST $ do
+            mary <- A.new_ n
+            step ary0 mary 0 0 n size
+      where
+        step :: A.Array (Leaf k v1) -> A.MArray s (Leaf k v2)
+             -> Int -> Int -> Int -> Int
+             -> ST s (Int, HashMap k v2)
+        step !ary !mary i !j n !sz
+            | i >= n    = case j of
+                0 -> return (sz, Empty)
+                1 -> do l <- A.read mary 0
+                        return $! (sz, Leaf h l)
+                _ | i == j -> do ary2 <- A.unsafeFreeze mary
+                                 return $! (sz, Collision h ary2)
+                  | otherwise -> do ary2 <- trim mary j
+                                    return $! (sz, Collision h ary2)
+            | Just el <- onColl (A.index ary i)
+                = A.write mary j el >> step ary mary (i+1) (j+1) n (sz + 1)
+            | otherwise = step ary mary (i+1) j n sz
+{-# INLINE filterMapAuxInternal #-}
+
 -- | /O(n)/ Filter this map by retaining only elements which values
 -- satisfy a predicate.
 filter :: (v -> Bool) -> HashMap k v -> HashMap k v
 filter p = filterWithKey (\_ v -> p v)
 {-# INLINE filter #-}
 
+-- | /O(n)/ Filter this map by retaining only elements which values
+-- satisfy a predicate.
+-- Returns a tuple with the new size of the result hashmap, and the hashmap
+-- itself.
+filterInternal :: (v -> Bool) -> HashMap k v -> (Int, HashMap k v)
+filterInternal p = filterWithKeyInternal (\_ v -> p v)
+{-# INLINE filterInternal #-}
+
 ------------------------------------------------------------------------
 -- * Conversions