@@ -294,6 +294,13 @@ module Data.IntMap.Internal (
294294 , Stack (.. )
295295 , ascLinkTop
296296 , ascLinkAll
297+ , IntMapBuilder (.. )
298+ , BStack (.. )
299+ , emptyB
300+ , insertB
301+ , finishB
302+ , moveToB
303+ , MoveResult (.. )
297304
298305 -- * Used by "IntMap.Merge.Lazy" and "IntMap.Merge.Strict"
299306 , mapWhenMissing
@@ -3321,20 +3328,25 @@ foldlFB = foldlWithKey
33213328
33223329
33233330-- | \(O(n \min(n,W))\). Create a map from a list of key\/value pairs.
3331+ -- If the list contains more than one value for the same key, the last value
3332+ -- for the key is retained.
3333+ --
3334+ -- If the keys are in sorted order, ascending or descending, this function
3335+ -- takes \(O(n)\) time.
33243336--
33253337-- > fromList [] == empty
33263338-- > fromList [(5,"a"), (3,"b"), (5, "c")] == fromList [(5,"c"), (3,"b")]
33273339-- > fromList [(5,"c"), (3,"b"), (5, "a")] == fromList [(5,"a"), (3,"b")]
33283340
33293341fromList :: [(Key ,a )] -> IntMap a
3330- fromList xs
3331- = Foldable. foldl' ins empty xs
3332- where
3333- ins t (k,x) = insert k x t
3342+ fromList xs = finishB (Foldable. foldl' (\ b (kx,x) -> insertB kx x b) emptyB xs)
33343343{-# INLINE fromList #-} -- Inline for list fusion
33353344
33363345-- | \(O(n \min(n,W))\). Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
33373346--
3347+ -- If the keys are in sorted order, ascending or descending, this function
3348+ -- takes \(O(n)\) time.
3349+ --
33383350-- > fromListWith (++) [(5,"a"), (5,"b"), (3,"x"), (5,"c")] == fromList [(3, "x"), (5, "cba")]
33393351-- > fromListWith (++) [] == empty
33403352--
@@ -3376,17 +3388,18 @@ fromListWith f xs
33763388
33773389-- | \(O(n \min(n,W))\). Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'.
33783390--
3391+ -- If the keys are in sorted order, ascending or descending, this function
3392+ -- takes \(O(n)\) time.
3393+ --
33793394-- > let f key new_value old_value = show key ++ ":" ++ new_value ++ "|" ++ old_value
33803395-- > fromListWithKey f [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"c")] == fromList [(3, "3:a|b"), (5, "5:c|5:b|a")]
33813396-- > fromListWithKey f [] == empty
33823397--
33833398-- Also see the performance note on 'fromListWith'.
33843399
33853400fromListWithKey :: (Key -> a -> a -> a ) -> [(Key ,a )] -> IntMap a
3386- fromListWithKey f xs
3387- = Foldable. foldl' ins empty xs
3388- where
3389- ins t (k,x) = insertWithKey f k x t
3401+ fromListWithKey f xs =
3402+ finishB (Foldable. foldl' (\ b (kx,x) -> insertWithB (f kx) kx x b) emptyB xs)
33903403{-# INLINE fromListWithKey #-} -- Inline for list fusion
33913404
33923405-- | \(O(n)\). Build a map from a list of key\/value pairs where
@@ -3491,6 +3504,137 @@ ascLinkStack stk !rk r = case stk of
34913504 where
34923505 p = mask rk m
34933506
3507+ {- -------------------------------------------------------------------
3508+ IntMapBuilder
3509+ --------------------------------------------------------------------}
3510+
3511+ -- Note [IntMapBuilder]
3512+ -- ~~~~~~~~~~~~~~~~~~~~
3513+ -- IntMapBuilder serves as an accumulator for element-by-element construction
3514+ -- of an IntMap. It can be used in folds to construct IntMaps. This plays nicely
3515+ -- with list fusion when the structure folded over is a list, as in fromList and
3516+ -- friends.
3517+ --
3518+ -- An IntMapBuilder is either empty (BNil) or has the recently inserted Tip
3519+ -- together with a stack of trees (BTip). The structure is effectively a
3520+ -- [zipper](https://en.wikipedia.org/wiki/Zipper_(data_structure)). It always
3521+ -- has its "focus" at the last inserted entry. To insert a new entry, we need
3522+ -- to move the focus to the new entry. To do this we move up the stack to the
3523+ -- lowest common ancestor of the currest position and the position of the
3524+ -- new key (implemented as moveUpB), then down to the position of the new key
3525+ -- (implemented as moveDownB).
3526+ --
3527+ -- When we are done inserting entries, we link the trees up the stack and get
3528+ -- the final result.
3529+ --
3530+ -- The advantage of this implementation is that we take the shortest path in
3531+ -- the tree from one key to the next. Unlike `insert`, we don't need to move
3532+ -- up to the root after every insertion. This is very beneficial when we have
3533+ -- runs of sorted keys, without many keys already in the tree in that range.
3534+ -- If the keys are fully sorted, inserting them all takes O(n) time instead
3535+ -- of O(n min(n,W)). But these benefits come at a small cost: when moving up
3536+ -- the tree we have to check at every point if it is time to move down. These
3537+ -- checks are absent in `insert`. So, in case we need to move up quite a lot,
3538+ -- repeated `insert` is slightly faster, but the trade-off is worthwhile since
3539+ -- such cases are pathological.
3540+
3541+ data IntMapBuilder a
3542+ = BNil
3543+ | BTip {- # UNPACK #-} Int a ! (BStack a )
3544+
3545+ -- BLeft: the IntMap is the left child
3546+ -- BRight: the IntMap is the right child
3547+ data BStack a
3548+ = BNada
3549+ | BLeft {- # UNPACK #-} Prefix ! (IntMap a ) ! (BStack a )
3550+ | BRight {- # UNPACK #-} Prefix ! (IntMap a ) ! (BStack a )
3551+
3552+ -- Empty builder.
3553+ emptyB :: IntMapBuilder a
3554+ emptyB = BNil
3555+
3556+ -- Insert a key and value. Replaces the old value if one already exists for
3557+ -- the key.
3558+ insertB :: Key -> a -> IntMapBuilder a -> IntMapBuilder a
3559+ insertB ! ky y b = case b of
3560+ BNil -> BTip ky y BNada
3561+ BTip kx x stk -> case moveToB ky kx x stk of
3562+ MoveResult _ stk' -> BTip ky y stk'
3563+ {-# INLINE insertB #-}
3564+
3565+ -- Insert a key and value. The new value is combined with the old value if one
3566+ -- already exists for the key.
3567+ insertWithB :: (a -> a -> a ) -> Key -> a -> IntMapBuilder a -> IntMapBuilder a
3568+ insertWithB f ! ky y b = case b of
3569+ BNil -> BTip ky y BNada
3570+ BTip kx x stk -> case moveToB ky kx x stk of
3571+ MoveResult m stk' -> case m of
3572+ Nothing -> BTip ky y stk'
3573+ Just x' -> BTip ky (f y x') stk'
3574+ {-# INLINE insertWithB #-}
3575+
3576+ -- GHC >=9.6 supports unpacking sums, so we unpack the Maybe and avoid
3577+ -- allocating Justs. GHC optimizes the workers for moveUpB and moveDownB to
3578+ -- return (# (# (# #) | a #), BStack a #).
3579+ data MoveResult a
3580+ = MoveResult
3581+ #if __GLASGOW_HASKELL__ >= 906
3582+ {-# UNPACK #-}
3583+ #endif
3584+ ! (Maybe a)
3585+ ! (BStack a)
3586+
3587+ moveToB :: Key -> Key -> a -> BStack a -> MoveResult a
3588+ moveToB ! ky ! kx x ! stk
3589+ | kx == ky = MoveResult (Just x) stk
3590+ | otherwise = moveUpB ky kx (Tip kx x) stk
3591+ -- Don't inline this; there is no benefit according to benchmarks.
3592+ {-# NOINLINE moveToB #-}
3593+
3594+ moveUpB :: Key -> Key -> IntMap a -> BStack a -> MoveResult a
3595+ moveUpB ! ky ! kx ! tx stk = case stk of
3596+ BNada -> MoveResult Nothing (linkB ky kx tx BNada )
3597+ BLeft p l stk'
3598+ | nomatch ky p -> moveUpB ky kx (Bin p l tx) stk'
3599+ | left ky p -> moveDownB ky l (BRight p tx stk')
3600+ | otherwise -> MoveResult Nothing (linkB ky kx tx stk)
3601+ BRight p r stk'
3602+ | nomatch ky p -> moveUpB ky kx (Bin p tx r) stk'
3603+ | left ky p -> MoveResult Nothing (linkB ky kx tx stk)
3604+ | otherwise -> moveDownB ky r (BLeft p tx stk')
3605+
3606+ moveDownB :: Key -> IntMap a -> BStack a -> MoveResult a
3607+ moveDownB ! ky tx ! stk = case tx of
3608+ Bin p l r
3609+ | nomatch ky p -> MoveResult Nothing (linkB ky (unPrefix p) tx stk)
3610+ | left ky p -> moveDownB ky l (BRight p r stk)
3611+ | otherwise -> moveDownB ky r (BLeft p l stk)
3612+ Tip kx x
3613+ | kx == ky -> MoveResult (Just x) stk
3614+ | otherwise -> MoveResult Nothing (linkB ky kx tx stk)
3615+ Nil -> error " moveDownB Tip"
3616+
3617+ linkB :: Key -> Key -> IntMap a -> BStack a -> BStack a
3618+ linkB ky kx tx stk
3619+ | i2w ky < i2w kx = BRight p tx stk
3620+ | otherwise = BLeft p tx stk
3621+ where
3622+ p = branchPrefix ky kx
3623+ {-# INLINE linkB #-}
3624+
3625+ -- Finalize the builder into a Map.
3626+ finishB :: IntMapBuilder a -> IntMap a
3627+ finishB b = case b of
3628+ BNil -> Nil
3629+ BTip kx x stk -> finishUpB (Tip kx x) stk
3630+ {-# INLINABLE finishB #-}
3631+
3632+ finishUpB :: IntMap a -> BStack a -> IntMap a
3633+ finishUpB ! t stk = case stk of
3634+ BNada -> t
3635+ BLeft p l stk' -> finishUpB (Bin p l t) stk'
3636+ BRight p r stk' -> finishUpB (Bin p t r) stk'
3637+
34943638{- -------------------------------------------------------------------
34953639 Eq
34963640--------------------------------------------------------------------}
0 commit comments