Build Set and Map more efficiently and consistently

containers-tests/benchmarks/Map.hs

@@ -21,7 +21,7 @@ main = do
         m_even = M.fromAscList elems_even :: M.Map Int Int
         m_odd = M.fromAscList elems_odd :: M.Map Int Int
     evaluate $ rnf [m, m_even, m_odd]
-    evaluate $ rnf elems_rev
+    evaluate $ rnf [elems_rev, elemsAsc, elemsDesc]
     defaultMain
         [ bench "lookup absent" $ whnf (lookup evens) m_odd
         , bench "lookup present" $ whnf (lookup evens) m_even
@@ -71,7 +71,7 @@ main = do
         , bench "insertLookupWithKey' present" $ whnf (insLookupWithKey' elems_even) m_even
         , bench "mapWithKey" $ whnf (M.mapWithKey (+)) m
         , bench "foldlWithKey" $ whnf (ins elems) m
-        , bench "foldlWithKey'" $ whnf (M.foldlWithKey' sum 0) m
+        , bench "foldlWithKey'" $ whnf (M.foldlWithKey' sumkv 0) m
         , bench "foldrWithKey" $ whnf (M.foldrWithKey consPair []) m
         , bench "foldrWithKey'" $ whnf (M.foldrWithKey' consPair []) m
         , bench "update absent" $ whnf (upd Just evens) m_odd
@@ -87,16 +87,33 @@ main = do
         , bench "difference" $ whnf (M.difference m) m_even
         , bench "intersection" $ whnf (M.intersection m) m_even
         , bench "split" $ whnf (M.split (bound `div` 2)) m
-        , bench "fromList" $ whnf M.fromList elems
-        , bench "fromList-desc" $ whnf M.fromList (reverse elems)
-        , bench "fromAscList" $ whnf M.fromAscList elems
+        , bench "fromList-asc" $ whnf M.fromList elems
+        , bench "fromList-asc:fusion" $ whnf (\n -> M.fromList [(i,i) | i <- [1..n]]) bound
+        , bench "fromList-desc" $ whnf M.fromList elems_rev
+        , bench "fromList-desc:fusion" $ whnf (\n -> M.fromList [(i,i) | i <- [n,n-1..1]]) bound
+        , bench "fromListWith-asc" $ whnf (M.fromListWith (+)) elemsAsc
+        , bench "fromListWith-asc:fusion" $
+            whnf (\n -> M.fromListWith (+) [(i `div` 2, i) | i <- [1..n]]) bound
+        , bench "fromListWith-desc" $ whnf (M.fromListWith (+)) elems_rev
+        , bench "fromListWith-desc:fusion" $
+            whnf (\n -> M.fromListWith (+) [(i `div` 2, i) | i <- [n,n-1..1]]) bound
+        , bench "fromListWithKey-asc" $ whnf (M.fromListWithKey sumkv) elemsAsc
+        , bench "fromListWithKey-asc:fusion" $
+            whnf (\n -> M.fromListWithKey sumkv [(i `div` 2, i) | i <- [1..n]]) bound
+        , bench "fromListWithKey-desc" $ whnf (M.fromListWithKey sumkv) elems_rev
+        , bench "fromListWithKey-desc:fusion" $
+            whnf (\n -> M.fromListWithKey sumkv [(i `div` 2, i) | i <- [n,n-1..1]]) bound
         , bench "fromDistinctAscList" $ whnf M.fromDistinctAscList elems
         , bench "fromDistinctAscList:fusion" $ whnf (\n -> M.fromDistinctAscList [(i,i) | i <- [1..n]]) bound
         , bench "fromDistinctDescList" $ whnf M.fromDistinctDescList elems_rev
         , bench "fromDistinctDescList:fusion" $ whnf (\n -> M.fromDistinctDescList [(i,i) | i <- [n,n-1..1]]) bound
         , bench "minView" $ whnf (\m' -> case M.minViewWithKey m' of {Nothing -> 0; Just ((k,v),m'') -> k+v+M.size m''}) (M.fromAscList $ zip [1..10::Int] [100..110::Int])
         , bench "eq" $ whnf (\m' -> m' == m') m -- worst case, compares everything
         , bench "compare" $ whnf (\m' -> compare m' m') m -- worst case, compares everything
+        , bench "mapKeys:asc" $ whnf (M.mapKeys (+1)) m
+        , bench "mapKeys:desc" $ whnf (M.mapKeys (negate . (+1))) m
+        , bench "mapKeysWith:asc" $ whnf (M.mapKeysWith (+) (`div` 2)) m
+        , bench "mapKeysWith:desc" $ whnf (M.mapKeysWith (+) (negate . (`div` 2))) m
         ]
   where
     bound = 2^12
@@ -108,7 +125,11 @@ main = do
     evens = [2,4..bound]
     odds = [1,3..bound]
     values = [1..bound]
-    sum k v1 v2 = k + v1 + v2
+    keysAsc = map (`div` 2) keys -- [0,1,1,2,2..]
+    elemsAsc = zip keysAsc values
+    keysDesc = map (`div` 2) (reverse keys) -- [..2,2,1,1,0]
+    elemsDesc = zip keysDesc values
+    sumkv k v1 v2 = k + v1 + v2
     consPair k v xs = (k, v) : xs
 
 add3 :: Int -> Int -> Int -> Int

containers-tests/benchmarks/Set.hs

@@ -18,7 +18,8 @@ main = do
     defaultMain
         [ bench "member" $ whnf (member elems) s
         , bench "insert" $ whnf (ins elems) S.empty
-        , bench "map" $ whnf (S.map (+ 1)) s
+        , bench "map:asc" $ whnf (S.map (+ 1)) s
+        , bench "map:desc" $ whnf (S.map (negate . (+ 1))) s
         , bench "filter" $ whnf (S.filter ((== 0) . (`mod` 2))) s
         , bench "partition" $ whnf (S.partition ((== 0) . (`mod` 2))) s
         , bench "fold" $ whnf (S.fold (:) []) s
@@ -31,8 +32,10 @@ main = do
         , bench "union" $ whnf (S.union s_even) s_odd
         , bench "difference" $ whnf (S.difference s) s_even
         , bench "intersection" $ whnf (S.intersection s) s_even
-        , bench "fromList" $ whnf S.fromList elems
-        , bench "fromList-desc" $ whnf S.fromList (reverse elems)
+        , bench "fromList-asc" $ whnf S.fromList elems
+        , bench "fromList-asc:fusion" $ whnf (\n -> S.fromList [1..n]) bound
+        , bench "fromList-desc" $ whnf S.fromList elems_rev
+        , bench "fromList-desc:fusion" $ whnf (\n -> S.fromList [n,n-1..1]) bound
         , bench "fromAscList" $ whnf S.fromAscList elems
         , bench "fromDistinctAscList" $ whnf S.fromDistinctAscList elems
         , bench "fromDistinctAscList:fusion" $ whnf (\n -> S.fromDistinctAscList [1..n]) bound

containers-tests/tests/map-properties.hs

@@ -193,6 +193,8 @@ main = defaultMain $ testGroup "map-properties"
          , testProperty "toDescList"           prop_descList
          , testProperty "toAscList+toDescList" prop_ascDescList
          , testProperty "fromList"             prop_fromList
+         , testProperty "fromListWith"         prop_fromListWith
+         , testProperty "fromListWithKey"      prop_fromListWithKey
          , testProperty "alter"                prop_alter
          , testProperty "alterF/alter"         prop_alterF_alter
          , testProperty "alterF/alter/noRULES" prop_alterF_alter_noRULES
@@ -219,7 +221,8 @@ main = defaultMain $ testGroup "map-properties"
          , testProperty "partition"            prop_partition
          , testProperty "map"                  prop_map
          , testProperty "fmap"                 prop_fmap
-         , testProperty "mapkeys"              prop_mapkeys
+         , testProperty "mapKeys"              prop_mapKeys
+         , testProperty "mapKeysWith"          prop_mapKeysWith
          , testProperty "split"                prop_splitModel
          , testProperty "fold"                 prop_fold
          , testProperty "foldMap"              prop_foldMap
@@ -1316,6 +1319,16 @@ prop_fromDistinctAscList xs =
     t = fromDistinctAscList nub_sort_xs
     nub_sort_xs = List.map List.head $ List.groupBy ((==) `on` fst) $ List.sortBy (comparing fst) xs
 
+prop_fromListWith :: Fun (A, A) A -> [(Int, A)] -> Property
+prop_fromListWith f kxs =
+  fromListWith (apply2 f) kxs ===
+  List.foldl' (\m (kx, x) -> insertWith (apply2 f) kx x m) empty kxs
+
+prop_fromListWithKey :: Fun (Int, A, A) A -> [(Int, A)] -> Property
+prop_fromListWithKey f kxs =
+  fromListWithKey (apply3 f) kxs ===
+  List.foldl' (\m (kx, x) -> insertWithKey (apply3 f) kx x m) empty kxs
+
 ----------------------------------------------------------------
 
 prop_alter :: UMap -> Int -> Bool
@@ -1509,11 +1522,15 @@ prop_fmap f ys = length ys > 0 ==>
       m  = fromList xs
   in  fmap (apply f) m == fromList [ (a, (apply f) b) | (a,b) <- xs ]
 
-prop_mapkeys :: Fun Int Int -> [(Int, Int)] -> Property
-prop_mapkeys f ys = length ys > 0 ==>
-  let xs = List.nubBy ((==) `on` fst) ys
-      m  = fromList xs
-  in  mapKeys (apply f) m == (fromList $ List.nubBy ((==) `on` fst) $ reverse [ (apply f a, b) | (a,b) <- sort xs])
+prop_mapKeys :: Fun Int Int -> Map Int A -> Property
+prop_mapKeys f m =
+  mapKeys (apply f) m ===
+  fromList (fmap (\(kx,x) -> (apply f kx, x)) (toList m))
+
+prop_mapKeysWith :: Fun (A, A) A -> Fun Int Int -> Map Int A -> Property
+prop_mapKeysWith f g m =
+  mapKeysWith (apply2 f) (apply g) m ===
+  fromListWith (apply2 f) (fmap (\(kx,x) -> (apply g kx, x)) (toList m))
 
 prop_splitModel :: Int -> [(Int, Int)] -> Property
 prop_splitModel n ys = length ys > 0 ==>

containers/src/Data/Map/Internal.hs

@@ -366,6 +366,9 @@ module Data.Map.Internal (
     , Identity(..)
     , Stack(..)
     , foldl'Stack
+    , MapBuilder(..)
+    , emptyB
+    , finishB
 
     -- Used by Map.Merge.Lazy
     , mapWhenMissing
@@ -387,7 +390,6 @@ import Data.Semigroup (Semigroup((<>)))
 #endif
 import Control.Applicative (Const (..))
 import Control.DeepSeq (NFData(rnf))
-import Data.Bits (shiftL, shiftR)
 import qualified Data.Foldable as Foldable
 import Data.Bifoldable
 import Utils.Containers.Internal.Prelude hiding
@@ -3242,7 +3244,7 @@ mapAccumRWithKey f a (Bin sx kx x l r) =
 -- > mapKeys (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "c"
 
 mapKeys :: Ord k2 => (k1->k2) -> Map k1 a -> Map k2 a
-mapKeys f = fromList . foldrWithKey (\k x xs -> (f k, x) : xs) []
+mapKeys f = finishB . foldlWithKey' (\b kx x -> insertB (f kx) x b) emptyB
 #if __GLASGOW_HASKELL__
 {-# INLINABLE mapKeys #-}
 #endif
@@ -3261,7 +3263,8 @@ mapKeys f = fromList . foldrWithKey (\k x xs -> (f k, x) : xs) []
 -- Also see the performance note on 'fromListWith'.
 
 mapKeysWith :: Ord k2 => (a -> a -> a) -> (k1->k2) -> Map k1 a -> Map k2 a
-mapKeysWith c f = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) []
+mapKeysWith c f =
+  finishB . foldlWithKey' (\b kx x -> insertWithB c (f kx) x b) emptyB
 #if __GLASGOW_HASKELL__
 {-# INLINABLE mapKeysWith #-}
 #endif
@@ -3510,46 +3513,9 @@ instance (Ord k) => GHCExts.IsList (Map k v) where
 -- > fromList [(5,"a"), (3,"b"), (5, "c")] == fromList [(5,"c"), (3,"b")]
 -- > fromList [(5,"c"), (3,"b"), (5, "a")] == fromList [(5,"a"), (3,"b")]
 
--- For some reason, when 'singleton' is used in fromList or in
--- create, it is not inlined, so we inline it manually.
 fromList :: Ord k => [(k,a)] -> Map k a
-fromList [] = Tip
-fromList [(kx, x)] = Bin 1 kx x Tip Tip
-fromList ((kx0, x0) : xs0) | not_ordered kx0 xs0 = fromList' (Bin 1 kx0 x0 Tip Tip) xs0
-                           | otherwise = go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0
-  where
-    not_ordered _ [] = False
-    not_ordered kx ((ky,_) : _) = kx >= ky
-    {-# INLINE not_ordered #-}
-
-    fromList' t0 xs = Foldable.foldl' ins t0 xs
-      where ins t (k,x) = insert k x t
-
-    go !_ t [] = t
-    go _ t [(kx, x)] = insertMax kx x t
-    go s l xs@((kx, x) : xss) | not_ordered kx xss = fromList' l xs
-                              | otherwise = case create s xss of
-                                  (r, ys, []) -> go (s `shiftL` 1) (link kx x l r) ys
-                                  (r, _,  ys) -> fromList' (link kx x l r) ys
-
-    -- The create is returning a triple (tree, xs, ys). Both xs and ys
-    -- represent not yet processed elements and only one of them can be nonempty.
-    -- If ys is nonempty, the keys in ys are not ordered with respect to tree
-    -- and must be inserted using fromList'. Otherwise the keys have been
-    -- ordered so far.
-    create !_ [] = (Tip, [], [])
-    create s xs@(xp : xss)
-      | s == 1 = case xp of (kx, x) | not_ordered kx xss -> (Bin 1 kx x Tip Tip, [], xss)
-                                    | otherwise -> (Bin 1 kx x Tip Tip, xss, [])
-      | otherwise = case create (s `shiftR` 1) xs of
-                      res@(_, [], _) -> res
-                      (l, [(ky, y)], zs) -> (insertMax ky y l, [], zs)
-                      (l, ys@((ky, y):yss), _) | not_ordered ky yss -> (l, [], ys)
-                                               | otherwise -> case create (s `shiftR` 1) yss of
-                                                   (r, zs, ws) -> (link ky y l r, zs, ws)
-#if __GLASGOW_HASKELL__
-{-# INLINABLE fromList #-}
-#endif
+fromList xs = finishB (Foldable.foldl' (\b (kx, x) -> insertB kx x b) emptyB xs)
+{-# INLINE fromList #-} -- INLINE for fusion
 
 -- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
 --
@@ -3588,11 +3554,9 @@ fromList ((kx0, x0) : xs0) | not_ordered kx0 xs0 = fromList' (Bin 1 kx0 x0 Tip T
 -- > fromListWith (++) $ reverse $ map (\(k, v) -> (k, [v])) someListOfTuples
 
 fromListWith :: Ord k => (a -> a -> a) -> [(k,a)] -> Map k a
-fromListWith f xs
-  = fromListWithKey (\_ x y -> f x y) xs
-#if __GLASGOW_HASKELL__
-{-# INLINABLE fromListWith #-}
-#endif
+fromListWith f xs =
+  finishB (Foldable.foldl' (\b (kx, x) -> insertWithB f kx x b) emptyB xs)
+{-# INLINE fromListWith #-}  -- INLINE for fusion
 
 -- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'.
 --
@@ -3603,13 +3567,9 @@ fromListWith f xs
 -- Also see the performance note on 'fromListWith'.
 
 fromListWithKey :: Ord k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
-fromListWithKey f xs
-  = Foldable.foldl' ins empty xs
-  where
-    ins t (k,x) = insertWithKey f k x t
-#if __GLASGOW_HASKELL__
-{-# INLINABLE fromListWithKey #-}
-#endif
+fromListWithKey f xs =
+  finishB (Foldable.foldl' (\b (kx, x) -> insertWithB (f kx) kx x b) emptyB xs)
+{-# INLINE fromListWithKey #-}  -- INLINE for fusion
 
 -- | \(O(n)\). Convert the map to a list of key\/value pairs. Subject to list fusion.
 --
@@ -4004,6 +3964,57 @@ splitMember k0 m = case go k0 m of
 
 data StrictTriple a b c = StrictTriple !a !b !c
 
+{--------------------------------------------------------------------
+  MapBuilder
+--------------------------------------------------------------------}
+
+-- See Note [SetBuilder] in Data.Set.Internal
+
+data MapBuilder k a
+  = BAsc !(Stack k a)
+  | BMap !(Map k a)
+
+-- Empty builder.
+emptyB :: MapBuilder k a
+emptyB = BAsc Nada
+
+-- Insert a key and value. Replaces the old value if one already exists for
+-- the key.
+insertB :: Ord k => k -> a -> MapBuilder k a -> MapBuilder k a
+insertB !ky y b = case b of
+  BAsc stk -> case stk of
+    Push kx x l stk' -> case compare ky kx of
+      LT -> BMap (insert ky y (ascLinkAll stk))
+      EQ -> BAsc (Push ky y l stk')
+      GT -> case l of
+        Tip -> BAsc (ascLinkTop stk' 1 (singleton kx x) ky y)
+        Bin{} -> BAsc (Push ky y Tip stk)
+    Nada -> BAsc (Push ky y Tip Nada)
+  BMap m -> BMap (insert ky y m)
+{-# INLINE insertB #-}
+
+-- Insert a key and value. The new value is combined with the old value if one
+-- already exists for the key.
+insertWithB
+  :: Ord k => (a -> a -> a) -> k -> a -> MapBuilder k a -> MapBuilder k a
+insertWithB f !ky y b = case b of
+  BAsc stk -> case stk of
+    Push kx x l stk' -> case compare ky kx of
+      LT -> BMap (insertWith f ky y (ascLinkAll stk))
+      EQ -> BAsc (Push ky (f y x) l stk')
+      GT -> case l of
+        Tip -> BAsc (ascLinkTop stk' 1 (singleton kx x) ky y)
+        Bin{} -> BAsc (Push ky y Tip stk)
+    Nada -> BAsc (Push ky y Tip Nada)
+  BMap m -> BMap (insertWith f ky y m)
+{-# INLINE insertWithB #-}
+
+-- Finalize the builder into a Map.
+finishB :: MapBuilder k a -> Map k a
+finishB (BAsc stk) = ascLinkAll stk
+finishB (BMap m) = m
+{-# INLINABLE finishB #-}
+
 {--------------------------------------------------------------------
   Utility functions that maintain the balance properties of the tree.
   All constructors assume that all values in [l] < [k] and all values