Also apply Bin 1 check to filter. It makes it 20% faster

haskell · Nov 15, 2024 · e75e359 · e75e359
1 parent 2f1eebc
commit e75e359
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/containers-tests/benchmarks/Map.hs b/containers-tests/benchmarks/Map.hs
@@ -27,6 +27,8 @@ main = do
         , bench "lookup present" $ whnf (lookup evens) m_even
         , bench "map" $ whnf (M.map (+ 1)) m
         , bench "map really" $ nf (M.map (+ 2)) m
+        , bench "filter" $ whnf (M.filter even) m
+        , bench "filter really" $ nf (M.filter even) m
         , bench "<$" $ whnf ((1 :: Int) <$) m
         , bench "<$ really" $ nf ((2 :: Int) <$) m
         , bench "alterF lookup absent" $ whnf (atLookup evens) m_odd

diff --git a/containers/src/Data/Map/Internal.hs b/containers/src/Data/Map/Internal.hs
@@ -2966,6 +2966,9 @@ filter p m
 
 filterWithKey :: (k -> a -> Bool) -> Map k a -> Map k a
 filterWithKey _ Tip = Tip
+filterWithKey p t@(Bin 1 kx x _ _)
+  | p kx x    = t
+  | otherwise = Tip
 filterWithKey p t@(Bin _ kx x l r)
   | p kx x    = if pl `ptrEq` l && pr `ptrEq` r
                 then t