Add link{L,R} for Set and Map (#1141)

linkL restores balance when the left tree is too heavy for the right
tree, but unlike link it skips checking the other way around. linkR
handles the opposite. linkL and linkR are used where we know the
direction in which we need to rebalance.

linkL and linkR delegate to linkL_ and linkR_, which do the actual work.
link also delegates to linkL_ and linkR_, since the direction does not
change when linking two trees. This single-direction linking algorithm
can be seen in Figure 4 of "Parallel Ordered Sets Using Join". Similarly,
link2/merge delegates to link2L_,link2_R/mergeL_,mergeR_.

Benchmarks show a decrease in time and allocations, especially in set-set
operations like union. Inspection of the Core of the previous link
implementation revealed that it reboxes trees in some cases, which
explains the decrease in allocations.

Author: meooow25

containers/src/Data/Map/Internal.hs

@@ -1558,7 +1558,7 @@ take i0 m0 = go i0 m0
     go i (Bin _ kx x l r) =
       case compare i sizeL of
         LT -> go i l
-        GT -> link kx x l (go (i - sizeL - 1) r)
+        GT -> linkL kx x l (go (i - sizeL - 1) r)
         EQ -> l
       where sizeL = size l
 
@@ -1578,7 +1578,7 @@ drop i0 m0 = go i0 m0
     go !_ Tip = Tip
     go i (Bin _ kx x l r) =
       case compare i sizeL of
-        LT -> link kx x (go i l) r
+        LT -> linkR kx x (go i l) r
         GT -> go (i - sizeL - 1) r
         EQ -> insertMin kx x r
       where sizeL = size l
@@ -1600,9 +1600,9 @@ splitAt i0 m0
     go i (Bin _ kx x l r)
       = case compare i sizeL of
           LT -> case go i l of
-                  ll :*: lr -> ll :*: link kx x lr r
+                  ll :*: lr -> ll :*: linkR kx x lr r
           GT -> case go (i - sizeL - 1) r of
-                  rl :*: rr -> link kx x l rl :*: rr
+                  rl :*: rr -> linkL kx x l rl :*: rr
           EQ -> l :*: insertMin kx x r
       where sizeL = size l
 
@@ -3034,7 +3034,7 @@ filterWithKeyA p t@(Bin _ kx x l r) =
 takeWhileAntitone :: (k -> Bool) -> Map k a -> Map k a
 takeWhileAntitone _ Tip = Tip
 takeWhileAntitone p (Bin _ kx x l r)
-  | p kx = link kx x l (takeWhileAntitone p r)
+  | p kx = linkL kx x l (takeWhileAntitone p r)
   | otherwise = takeWhileAntitone p l
 
 -- | \(O(\log n)\). Drop while a predicate on the keys holds.
@@ -3052,7 +3052,7 @@ dropWhileAntitone :: (k -> Bool) -> Map k a -> Map k a
 dropWhileAntitone _ Tip = Tip
 dropWhileAntitone p (Bin _ kx x l r)
   | p kx = dropWhileAntitone p r
-  | otherwise = link kx x (dropWhileAntitone p l) r
+  | otherwise = linkR kx x (dropWhileAntitone p l) r
 
 -- | \(O(\log n)\). Divide a map at the point where a predicate on the keys stops holding.
 -- The user is responsible for ensuring that for all keys @j@ and @k@ in the map,
@@ -3075,8 +3075,8 @@ spanAntitone p0 m = toPair (go p0 m)
   where
     go _ Tip = Tip :*: Tip
     go p (Bin _ kx x l r)
-      | p kx = let u :*: v = go p r in link kx x l u :*: v
-      | otherwise = let u :*: v = go p l in u :*: link kx x v r
+      | p kx = let u :*: v = go p r in linkL kx x l u :*: v
+      | otherwise = let u :*: v = go p l in u :*: linkR kx x v r
 
 -- | \(O(n)\). Partition the map according to a predicate. The first
 -- map contains all elements that satisfy the predicate, the second all
@@ -3842,7 +3842,7 @@ ascLinkTop (Push kx x l@(Bin lsz _ _ _ _) stk) !rsz r ky y
 ascLinkTop stk !_ l kx x = Push kx x l stk
 
 ascLinkAll :: Stack k a -> Map k a
-ascLinkAll stk = foldl'Stack (\r kx x l -> link kx x l r) Tip stk
+ascLinkAll stk = foldl'Stack (\r kx x l -> linkL kx x l r) Tip stk
 {-# INLINABLE ascLinkAll #-}
 
 -- | \(O(n)\). Build a map from a descending list of distinct elements in linear time.
@@ -3875,7 +3875,7 @@ descLinkTop ky y !_ r stk = Push ky y r stk
 {-# INLINABLE descLinkTop #-}
 
 descLinkAll :: Stack k a -> Map k a
-descLinkAll stk = foldl'Stack (\l kx x r -> link kx x l r) Tip stk
+descLinkAll stk = foldl'Stack (\l kx x r -> linkR kx x l r) Tip stk
 {-# INLINABLE descLinkAll #-}
 
 data Stack k a = Push !k a !(Map k a) !(Stack k a) | Nada
@@ -3939,8 +3939,8 @@ split !k0 t0 = toPair $ go k0 t0
       case t of
         Tip            -> Tip :*: Tip
         Bin _ kx x l r -> case compare k kx of
-          LT -> let (lt :*: gt) = go k l in lt :*: link kx x gt r
-          GT -> let (lt :*: gt) = go k r in link kx x l lt :*: gt
+          LT -> let (lt :*: gt) = go k l in lt :*: linkR kx x gt r
+          GT -> let (lt :*: gt) = go k r in linkL kx x l lt :*: gt
           EQ -> (l :*: r)
 #if __GLASGOW_HASKELL__
 {-# INLINABLE split #-}
@@ -3964,10 +3964,10 @@ splitLookup k0 m = case go k0 m of
         Tip            -> StrictTriple Tip Nothing Tip
         Bin _ kx x l r -> case compare k kx of
           LT -> let StrictTriple lt z gt = go k l
-                    !gt' = link kx x gt r
+                    !gt' = linkR kx x gt r
                 in StrictTriple lt z gt'
           GT -> let StrictTriple lt z gt = go k r
-                    !lt' = link kx x l lt
+                    !lt' = linkL kx x l lt
                 in StrictTriple lt' z gt
           EQ -> StrictTriple l (Just x) r
 #if __GLASGOW_HASKELL__
@@ -3988,10 +3988,10 @@ splitMember k0 m = case go k0 m of
         Tip            -> StrictTriple Tip False Tip
         Bin _ kx x l r -> case compare k kx of
           LT -> let StrictTriple lt z gt = go k l
-                    !gt' = link kx x gt r
+                    !gt' = linkR kx x gt r
                 in StrictTriple lt z gt'
           GT -> let StrictTriple lt z gt = go k r
-                    !lt' = link kx x l lt
+                    !lt' = linkL kx x l lt
                 in StrictTriple lt' z gt
           EQ -> StrictTriple l True r
 #if __GLASGOW_HASKELL__
@@ -4079,11 +4079,38 @@ finishB (BMap m) = m
 link :: k -> a -> Map k a -> Map k a -> Map k a
 link kx x Tip r  = insertMin kx x r
 link kx x l Tip  = insertMax kx x l
-link kx x l@(Bin sizeL ky y ly ry) r@(Bin sizeR kz z lz rz)
-  | delta*sizeL < sizeR  = balanceL kz z (link kx x l lz) rz
-  | delta*sizeR < sizeL  = balanceR ky y ly (link kx x ry r)
-  | otherwise            = bin kx x l r
-
+link kx x l@(Bin lsz lkx lx ll lr) r@(Bin rsz rkx rx rl rr)
+  | delta*lsz < rsz = balanceL rkx rx (linkR_ kx x lsz l rl) rr
+  | delta*rsz < lsz = balanceR lkx lx ll (linkL_ kx x lr rsz r)
+  | otherwise       = Bin (1+lsz+rsz) kx x l r
+
+-- Variant of link. Restores balance when the left tree may be too large for the
+-- right tree, but not the other way around.
+linkL :: k -> a -> Map k a -> Map k a -> Map k a
+linkL kx x l r = case r of
+  Tip -> insertMax kx x l
+  Bin rsz _ _ _ _ -> linkL_ kx x l rsz r
+
+linkL_ :: k -> a -> Map k a -> Int -> Map k a -> Map k a
+linkL_ kx x l !rsz r = case l of
+  Bin lsz lkx lx ll lr
+    | delta*rsz < lsz -> balanceR lkx lx ll (linkL_ kx x lr rsz r)
+    | otherwise -> Bin (1+lsz+rsz) kx x l r
+  Tip -> Bin (1+rsz) kx x Tip r
+
+-- Variant of link. Restores balance when the right tree may be too large for
+-- the left tree, but not the other way around.
+linkR :: k -> a -> Map k a -> Map k a -> Map k a
+linkR kx x l r = case l of
+  Tip -> insertMin kx x r
+  Bin lsz _ _ _ _ -> linkR_ kx x lsz l r
+
+linkR_ :: k -> a -> Int -> Map k a -> Map k a -> Map k a
+linkR_ kx x !lsz l r = case r of
+  Bin rsz rkx rx rl rr
+    | delta*lsz < rsz -> balanceL rkx rx (linkR_ kx x lsz l rl) rr
+    | otherwise -> Bin (1+lsz+rsz) kx x l r
+  Tip -> Bin (1+lsz) kx x l Tip
 
 -- insertMin and insertMax don't perform potentially expensive comparisons.
 insertMax,insertMin :: k -> a -> Map k a -> Map k a
@@ -4105,10 +4132,24 @@ insertMin kx x t
 link2 :: Map k a -> Map k a -> Map k a
 link2 Tip r   = r
 link2 l Tip   = l
-link2 l@(Bin sizeL kx x lx rx) r@(Bin sizeR ky y ly ry)
-  | delta*sizeL < sizeR = balanceL ky y (link2 l ly) ry
-  | delta*sizeR < sizeL = balanceR kx x lx (link2 rx r)
-  | otherwise           = glue l r
+link2 l@(Bin lsz lkx lx ll lr) r@(Bin rsz rkx rx rl rr)
+  | delta*lsz < rsz = balanceL rkx rx (link2R_ lsz l rl) rr
+  | delta*rsz < lsz = balanceR lkx lx ll (link2L_ lr rsz r)
+  | otherwise = glue l r
+
+link2L_ :: Map k a -> Int -> Map k a -> Map k a
+link2L_ l !rsz r = case l of
+  Bin lsz lkx lx ll lr
+    | delta*rsz < lsz -> balanceR lkx lx ll (link2L_ lr rsz r)
+    | otherwise -> glue l r
+  Tip -> r
+
+link2R_ :: Int -> Map k a -> Map k a -> Map k a
+link2R_ !lsz l r = case r of
+  Bin rsz rkx rx rl rr
+    | delta*lsz < rsz -> balanceL rkx rx (link2R_ lsz l rl) rr
+    | otherwise -> glue l r
+  Tip -> l
 
 {--------------------------------------------------------------------
   [glue l r]: glues two trees together.

containers/src/Data/Set/Internal.hs

@@ -1231,7 +1231,7 @@ ascLinkTop (Push x l@(Bin lsz _ _ _) stk) !rsz r y
 ascLinkTop stk !_ r y = Push y r stk
 
 ascLinkAll :: Stack a -> Set a
-ascLinkAll stk = foldl'Stack (\r x l -> link x l r) Tip stk
+ascLinkAll stk = foldl'Stack (\r x l -> linkL x l r) Tip stk
 {-# INLINABLE ascLinkAll #-}
 
 -- | \(O(n)\). Build a set from a descending list of distinct elements in linear time.
@@ -1259,7 +1259,7 @@ descLinkTop x !lsz l (Push y r@(Bin rsz _ _ _) stk)
 descLinkTop y !_ r stk = Push y r stk
 
 descLinkAll :: Stack a -> Set a
-descLinkAll stk = foldl'Stack (\l x r -> link x l r) Tip stk
+descLinkAll stk = foldl'Stack (\l x r -> linkR x l r) Tip stk
 {-# INLINABLE descLinkAll #-}
 
 data Stack a = Push !a !(Set a) !(Stack a) | Nada
@@ -1416,8 +1416,8 @@ splitS :: Ord a => a -> Set a -> StrictPair (Set a) (Set a)
 splitS _ Tip = (Tip :*: Tip)
 splitS x (Bin _ y l r)
       = case compare x y of
-          LT -> let (lt :*: gt) = splitS x l in (lt :*: link y gt r)
-          GT -> let (lt :*: gt) = splitS x r in (link y l lt :*: gt)
+          LT -> let (lt :*: gt) = splitS x l in (lt :*: linkR y gt r)
+          GT -> let (lt :*: gt) = splitS x r in (linkL y l lt :*: gt)
           EQ -> (l :*: r)
 {-# INLINABLE splitS #-}
 
@@ -1428,10 +1428,10 @@ splitMember _ Tip = (Tip, False, Tip)
 splitMember x (Bin _ y l r)
    = case compare x y of
        LT -> let (lt, found, gt) = splitMember x l
-                 !gt' = link y gt r
+                 !gt' = linkR y gt r
              in (lt, found, gt')
        GT -> let (lt, found, gt) = splitMember x r
-                 !lt' = link y l lt
+                 !lt' = linkL y l lt
              in (lt', found, gt)
        EQ -> (l, True, r)
 #if __GLASGOW_HASKELL__
@@ -1558,7 +1558,7 @@ take i0 m0 = go i0 m0
     go i (Bin _ x l r) =
       case compare i sizeL of
         LT -> go i l
-        GT -> link x l (go (i - sizeL - 1) r)
+        GT -> linkL x l (go (i - sizeL - 1) r)
         EQ -> l
       where sizeL = size l
 
@@ -1578,7 +1578,7 @@ drop i0 m0 = go i0 m0
     go !_ Tip = Tip
     go i (Bin _ x l r) =
       case compare i sizeL of
-        LT -> link x (go i l) r
+        LT -> linkR x (go i l) r
         GT -> go (i - sizeL - 1) r
         EQ -> insertMin x r
       where sizeL = size l
@@ -1598,9 +1598,9 @@ splitAt i0 m0
     go i (Bin _ x l r)
       = case compare i sizeL of
           LT -> case go i l of
-                  ll :*: lr -> ll :*: link x lr r
+                  ll :*: lr -> ll :*: linkR x lr r
           GT -> case go (i - sizeL - 1) r of
-                  rl :*: rr -> link x l rl :*: rr
+                  rl :*: rr -> linkL x l rl :*: rr
           EQ -> l :*: insertMin x r
       where sizeL = size l
 
@@ -1618,7 +1618,7 @@ splitAt i0 m0
 takeWhileAntitone :: (a -> Bool) -> Set a -> Set a
 takeWhileAntitone _ Tip = Tip
 takeWhileAntitone p (Bin _ x l r)
-  | p x = link x l (takeWhileAntitone p r)
+  | p x = linkL x l (takeWhileAntitone p r)
   | otherwise = takeWhileAntitone p l
 
 -- | \(O(\log n)\). Drop while a predicate on the elements holds.
@@ -1636,7 +1636,7 @@ dropWhileAntitone :: (a -> Bool) -> Set a -> Set a
 dropWhileAntitone _ Tip = Tip
 dropWhileAntitone p (Bin _ x l r)
   | p x = dropWhileAntitone p r
-  | otherwise = link x (dropWhileAntitone p l) r
+  | otherwise = linkR x (dropWhileAntitone p l) r
 
 -- | \(O(\log n)\). Divide a set at the point where a predicate on the elements stops holding.
 -- The user is responsible for ensuring that for all elements @j@ and @k@ in the set,
@@ -1659,8 +1659,8 @@ spanAntitone p0 m = toPair (go p0 m)
   where
     go _ Tip = Tip :*: Tip
     go p (Bin _ x l r)
-      | p x = let u :*: v = go p r in link x l u :*: v
-      | otherwise = let u :*: v = go p l in u :*: link x v r
+      | p x = let u :*: v = go p r in linkL x l u :*: v
+      | otherwise = let u :*: v = go p l in u :*: linkR x v r
 
 {--------------------------------------------------------------------
   SetBuilder
@@ -1740,11 +1740,38 @@ finishB (BSet s) = s
 link :: a -> Set a -> Set a -> Set a
 link x Tip r  = insertMin x r
 link x l Tip  = insertMax x l
-link x l@(Bin sizeL y ly ry) r@(Bin sizeR z lz rz)
-  | delta*sizeL < sizeR  = balanceL z (link x l lz) rz
-  | delta*sizeR < sizeL  = balanceR y ly (link x ry r)
-  | otherwise            = bin x l r
-
+link x l@(Bin lsz lx ll lr) r@(Bin rsz rx rl rr)
+  | delta*lsz < rsz = balanceL rx (linkR_ x lsz l rl) rr
+  | delta*rsz < lsz = balanceR lx ll (linkL_ x lr rsz r)
+  | otherwise = Bin (1+lsz+rsz) x l r
+
+-- Variant of link. Restores balance when the left tree may be too large for the
+-- right tree, but not the other way around.
+linkL :: a -> Set a -> Set a -> Set a
+linkL x l r = case r of
+  Tip -> insertMax x l
+  Bin rsz _ _ _ -> linkL_ x l rsz r
+
+linkL_ :: a -> Set a -> Int -> Set a -> Set a
+linkL_ x l !rsz r = case l of
+  Bin lsz lx ll lr
+    | delta*rsz < lsz -> balanceR lx ll (linkL_ x lr rsz r)
+    | otherwise -> Bin (1+lsz+rsz) x l r
+  Tip -> Bin (1+rsz) x Tip r
+
+-- Variant of link. Restores balance when the right tree may be too large for
+-- the left tree, but not the other way around.
+linkR :: a -> Set a -> Set a -> Set a
+linkR x l r = case l of
+  Tip -> insertMin x r
+  Bin lsz _ _ _ -> linkR_ x lsz l r
+
+linkR_ :: a -> Int -> Set a -> Set a -> Set a
+linkR_ x !lsz l r = case r of
+  Bin rsz rx rl rr
+    | delta*lsz < rsz -> balanceL rx (linkR_ x lsz l rl) rr
+    | otherwise -> Bin (1+lsz+rsz) x l r
+  Tip -> Bin (1+lsz) x l Tip
 
 -- insertMin and insertMax don't perform potentially expensive comparisons.
 insertMax,insertMin :: a -> Set a -> Set a
@@ -1766,10 +1793,24 @@ insertMin x t
 merge :: Set a -> Set a -> Set a
 merge Tip r   = r
 merge l Tip   = l
-merge l@(Bin sizeL x lx rx) r@(Bin sizeR y ly ry)
-  | delta*sizeL < sizeR = balanceL y (merge l ly) ry
-  | delta*sizeR < sizeL = balanceR x lx (merge rx r)
-  | otherwise           = glue l r
+merge l@(Bin lsz lx ll lr) r@(Bin rsz rx rl rr)
+  | delta*lsz < rsz = balanceL rx (mergeR_ lsz l rl) rr
+  | delta*rsz < lsz = balanceR lx ll (mergeL_ lr rsz r)
+  | otherwise = glue l r
+
+mergeL_ :: Set a -> Int -> Set a -> Set a
+mergeL_ l !rsz r = case l of
+  Bin lsz lx ll lr
+    | delta*rsz < lsz -> balanceR lx ll (mergeL_ lr rsz r)
+    | otherwise -> glue l r
+  Tip -> r
+
+mergeR_ :: Int -> Set a -> Set a -> Set a
+mergeR_ !lsz l r = case r of
+  Bin rsz rx rl rr
+    | delta*lsz < rsz -> balanceL rx (mergeR_ lsz l rl) rr
+    | otherwise -> glue l r
+  Tip -> l
 
 {--------------------------------------------------------------------
   [glue l r]: glues two trees together.