Skip to content

Commit 4c05817

Browse files
meooow25meooow25
committed
Build Set and Map more efficiently
Use "Builder"s to implement some Set and Map construction functions. As a result, some have become good consumers in terms of list fusion, and all are now O(n) for non-decreasing input. Fusible Fusible O(n) for O(n) for before after before after Set.fromList No Yes Strict incr Non-decr Set.map - - Strict incr Non-decr Map.fromList No Yes Strict incr Non-decr Map.fromListWith Yes Yes Never Non-decr Map.fromListWithKey Yes Yes Never Non-decr Map.mapKeys - - Strict incr Non-decr Map.mapKeysWith - - Never Non-decr
1 parent c2b3c15 commit 4c05817

File tree

3 files changed

+155
-150
lines changed

3 files changed

+155
-150
lines changed

‎containers/src/Data/Map/Internal.hs

Lines changed: 66 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -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
389393
import Control.Applicative (Const (..))
390394
import Control.DeepSeq (NFData(rnf))
391-
import Data.Bits (shiftL, shiftR)
392395
import qualified Data.Foldable as Foldable
393396
import Data.Bifoldable
394397
import Utils.Containers.Internal.Prelude hiding
@@ -3256,7 +3259,7 @@ mapAccumRWithKey f a (Bin sx kx x l r) =
32563259
-- > mapKeys (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "c"
32573260

32583261
mapKeys :: Ord k2 => (k1->k2) -> Map k1 a -> Map k2 a
3259-
mapKeys f = fromList . foldrWithKey (\k x xs -> (f k, x) : xs) []
3262+
mapKeys f m = finishB (foldlWithKey' (\b kx x -> insertB (f kx) x b) emptyB m)
32603263
#if __GLASGOW_HASKELL__
32613264
{-# INLINABLE mapKeys #-}
32623265
#endif
@@ -3275,7 +3278,8 @@ mapKeys f = fromList . foldrWithKey (\k x xs -> (f k, x) : xs) []
32753278
-- Also see the performance note on 'fromListWith'.
32763279

32773280
mapKeysWith :: 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) []
3281+
mapKeysWith c f m =
3282+
finishB (foldlWithKey' (\b kx x -> insertWithB c (f kx) x b) emptyB m)
32793283
#if __GLASGOW_HASKELL__
32803284
{-# INLINABLE mapKeysWith #-}
32813285
#endif
@@ -3526,46 +3530,9 @@ instance (Ord k) => GHCExts.IsList (Map k v) where
35263530
-- > fromList [(5,"a"), (3,"b"), (5, "c")] == fromList [(5,"c"), (3,"b")]
35273531
-- > fromList [(5,"c"), (3,"b"), (5, "a")] == fromList [(5,"a"), (3,"b")]
35283532

3529-
-- For some reason, when 'singleton' is used in fromList or in
3530-
-- create, it is not inlined, so we inline it manually.
35313533
fromList :: 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
3534+
fromList xs = finishB (Foldable.foldl' (\b (kx, x) -> insertB kx x b) emptyB xs)
3535+
{-# INLINE fromList #-} -- INLINE for fusion
35693536

35703537
-- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
35713538
--
@@ -3604,11 +3571,9 @@ fromList ((kx0, x0) : xs0) | not_ordered kx0 xs0 = fromList' (Bin 1 kx0 x0 Tip T
36043571
-- > fromListWith (++) $ reverse $ map (\(k, v) -> (k, [v])) someListOfTuples
36053572

36063573
fromListWith :: 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
3574+
fromListWith f xs =
3575+
finishB (Foldable.foldl' (\b (kx, x) -> insertWithB f kx x b) emptyB xs)
3576+
{-# INLINE fromListWith #-} -- INLINE for fusion
36123577

36133578
-- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'.
36143579
--
@@ -3619,13 +3584,9 @@ fromListWith f xs
36193584
-- Also see the performance note on 'fromListWith'.
36203585

36213586
fromListWithKey :: 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
3587+
fromListWithKey f xs =
3588+
finishB (Foldable.foldl' (\b (kx, x) -> insertWithB (f kx) kx x b) emptyB xs)
3589+
{-# INLINE fromListWithKey #-} -- INLINE for fusion
36293590

36303591
-- | \(O(n)\). Convert the map to a list of key\/value pairs. Subject to list fusion.
36313592
--
@@ -4020,6 +3981,57 @@ splitMember k0 m = case go k0 m of
40203981

40213982
data StrictTriple a b c = StrictTriple !a !b !c
40223983

3984+
{--------------------------------------------------------------------
3985+
MapBuilder
3986+
--------------------------------------------------------------------}
3987+
3988+
-- See Note [SetBuilder] in Data.Set.Internal
3989+
3990+
data MapBuilder k a
3991+
= BAsc !(Stack k a)
3992+
| BMap !(Map k a)
3993+
3994+
-- Empty builder.
3995+
emptyB :: MapBuilder k a
3996+
emptyB = BAsc Nada
3997+
3998+
-- Insert a key and value. Replaces the old value if one already exists for
3999+
-- the key.
4000+
insertB :: Ord k => k -> a -> MapBuilder k a -> MapBuilder k a
4001+
insertB !ky y b = case b of
4002+
BAsc stk -> case stk of
4003+
Push kx x l stk' -> case compare ky kx of
4004+
LT -> BMap (insert ky y (ascLinkAll stk))
4005+
EQ -> BAsc (Push ky y l stk')
4006+
GT -> case l of
4007+
Tip -> BAsc (ascLinkTop stk' 1 (singleton kx x) ky y)
4008+
Bin{} -> BAsc (Push ky y Tip stk)
4009+
Nada -> BAsc (Push ky y Tip Nada)
4010+
BMap m -> BMap (insert ky y m)
4011+
{-# INLINE insertB #-}
4012+
4013+
-- Insert a key and value. The new value is combined with the old value if one
4014+
-- already exists for the key.
4015+
insertWithB
4016+
:: Ord k => (a -> a -> a) -> k -> a -> MapBuilder k a -> MapBuilder k a
4017+
insertWithB f !ky y b = case b of
4018+
BAsc stk -> case stk of
4019+
Push kx x l stk' -> case compare ky kx of
4020+
LT -> BMap (insertWith f ky y (ascLinkAll stk))
4021+
EQ -> BAsc (Push ky (f y x) l stk')
4022+
GT -> case l of
4023+
Tip -> BAsc (ascLinkTop stk' 1 (singleton kx x) ky y)
4024+
Bin{} -> BAsc (Push ky y Tip stk)
4025+
Nada -> BAsc (Push ky y Tip Nada)
4026+
BMap m -> BMap (insertWith f ky y m)
4027+
{-# INLINE insertWithB #-}
4028+
4029+
-- Finalize the builder into a Map.
4030+
finishB :: MapBuilder k a -> Map k a
4031+
finishB (BAsc stk) = ascLinkAll stk
4032+
finishB (BMap m) = m
4033+
{-# INLINABLE finishB #-}
4034+
40234035
{--------------------------------------------------------------------
40244036
Utility functions that maintain the balance properties of the tree.
40254037
All constructors assume that all values in [l] < [k] and all values

‎containers/src/Data/Map/Strict/Internal.hs

Lines changed: 37 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -333,6 +333,10 @@ import Data.Map.Internal
333333
, descLinkTop
334334
, descLinkAll
335335
, Stack (..)
336+
, MapBuilder(..)
337+
, emptyB
338+
, insertB
339+
, finishB
336340
, (!)
337341
, (!?)
338342
, (\\)
@@ -375,7 +379,6 @@ import Data.Map.Internal
375379
, foldrWithKey
376380
, foldrWithKey'
377381
, glue
378-
, insertMax
379382
, intersection
380383
, isProperSubmapOf
381384
, isProperSubmapOfBy
@@ -430,7 +433,6 @@ import qualified Data.Set.Internal as Set
430433
import qualified Data.Map.Internal as L
431434
import Utils.Containers.Internal.StrictPair
432435

433-
import Data.Bits (shiftL, shiftR)
434436
#ifdef __GLASGOW_HASKELL__
435437
import Data.Coerce
436438
#endif
@@ -1448,7 +1450,8 @@ mapAccumRWithKey f a (Bin sx kx x l r) =
14481450
-- Also see the performance note on 'fromListWith'.
14491451

14501452
mapKeysWith :: Ord k2 => (a -> a -> a) -> (k1->k2) -> Map k1 a -> Map k2 a
1451-
mapKeysWith c f = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) []
1453+
mapKeysWith c f m =
1454+
finishB (foldlWithKey' (\b kx x -> insertWithB c (f kx) x b) emptyB m)
14521455
#if __GLASGOW_HASKELL__
14531456
{-# INLINABLE mapKeysWith #-}
14541457
#endif
@@ -1489,46 +1492,10 @@ fromArgSet (Set.Bin sz (Arg x v) l r) = v `seq` Bin sz x v (fromArgSet l) (fromA
14891492
-- > fromList [(5,"a"), (3,"b"), (5, "c")] == fromList [(5,"c"), (3,"b")]
14901493
-- > fromList [(5,"c"), (3,"b"), (5, "a")] == fromList [(5,"a"), (3,"b")]
14911494

1492-
-- For some reason, when 'singleton' is used in fromList or in
1493-
-- create, it is not inlined, so we inline it manually.
14941495
fromList :: Ord k => [(k,a)] -> Map k a
1495-
fromList [] = Tip
1496-
fromList [(kx, x)] = x `seq` Bin 1 kx x Tip Tip
1497-
fromList ((kx0, x0) : xs0) | not_ordered kx0 xs0 = x0 `seq` fromList' (Bin 1 kx0 x0 Tip Tip) xs0
1498-
| otherwise = x0 `seq` go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0
1499-
where
1500-
not_ordered _ [] = False
1501-
not_ordered kx ((ky,_) : _) = kx >= ky
1502-
{-# INLINE not_ordered #-}
1503-
1504-
fromList' t0 xs = Foldable.foldl' ins t0 xs
1505-
where ins t (k,x) = insert k x t
1506-
1507-
go !_ t [] = t
1508-
go _ t [(kx, x)] = x `seq` insertMax kx x t
1509-
go s l xs@((kx, x) : xss) | not_ordered kx xss = fromList' l xs
1510-
| otherwise = case create s xss of
1511-
(r, ys, []) -> x `seq` go (s `shiftL` 1) (link kx x l r) ys
1512-
(r, _, ys) -> x `seq` fromList' (link kx x l r) ys
1513-
1514-
-- The create is returning a triple (tree, xs, ys). Both xs and ys
1515-
-- represent not yet processed elements and only one of them can be nonempty.
1516-
-- If ys is nonempty, the keys in ys are not ordered with respect to tree
1517-
-- and must be inserted using fromList'. Otherwise the keys have been
1518-
-- ordered so far.
1519-
create !_ [] = (Tip, [], [])
1520-
create s xs@(xp : xss)
1521-
| s == 1 = case xp of (kx, x) | not_ordered kx xss -> x `seq` (Bin 1 kx x Tip Tip, [], xss)
1522-
| otherwise -> x `seq` (Bin 1 kx x Tip Tip, xss, [])
1523-
| otherwise = case create (s `shiftR` 1) xs of
1524-
res@(_, [], _) -> res
1525-
(l, [(ky, y)], zs) -> y `seq` (insertMax ky y l, [], zs)
1526-
(l, ys@((ky, y):yss), _) | not_ordered ky yss -> (l, [], ys)
1527-
| otherwise -> case create (s `shiftR` 1) yss of
1528-
(r, zs, ws) -> y `seq` (link ky y l r, zs, ws)
1529-
#if __GLASGOW_HASKELL__
1530-
{-# INLINABLE fromList #-}
1531-
#endif
1496+
fromList xs =
1497+
finishB (Foldable.foldl' (\b (kx, !x) -> insertB kx x b) emptyB xs)
1498+
{-# INLINE fromList #-} -- INLINE for fusion
15321499

15331500
-- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
15341501
--
@@ -1567,11 +1534,9 @@ fromList ((kx0, x0) : xs0) | not_ordered kx0 xs0 = x0 `seq` fromList' (Bin 1 kx0
15671534
-- > fromListWith (++) $ reverse $ map (\(k, v) -> (k, [v])) someListOfTuples
15681535

15691536
fromListWith :: Ord k => (a -> a -> a) -> [(k,a)] -> Map k a
1570-
fromListWith f xs
1571-
= fromListWithKey (\_ x y -> f x y) xs
1572-
#if __GLASGOW_HASKELL__
1573-
{-# INLINABLE fromListWith #-}
1574-
#endif
1537+
fromListWith f xs =
1538+
finishB (Foldable.foldl' (\b (kx, x) -> insertWithB f kx x b) emptyB xs)
1539+
{-# INLINE fromListWith #-} -- INLINE for fusion
15751540

15761541
-- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'.
15771542
--
@@ -1582,13 +1547,9 @@ fromListWith f xs
15821547
-- Also see the performance note on 'fromListWith'.
15831548

15841549
fromListWithKey :: Ord k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
1585-
fromListWithKey f xs
1586-
= Foldable.foldl' ins empty xs
1587-
where
1588-
ins t (k,x) = insertWithKey f k x t
1589-
#if __GLASGOW_HASKELL__
1590-
{-# INLINABLE fromListWithKey #-}
1591-
#endif
1550+
fromListWithKey f xs =
1551+
finishB (Foldable.foldl' (\b (kx, x) -> insertWithB (f kx) kx x b) emptyB xs)
1552+
{-# INLINE fromListWithKey #-} -- INLINE for fusion
15921553

15931554
{--------------------------------------------------------------------
15941555
Building trees from ascending/descending lists can be done in linear time.
@@ -1753,3 +1714,25 @@ fromDistinctDescList = descLinkAll . Foldable.foldl' next Nada
17531714
next (Push ky y Tip stk) (!kx, !x) = descLinkTop kx x 1 (singleton ky y) stk
17541715
next stk (!ky, !y) = Push ky y Tip stk
17551716
{-# INLINE fromDistinctDescList #-} -- INLINE for fusion
1717+
1718+
{--------------------------------------------------------------------
1719+
MapBuilder
1720+
--------------------------------------------------------------------}
1721+
1722+
-- Insert a key and value. The new value is combined with the old value if one
1723+
-- already exists for the key. Strict in the inserted value.
1724+
insertWithB
1725+
:: Ord k => (a -> a -> a) -> k -> a -> MapBuilder k a -> MapBuilder k a
1726+
insertWithB f !ky y b = case b of
1727+
BAsc stk -> case stk of
1728+
Push kx x l stk' -> case compare ky kx of
1729+
LT -> BMap (insertWith f ky y (ascLinkAll stk))
1730+
EQ -> BAsc (push' ky (f y x) l stk')
1731+
GT -> case l of
1732+
Tip -> y `seq` BAsc (ascLinkTop stk' 1 (singleton kx x) ky y)
1733+
Bin{} -> BAsc (push' ky y Tip stk)
1734+
Nada -> BAsc (push' ky y Tip Nada)
1735+
BMap m -> BMap (insertWith f ky y m)
1736+
where
1737+
push' kx !x = Push kx x
1738+
{-# INLINE insertWithB #-}

0 commit comments

Comments
 (0)