Map performance - use pattern match on size to reduce recursive function calls

containers-tests/benchmarks/Map.hs

@@ -27,6 +27,9 @@ 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 "partition" $ whnf (M.partition even) m
         , bench "<$" $ whnf ((1 :: Int) <$) m
         , bench "<$ really" $ nf ((2 :: Int) <$) m
         , bench "alterF lookup absent" $ whnf (atLookup evens) m_odd

containers/src/Data/Map/Internal.hs

@@ -83,7 +83,7 @@
 -- [Note: Using INLINABLE]
 -- ~~~~~~~~~~~~~~~~~~~~~~~
 -- It is crucial to the performance that the functions specialize on the Ord
--- type when possible. GHC 7.0 and higher does this by itself when it sees th
+-- type when possible. GHC 7.0 and higher does this by itself when it sees the
 -- unfolding of a function -- that is why all public functions are marked
 -- INLINABLE (that exposes the unfolding).
 
@@ -116,7 +116,7 @@
 -- floats out of its enclosing function and then it heap-allocates the
 -- dictionary and the argument. Maybe it floats out too late and strictness
 -- analyzer cannot see that these could be passed on stack.
---
+
 
 -- [Note: Order of constructors]
 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -127,6 +127,22 @@
 -- On GHC 7.0, reordering constructors from Tip | Bin to Bin | Tip
 -- improves the benchmark by up to 10% on x86.
 
