Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 72 additions & 25 deletions containers/src/Data/Map/Internal.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand All @@ -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

Expand Down Expand Up @@ -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.
Expand All @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 #-}
Expand All @@ -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__
Expand All @@ -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__
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand Down
93 changes: 70 additions & 23 deletions containers/src/Data/Set/Internal.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 #-}

Expand All @@ -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__
Expand Down Expand Up @@ -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

Expand All @@ -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
Expand All @@ -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

Expand All @@ -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.
Expand All @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand Down