diff --git a/containers/src/Data/Map/Internal.hs b/containers/src/Data/Map/Internal.hs index 9461894e4..63766f8f1 100644 --- a/containers/src/Data/Map/Internal.hs +++ b/containers/src/Data/Map/Internal.hs @@ -1532,7 +1532,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 @@ -1552,7 +1552,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 @@ -1574,9 +1574,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 @@ -3001,7 +3001,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. @@ -3019,7 +3019,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, @@ -3042,8 +3042,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 @@ -3809,7 +3809,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. @@ -3842,7 +3842,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 @@ -3906,8 +3906,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 #-} @@ -3931,10 +3931,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__ @@ -3955,10 +3955,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__ @@ -4046,11 +4046,40 @@ 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 + +-- Called 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 +{-# INLINE linkL #-} + +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 + +-- Called 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 +{-# INLINE linkR #-} + +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 @@ -4072,10 +4101,28 @@ 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 + +-- Called when the left tree may be too large for the right tree, but not the +-- other way around. +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 + +-- Called when the right tree may be too large for the left tree, but not the +-- other way around. +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. diff --git a/containers/src/Data/Set/Internal.hs b/containers/src/Data/Set/Internal.hs index 3a8162ae4..9ce417987 100644 --- a/containers/src/Data/Set/Internal.hs +++ b/containers/src/Data/Set/Internal.hs @@ -1225,7 +1225,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. @@ -1253,7 +1253,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 @@ -1410,8 +1410,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 #-} @@ -1422,10 +1422,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__ @@ -1546,7 +1546,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 @@ -1566,7 +1566,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 @@ -1586,9 +1586,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 @@ -1606,7 +1606,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. @@ -1624,7 +1624,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, @@ -1647,8 +1647,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 @@ -1728,11 +1728,40 @@ 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 + +-- Called 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 +{-# INLINE linkL #-} + +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 + +-- Called 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 +{-# INLINE linkR #-} + +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 @@ -1754,10 +1783,28 @@ 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 + +-- Called when the left tree may be too large for the right tree, but not the +-- other way around. +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 + +-- Called when the right tree may be too large for the left tree, but not the +-- other way around. +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.