Use pattern match on 1 to reduce recursive function calls

Author: HuwCampbell

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
 

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