From 7e1b4f4678f01363b80a3ebf4d0dc8ecf0a7a529 Mon Sep 17 00:00:00 2001
From: Huw Campbell <hcampbell@simplemachines.com.au>
Date: Thu, 14 Nov 2024 10:36:41 +1100
Subject: [PATCH 1/4] Use pattern match on 1 to reduce recursive function calls

---
 containers/src/Data/Map/Internal.hs        | 36 +++++++++++++++++-----
 containers/src/Data/Map/Strict/Internal.hs | 12 +++++++-
 containers/src/Data/Set/Internal.hs        |  2 ++
 3 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/containers/src/Data/Map/Internal.hs b/containers/src/Data/Map/Internal.hs
index b230a574e..871a59981 100644
--- a/containers/src/Data/Map/Internal.hs
+++ b/containers/src/Data/Map/Internal.hs
@@ -3079,6 +3079,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 +3094,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 +3126,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 +3144,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 +3165,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 +3219,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 +3232,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 +3318,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 +3328,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 +3346,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 +3356,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 +3376,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 +3388,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 +3406,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 +3417,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 +4412,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
 
diff --git a/containers/src/Data/Map/Strict/Internal.hs b/containers/src/Data/Map/Strict/Internal.hs
index 21afe2d91..a8bbde101 100644
--- a/containers/src/Data/Map/Strict/Internal.hs
+++ b/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
diff --git a/containers/src/Data/Set/Internal.hs b/containers/src/Data/Set/Internal.hs
index a55877501..0c1c5ecff 100644
--- a/containers/src/Data/Set/Internal.hs
+++ b/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

From 2f1eebcba5ab9e4cf3f6f47ffeb27b43306cbf41 Mon Sep 17 00:00:00 2001
From: Huw Campbell <hcampbell@simplemachines.com.au>
Date: Fri, 15 Nov 2024 12:01:21 +1100
Subject: [PATCH 2/4] Add note explaining this optimisation

---
 containers/src/Data/Map/Internal.hs | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/containers/src/Data/Map/Internal.hs b/containers/src/Data/Map/Internal.hs
index 871a59981..063b47f46 100644
--- a/containers/src/Data/Map/Internal.hs
+++ b/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

From ec94da28eb85fe48ac989c6c1bb9a4b02b96b662 Mon Sep 17 00:00:00 2001
From: Huw Campbell <hcampbell@simplemachines.com.au>
Date: Fri, 15 Nov 2024 16:40:31 +1100
Subject: [PATCH 3/4] Also apply Bin 1 check to filter. It makes it 20% faster

---
 containers-tests/benchmarks/Map.hs  | 2 ++
 containers/src/Data/Map/Internal.hs | 5 +++++
 2 files changed, 7 insertions(+)

diff --git a/containers-tests/benchmarks/Map.hs b/containers-tests/benchmarks/Map.hs
index 67944277f..6e0d9e959 100644
--- 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
index 063b47f46..757d16cfb 100644
--- 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
@@ -2978,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

From 1a494efe16a611c85cd37c7c3da3e19faa66fc85 Mon Sep 17 00:00:00 2001
From: Huw Campbell <hcampbell@simplemachines.com.au>
Date: Tue, 19 Nov 2024 10:47:19 +1100
Subject: [PATCH 4/4] Apply to Bin 1 trick to partitionWithKey. Improves
 performance by 40%

---
 containers-tests/benchmarks/Map.hs  | 1 +
 containers/src/Data/Map/Internal.hs | 3 +++
 2 files changed, 4 insertions(+)

diff --git a/containers-tests/benchmarks/Map.hs b/containers-tests/benchmarks/Map.hs
index 6e0d9e959..2cc74aa5c 100644
--- a/containers-tests/benchmarks/Map.hs
+++ b/containers-tests/benchmarks/Map.hs
@@ -29,6 +29,7 @@ main = do
         , 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
diff --git a/containers/src/Data/Map/Internal.hs b/containers/src/Data/Map/Internal.hs
index 757d16cfb..487615c4d 100644
--- a/containers/src/Data/Map/Internal.hs
+++ b/containers/src/Data/Map/Internal.hs
@@ -3073,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