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 b230a574e..757d16cfb 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 @@ -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 @@ -3079,6 +3100,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 +3115,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 +3147,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 +3165,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 +3186,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 +3240,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 +3253,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 +3339,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 +3349,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 +3367,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 +3377,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 +3397,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 +3409,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 +3427,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 +3438,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 +4433,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