[commit: containers] master: Improve {Map, Set}.union. (cd96a90)

[Paolo Capriotti]

Repository : ssh://darcs.haskell.org//srv/darcs/packages/containers

On branch  : master

http://hackage.haskell.org/trac/ghc/changeset/cd96a9045df8d2f8d98d9ec6e13c4a2492f0836b

>---------------------------------------------------------------

commit cd96a9045df8d2f8d98d9ec6e13c4a2492f0836b
Author: Milan Straka <fox at ucw.cz>
Date:   Fri Apr 27 11:51:48 2012 +0200

    Improve {Map, Set}.union.
    
    Instead of having a special case set `union` set_of_size_1 in union,
    move it to hedgeUnion, so it can be used recursively. Benchmark shows up
    to 30% speedup.

>---------------------------------------------------------------

 Data/Map/Base.hs |   26 ++++++++++++++++++++++++--
 Data/Set/Base.hs |    9 +++++----
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/Data/Map/Base.hs b/Data/Map/Base.hs
index ae22b51..744ddad 100644
--- a/Data/Map/Base.hs
+++ b/Data/Map/Base.hs
@@ -672,6 +672,27 @@ insert = go
 {-# INLINE insert #-}
 #endif
 
+-- Insert a new key and value in the map if it is not already present.
+-- Used by `union`.
+
+-- See Note: Type of local 'go' function
+insertR :: Ord k => k -> a -> Map k a -> Map k a
+insertR = go
+  where
+    go :: Ord k => k -> a -> Map k a -> Map k a
+    STRICT_1_OF_3(go)
+    go kx x Tip = singleton kx x
+    go kx x t@(Bin sz ky y l r) =
+        case compare kx ky of
+            LT -> balanceL ky y (go kx x l) r
+            GT -> balanceR ky y l (go kx x r)
+            EQ -> t
+#if __GLASGOW_HASKELL__ >= 700
+{-# INLINABLE insertR #-}
+#else
+{-# INLINE insertR #-}
+#endif
+
 -- | /O(log n)/. Insert with a function, combining new value and old value.
 -- @'insertWith' f key value mp@
 -- will insert the pair (key, value) into @mp@ if key does
@@ -1216,8 +1237,6 @@ unionsWith f ts
 union :: Ord k => Map k a -> Map k a -> Map k a
 union Tip t2  = t2
 union t1 Tip  = t1
-union (Bin _ k x Tip Tip) t = insert k x t
-union t (Bin _ k x Tip Tip) = insertWith (\_ y->y) k x t
 union t1 t2 = hedgeUnionL NothingS NothingS t1 t2
 #if __GLASGOW_HASKELL__ >= 700
 {-# INLINABLE union #-}
@@ -1231,6 +1250,9 @@ hedgeUnionL _     _     t1 Tip
   = t1
 hedgeUnionL blo bhi Tip (Bin _ kx x l r)
   = join kx x (filterGt blo l) (filterLt bhi r)
+hedgeUnionL blo bhi t1 (Bin _ kx x Tip Tip)
+  = insertR kx x t1  -- According to benchmarks, this special case increases
+                     -- performance up to 30%. It does not help in difference or intersection.
 hedgeUnionL blo bhi (Bin _ kx x l r) t2
   = join kx x (hedgeUnionL blo bmi l (trim blo bmi t2))
               (hedgeUnionL bmi bhi r (trim bmi bhi t2))
diff --git a/Data/Set/Base.hs b/Data/Set/Base.hs
index a8573c7..0345bda 100644
--- a/Data/Set/Base.hs
+++ b/Data/Set/Base.hs
@@ -438,8 +438,8 @@ insert = go
 {-# INLINE insert #-}
 #endif
 
--- Insert an element to the set only if it is not in the set. Used by
--- `union`.
+-- Insert an element to the set only if it is not in the set.
+-- Used by `union`.
 
 -- See Note: Type of local 'go' function
 insertR :: Ord a => a -> Set a -> Set a
@@ -554,8 +554,6 @@ unions = foldlStrict union empty
 union :: Ord a => Set a -> Set a -> Set a
 union Tip t2  = t2
 union t1 Tip  = t1
-union (Bin _ x Tip Tip) t = insert x t
-union t (Bin _ x Tip Tip) = insertR x t
 union t1 t2 = hedgeUnion NothingS NothingS t1 t2
 #if __GLASGOW_HASKELL__ >= 700
 {-# INLINABLE union #-}
@@ -567,6 +565,9 @@ hedgeUnion _     _     t1 Tip
   = t1
 hedgeUnion blo bhi Tip (Bin _ x l r)
   = join x (filterGt blo l) (filterLt bhi r)
+hedgeUnion blo bhi t1 (Bin _ x Tip Tip)
+  = insertR x t1     -- According to benchmarks, this special case increases
+                     -- performance up to 30%. It does not help in difference or intersection.
 hedgeUnion blo bhi (Bin _ x l r) t2
   = join x (hedgeUnion blo bmi l (trim blo bmi t2))
            (hedgeUnion bmi bhi r (trim bmi bhi t2))