@@ -367,6 +367,10 @@ module Data.Map.Internal (
367367 , Identity (.. )
368368 , Stack (.. )
369369 , foldl'Stack
370+ , MapBuilder (.. )
371+ , emptyB
372+ , insertB
373+ , finishB
370374
371375 -- Used by Map.Merge.Lazy
372376 , mapWhenMissing
@@ -388,7 +392,6 @@ import Data.Semigroup (Semigroup((<>)))
388392#endif
389393import Control.Applicative (Const (.. ))
390394import Control.DeepSeq (NFData (rnf ),NFData1 (liftRnf ),NFData2 (liftRnf2 ))
391- import Data.Bits (shiftL , shiftR )
392395import qualified Data.Foldable as Foldable
393396import Data.Bifoldable
394397import Utils.Containers.Internal.Prelude hiding
@@ -3247,6 +3250,8 @@ mapAccumRWithKey f a (Bin sx kx x l r) =
32473250-- | \(O(n \log n)\).
32483251-- @'mapKeys' f s@ is the map obtained by applying @f@ to each key of @s@.
32493252--
3253+ -- If `f` is monotonically non-decreasing, this function takes \(O(n)\) time.
3254+ --
32503255-- The size of the result may be smaller if @f@ maps two or more distinct
32513256-- keys to the same new key. In this case the value at the greatest of the
32523257-- original keys is retained.
@@ -3256,14 +3261,16 @@ mapAccumRWithKey f a (Bin sx kx x l r) =
32563261-- > mapKeys (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "c"
32573262
32583263mapKeys :: Ord k2 => (k1 -> k2 ) -> Map k1 a -> Map k2 a
3259- mapKeys f = fromList . foldrWithKey ( \ k x xs -> (f k, x) : xs) []
3264+ mapKeys f m = finishB (foldlWithKey' ( \ b kx x -> insertB (f kx) x b) emptyB m)
32603265#if __GLASGOW_HASKELL__
32613266{-# INLINABLE mapKeys #-}
32623267#endif
32633268
32643269-- | \(O(n \log n)\).
32653270-- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@.
32663271--
3272+ -- If `f` is monotonically non-decreasing, this function takes \(O(n)\) time.
3273+ --
32673274-- The size of the result may be smaller if @f@ maps two or more distinct
32683275-- keys to the same new key. In this case the associated values will be
32693276-- combined using @c@. The value at the greater of the two original keys
@@ -3275,7 +3282,8 @@ mapKeys f = fromList . foldrWithKey (\k x xs -> (f k, x) : xs) []
32753282-- Also see the performance note on 'fromListWith'.
32763283
32773284mapKeysWith :: Ord k2 => (a -> a -> a ) -> (k1 -> k2 ) -> Map k1 a -> Map k2 a
3278- mapKeysWith c f = fromListWith c . foldrWithKey (\ k x xs -> (f k, x) : xs) []
3285+ mapKeysWith c f m =
3286+ finishB (foldlWithKey' (\ b kx x -> insertWithB c (f kx) x b) emptyB m)
32793287#if __GLASGOW_HASKELL__
32803288{-# INLINABLE mapKeysWith #-}
32813289#endif
@@ -3520,55 +3528,20 @@ instance (Ord k) => GHCExts.IsList (Map k v) where
35203528-- If the list contains more than one value for the same key, the last value
35213529-- for the key is retained.
35223530--
3523- -- If the keys of the list are ordered, a linear-time implementation is used .
3531+ -- If the keys are in non-decreasing order, this function takes \(O(n)\) time .
35243532--
35253533-- > fromList [] == empty
35263534-- > fromList [(5,"a"), (3,"b"), (5, "c")] == fromList [(5,"c"), (3,"b")]
35273535-- > fromList [(5,"c"), (3,"b"), (5, "a")] == fromList [(5,"a"), (3,"b")]
35283536
3529- -- For some reason, when 'singleton' is used in fromList or in
3530- -- create, it is not inlined, so we inline it manually.
35313537fromList :: Ord k => [(k ,a )] -> Map k a
3532- fromList [] = Tip
3533- fromList [(kx, x)] = Bin 1 kx x Tip Tip
3534- fromList ((kx0, x0) : xs0) | not_ordered kx0 xs0 = fromList' (Bin 1 kx0 x0 Tip Tip ) xs0
3535- | otherwise = go (1 :: Int Bin 1 kx0 x0 Tip Tip ) xs0
3536- where
3537- not_ordered _ [] = False
3538- not_ordered kx ((ky,_) : _) = kx >= ky
3539- {-# INLINE not_ordered #-}
3540-
3541- fromList' t0 xs = Foldable. foldl' ins t0 xs
3542- where ins t (k,x) = insert k x t
3543-
3544- go ! _ t [] = t
3545- go _ t [(kx, x)] = insertMax kx x t
3546- go s l xs@ ((kx, x) : xss) | not_ordered kx xss = fromList' l xs
3547- | otherwise = case create s xss of
3548- (r, ys, [] ) -> go (s `shiftL` 1 ) (link kx x l r) ys
3549- (r, _, ys) -> fromList' (link kx x l r) ys
3550-
3551- -- The create is returning a triple (tree, xs, ys). Both xs and ys
3552- -- represent not yet processed elements and only one of them can be nonempty.
3553- -- If ys is nonempty, the keys in ys are not ordered with respect to tree
3554- -- and must be inserted using fromList'. Otherwise the keys have been
3555- -- ordered so far.
3556- create ! _ [] = (Tip , [] , [] )
3557- create s xs@ (xp : xss)
3558- | s == 1 = case xp of (kx, x) | not_ordered kx xss -> (Bin 1 kx x Tip Tip , [] , xss)
3559- | otherwise -> (Bin 1 kx x Tip Tip , xss, [] )
3560- | otherwise = case create (s `shiftR` 1 ) xs of
3561- res@ (_, [] , _) -> res
3562- (l, [(ky, y)], zs) -> (insertMax ky y l, [] , zs)
3563- (l, ys@ ((ky, y): yss), _) | not_ordered ky yss -> (l, [] , ys)
3564- | otherwise -> case create (s `shiftR` 1 ) yss of
3565- (r, zs, ws) -> (link ky y l r, zs, ws)
3566- #if __GLASGOW_HASKELL__
3567- {-# INLINABLE fromList #-}
3568- #endif
3538+ fromList xs = finishB (Foldable. foldl' (\ b (kx, x) -> insertB kx x b) emptyB xs)
3539+ {-# INLINE fromList #-} -- INLINE for fusion
35693540
35703541-- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
35713542--
3543+ -- If the keys are in non-decreasing order, this function takes \(O(n)\) time.
3544+ --
35723545-- > fromListWith (++) [(5,"a"), (5,"b"), (3,"x"), (5,"c")] == fromList [(3, "x"), (5, "cba")]
35733546-- > fromListWith (++) [] == empty
35743547--
@@ -3604,28 +3577,24 @@ fromList ((kx0, x0) : xs0) | not_ordered kx0 xs0 = fromList' (Bin 1 kx0 x0 Tip T
36043577-- > fromListWith (++) $ reverse $ map (\(k, v) -> (k, [v])) someListOfTuples
36053578
36063579fromListWith :: Ord k => (a -> a -> a ) -> [(k ,a )] -> Map k a
3607- fromListWith f xs
3608- = fromListWithKey (\ _ x y -> f x y) xs
3609- #if __GLASGOW_HASKELL__
3610- {-# INLINABLE fromListWith #-}
3611- #endif
3580+ fromListWith f xs =
3581+ finishB (Foldable. foldl' (\ b (kx, x) -> insertWithB f kx x b) emptyB xs)
3582+ {-# INLINE fromListWith #-} -- INLINE for fusion
36123583
36133584-- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'.
36143585--
3586+ -- If the keys are in non-decreasing order, this function takes \(O(n)\) time.
3587+ --
36153588-- > let f key new_value old_value = show key ++ ":" ++ new_value ++ "|" ++ old_value
36163589-- > fromListWithKey f [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"c")] == fromList [(3, "3:a|b"), (5, "5:c|5:b|a")]
36173590-- > fromListWithKey f [] == empty
36183591--
36193592-- Also see the performance note on 'fromListWith'.
36203593
36213594fromListWithKey :: Ord k => (k -> a -> a -> a ) -> [(k ,a )] -> Map k a
3622- fromListWithKey f xs
3623- = Foldable. foldl' ins empty xs
3624- where
3625- ins t (k,x) = insertWithKey f k x t
3626- #if __GLASGOW_HASKELL__
3627- {-# INLINABLE fromListWithKey #-}
3628- #endif
3595+ fromListWithKey f xs =
3596+ finishB (Foldable. foldl' (\ b (kx, x) -> insertWithB (f kx) kx x b) emptyB xs)
3597+ {-# INLINE fromListWithKey #-} -- INLINE for fusion
36293598
36303599-- | \(O(n)\). Convert the map to a list of key\/value pairs. Subject to list fusion.
36313600--
@@ -3971,6 +3940,57 @@ splitMember k0 m = case go k0 m of
39713940
39723941data StrictTriple a b c = StrictTriple ! a ! b ! c
39733942
3943+ {- -------------------------------------------------------------------
3944+ MapBuilder
3945+ --------------------------------------------------------------------}
3946+
3947+ -- See Note [SetBuilder] in Data.Set.Internal
3948+
3949+ data MapBuilder k a
3950+ = BAsc ! (Stack k a )
3951+ | BMap ! (Map k a )
3952+
3953+ -- Empty builder.
3954+ emptyB :: MapBuilder k a
3955+ emptyB = BAsc Nada
3956+
3957+ -- Insert a key and value. Replaces the old value if one already exists for
3958+ -- the key.
3959+ insertB :: Ord k => k -> a -> MapBuilder k a -> MapBuilder k a
3960+ insertB ! ky y b = case b of
3961+ BAsc stk -> case stk of
3962+ Push kx x l stk' -> case compare ky kx of
3963+ LT -> BMap (insert ky y (ascLinkAll stk))
3964+ EQ -> BAsc (Push ky y l stk')
3965+ GT -> case l of
3966+ Tip -> BAsc (ascLinkTop stk' 1 (singleton kx x) ky y)
3967+ Bin {} -> BAsc (Push ky y Tip stk)
3968+ Nada -> BAsc (Push ky y Tip Nada )
3969+ BMap m -> BMap (insert ky y m)
3970+ {-# INLINE insertB #-}
3971+
3972+ -- Insert a key and value. The new value is combined with the old value if one
3973+ -- already exists for the key.
3974+ insertWithB
3975+ :: Ord k => (a -> a -> a ) -> k -> a -> MapBuilder k a -> MapBuilder k a
3976+ insertWithB f ! ky y b = case b of
3977+ BAsc stk -> case stk of
3978+ Push kx x l stk' -> case compare ky kx of
3979+ LT -> BMap (insertWith f ky y (ascLinkAll stk))
3980+ EQ -> BAsc (Push ky (f y x) l stk')
3981+ GT -> case l of
3982+ Tip -> BAsc (ascLinkTop stk' 1 (singleton kx x) ky y)
3983+ Bin {} -> BAsc (Push ky y Tip stk)
3984+ Nada -> BAsc (Push ky y Tip Nada )
3985+ BMap m -> BMap (insertWith f ky y m)
3986+ {-# INLINE insertWithB #-}
3987+
3988+ -- Finalize the builder into a Map.
3989+ finishB :: MapBuilder k a -> Map k a
3990+ finishB (BAsc stk) = ascLinkAll stk
3991+ finishB (BMap m) = m
3992+ {-# INLINABLE finishB #-}
3993+
39743994{- -------------------------------------------------------------------
39753995 Utility functions that maintain the balance properties of the tree.
39763996 All constructors assume that all values in [l] < [k] and all values
0 commit comments