+
+-- [Note: Matching on Leafy Nodes]
+-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+-- In a balanced tree, at least two-thirds of Tip constructors are siblings
+-- of another Tip constructor. The parents of these cases can be quickly
+-- identified as the size value packed into their Bin constructors will equal
+-- 1. By specializing recursive functions which visit the whole tree to
+-- recognize this scenario, we can elide unnecessary function calls that would
+-- go on to match these Tip constructors but otherwise perform no useful work.
+-- This optimization can lead to performance improvements of approximately
+-- 30% to 35% for foldMap and foldl', and around 20% for mapMaybe.
+--
+-- Alternatives, like matching on the Tip constructors directly, or also
+-- trying to optimise cases where only one side a Tip are slower in practice.
+
+
 module Data.Map.Internal (
     -- * Map type
       Map(..)          -- instance Eq,Show,Read
@@ -2950,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
@@ -2962,6 +2981,8 @@ filterWithKey p t@(Bin _ kx x l r)
 -- predicate.
 filterWithKeyA :: Applicative f => (k -> a -> f Bool) -> Map k a -> f (Map k a)
 filterWithKeyA _ Tip = pure Tip
+filterWithKeyA p t@(Bin 1 kx x _ _) =
+  fmap (bool Tip t) (p kx x)
 filterWithKeyA p t@(Bin _ kx x l r) =
   liftA3 combine (filterWithKeyA p l) (p kx x) (filterWithKeyA p r)
   where
@@ -3052,6 +3073,9 @@ partitionWithKey :: (k -> a -> Bool) -> Map k a -> (Map k a,Map k a)
 partitionWithKey p0 t0 = toPair $ go p0 t0
   where
     go _ Tip = (Tip :*: Tip)
+    go p t@(Bin 1 kx x _ _)
+      | p kx x    = t :*: Tip
+      | otherwise = Tip :*: t
     go p t@(Bin _ kx x l r)
       | p kx x    = (if l1 `ptrEq` l && r1 `ptrEq` r
                      then t
@@ -3079,6 +3103,9 @@ mapMaybe f = mapMaybeWithKey (\_ x -> f x)
 
 mapMaybeWithKey :: (k -> a -> Maybe b) -> Map k a -> Map k b
 mapMaybeWithKey _ Tip = Tip
+mapMaybeWithKey f (Bin 1 kx x _ _) = case f kx x of
+  Just y  -> Bin 1 kx y Tip Tip
+  Nothing -> Tip
 mapMaybeWithKey f (Bin _ kx x l r) = case f kx x of
   Just y  -> link kx y (mapMaybeWithKey f l) (mapMaybeWithKey f r)
   Nothing -> link2 (mapMaybeWithKey f l) (mapMaybeWithKey f r)
@@ -3091,7 +3118,7 @@ traverseMaybeWithKey :: Applicative f
 traverseMaybeWithKey = go
   where
     go _ Tip = pure Tip
-    go f (Bin _ kx x Tip Tip) = maybe Tip (\x' -> Bin 1 kx x' Tip Tip) <$> f kx x
+    go f (Bin 1 kx x _ _) = maybe Tip (\x' -> Bin 1 kx x' Tip Tip) <$> f kx x
     go f (Bin _ kx x l r) = liftA3 combine (go f l) (f kx x) (go f r)
       where
         combine !l' mx !r' = case mx of
@@ -3123,7 +3150,7 @@ mapEither f m
 mapEitherWithKey :: (k -> a -> Either b c) -> Map k a -> (Map k b, Map k c)
 mapEitherWithKey f0 t0 = toPair $ go f0 t0
   where
-    go _ Tip = (Tip :*: Tip)
+    go _ Tip = Tip :*: Tip
     go f (Bin _ kx x l r) = case f kx x of
       Left y  -> link kx y l1 r1 :*: link2 l2 r2
       Right z -> link2 l1 r1 :*: link kx z l2 r2
@@ -3141,6 +3168,7 @@ mapEitherWithKey f0 t0 = toPair $ go f0 t0
 map :: (a -> b) -> Map k a -> Map k b
 map f = go where
   go Tip = Tip
+  go (Bin 1 kx x _ _) = Bin 1 kx (f x) Tip Tip
   go (Bin sx kx x l r) = Bin sx kx (f x) (go l) (go r)
 -- We use a `go` function to allow `map` to inline. This makes
 -- a big difference if someone uses `map (const x) m` instead
@@ -3161,6 +3189,7 @@ map f = go where
 
 mapWithKey :: (k -> a -> b) -> Map k a -> Map k b
 mapWithKey _ Tip = Tip
+mapWithKey f (Bin 1 kx x _ _) = Bin 1 kx (f kx x) Tip Tip
 mapWithKey f (Bin sx kx x l r) = Bin sx kx (f kx x) (mapWithKey f l) (mapWithKey f r)
 
 #ifdef __GLASGOW_HASKELL__
@@ -3214,6 +3243,9 @@ mapAccumWithKey f a t
 -- argument through the map in ascending order of keys.
 mapAccumL :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c)
 mapAccumL _ a Tip               = (a,Tip)
+mapAccumL f a (Bin 1 kx x _ _ ) =
+  let (a1,x') = f a kx x
+  in (a1,Bin 1 kx x' Tip Tip)
 mapAccumL f a (Bin sx kx x l r) =
   let (a1,l') = mapAccumL f a l
       (a2,x') = f a1 kx x
@@ -3224,6 +3256,9 @@ mapAccumL f a (Bin sx kx x l r) =
 -- argument through the map in descending order of keys.
 mapAccumRWithKey :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c)
 mapAccumRWithKey _ a Tip = (a,Tip)
+mapAccumRWithKey f a (Bin 1 kx x _ _) =
+  let (a0,x') = f a kx x
+  in (a0,Bin 1 kx x' Tip Tip)
 mapAccumRWithKey f a (Bin sx kx x l r) =
   let (a1,r') = mapAccumRWithKey f a r
       (a2,x') = f a1 kx x
@@ -3307,6 +3342,7 @@ foldr :: (a -> b -> b) -> b -> Map k a -> b
 foldr f z = go z
   where
     go z' Tip             = z'
+    go z' (Bin 1 _ x _ _) = f x z'
     go z' (Bin _ _ x l r) = go (f x (go z' r)) l
 {-# INLINE foldr #-}
 
@@ -3316,8 +3352,9 @@ foldr f z = go z
 foldr' :: (a -> b -> b) -> b -> Map k a -> b
 foldr' f z = go z
   where
-    go !z' Tip            = z'
-    go z' (Bin _ _ x l r) = go (f x $! go z' r) l
+    go !z' Tip             = z'
+    go !z' (Bin 1 _ x _ _) = f x z'
+    go z'  (Bin _ _ x l r) = go (f x $! go z' r) l
 {-# INLINE foldr' #-}
 
 -- | \(O(n)\). Fold the values in the map using the given left-associative
@@ -3333,6 +3370,7 @@ foldl :: (a -> b -> a) -> a -> Map k b -> a
 foldl f z = go z
   where
     go z' Tip             = z'
+    go z' (Bin 1 _ x _ _) = f z' x
     go z' (Bin _ _ x l r) = go (f (go z' l) x) r
 {-# INLINE foldl #-}
 
@@ -3342,8 +3380,9 @@ foldl f z = go z
 foldl' :: (a -> b -> a) -> a -> Map k b -> a
 foldl' f z = go z
   where
-    go !z' Tip            = z'
-    go z' (Bin _ _ x l r) =
+    go !z' Tip             = z'
+    go !z' (Bin 1 _ x _ _) = f z' x
+    go z' (Bin _ _ x l r)  =
       let !z'' = go z' l
       in go (f z'' x) r
 {-# INLINE foldl' #-}
@@ -3361,7 +3400,8 @@ foldl' f z = go z
 foldrWithKey :: (k -> a -> b -> b) -> b -> Map k a -> b
 foldrWithKey f z = go z
   where
-    go z' Tip             = z'
+    go z' Tip              = z'
+    go z' (Bin 1 kx x _ _) = f kx x z'
     go z' (Bin _ kx x l r) = go (f kx x (go z' r)) l
 {-# INLINE foldrWithKey #-}
 
@@ -3372,7 +3412,8 @@ foldrWithKey' :: (k -> a -> b -> b) -> b -> Map k a -> b
 foldrWithKey' f z = go z
   where
     go !z' Tip              = z'
-    go z' (Bin _ kx x l r) = go (f kx x $! go z' r) l
+    go !z' (Bin 1 kx x _ _) = f kx x z'
+    go z' (Bin _ kx x l r)  = go (f kx x $! go z' r) l
 {-# INLINE foldrWithKey' #-}
 
 -- | \(O(n)\). Fold the keys and values in the map using the given left-associative
@@ -3389,6 +3430,7 @@ foldlWithKey :: (a -> k -> b -> a) -> a -> Map k b -> a
 foldlWithKey f z = go z
   where
     go z' Tip              = z'
+    go z' (Bin 1 kx x _ _) = f z' kx x
     go z' (Bin _ kx x l r) = go (f (go z' l) kx x) r
 {-# INLINE foldlWithKey #-}
 
@@ -3399,6 +3441,7 @@ foldlWithKey' :: (a -> k -> b -> a) -> a -> Map k b -> a
 foldlWithKey' f z = go z
   where
     go !z' Tip             = z'
+    go !z' (Bin 1 kx x _ _) = f z' kx x
     go z' (Bin _ kx x l r) =
       let !z'' = go z' l
       in go (f z'' kx x) r
@@ -4393,6 +4436,7 @@ instance Functor (Map k) where
   fmap f m  = map f m
 #ifdef __GLASGOW_HASKELL__
   _ <$ Tip = Tip
+  a <$ (Bin 1 kx _ _ _) = Bin 1 kx a Tip Tip
   a <$ (Bin sx kx _ l r) = Bin sx kx a (a <$ l) (a <$ r)
 #endif
 

containers/src/Data/Map/Strict/Internal.hs

@@ -1271,6 +1271,9 @@ mapMaybe f = mapMaybeWithKey (\_ x -> f x)
 
 mapMaybeWithKey :: (k -> a -> Maybe b) -> Map k a -> Map k b
 mapMaybeWithKey _ Tip = Tip
+mapMaybeWithKey f (Bin 1 kx x _ _) = case f kx x of
+  Just y  -> y `seq` Bin 1 kx y Tip Tip
+  Nothing -> Tip
 mapMaybeWithKey f (Bin _ kx x l r) = case f kx x of
   Just y  -> y `seq` link kx y (mapMaybeWithKey f l) (mapMaybeWithKey f r)
   Nothing -> link2 (mapMaybeWithKey f l) (mapMaybeWithKey f r)
@@ -1284,7 +1287,7 @@ traverseMaybeWithKey :: Applicative f
 traverseMaybeWithKey = go
   where
     go _ Tip = pure Tip
-    go f (Bin _ kx x Tip Tip) = maybe Tip (\ !x' -> Bin 1 kx x' Tip Tip) <$> f kx x
+    go f (Bin 1 kx x _ _) = maybe Tip (\ !x' -> Bin 1 kx x' Tip Tip) <$> f kx x
     go f (Bin _ kx x l r) = liftA3 combine (go f l) (f kx x) (go f r)
       where
         combine !l' mx !r' = case mx of
@@ -1335,6 +1338,7 @@ map :: (a -> b) -> Map k a -> Map k b
 map f = go
   where
     go Tip = Tip
+    go (Bin 1 kx x _ _) = let !x' = f x in Bin 1 kx x' Tip Tip
     go (Bin sx kx x l r) = let !x' = f x in Bin sx kx x' (go l) (go r)
 -- We use `go` to let `map` inline. This is important if `f` is a constant
 -- function.
@@ -1354,6 +1358,9 @@ map f = go
 
 mapWithKey :: (k -> a -> b) -> Map k a -> Map k b
 mapWithKey _ Tip = Tip
+mapWithKey f (Bin 1 kx x _ _) =
+  let x' = f kx x
+  in x' `seq` Bin 1 kx x' Tip Tip
 mapWithKey f (Bin sx kx x l r) =
   let x' = f kx x
   in x' `seq` Bin sx kx x' (mapWithKey f l) (mapWithKey f r)
@@ -1416,6 +1423,9 @@ mapAccumWithKey f a t
 -- argument through the map in ascending order of keys.
 mapAccumL :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c)
 mapAccumL _ a Tip               = (a,Tip)
+mapAccumL f a (Bin 1 kx x _ _) =
+  let (a1,x') = f a kx x
+  in x' `seq` (a1,Bin 1 kx x' Tip Tip)
 mapAccumL f a (Bin sx kx x l r) =
   let (a1,l') = mapAccumL f a l
       (a2,x') = f a1 kx x

containers/src/Data/Set/Internal.hs

@@ -1062,6 +1062,7 @@ foldl :: (a -> b -> a) -> a -> Set b -> a
 foldl f z = go z
   where
     go z' Tip           = z'
+    go z' (Bin 1 x _ _) = f z' x
     go z' (Bin _ x l r) = go (f (go z' l) x) r
 {-# INLINE foldl #-}
 
@@ -1072,6 +1073,7 @@ foldl' :: (a -> b -> a) -> a -> Set b -> a
 foldl' f z = go z
   where
     go !z' Tip           = z'
+    go !z' (Bin 1 x _ _) = f z' x
     go z' (Bin _ x l r) =
       let !z'' = go z' l
       in go (f z'' x) r