-- |
-- Copyright   :  (c) Adrian Hey 2004,2005
-- License     :  BSD3
--
{-# OPTIONS_HADDOCK hide #-}
module Data.Tree.AVL.Join
(-- * Joining AVL trees
 join,concatAVL,flatConcat,
) where

import Prelude ()
import Data.Tree.AVL.Internals.Types(AVL(..))
import Data.Tree.AVL.Size(addSize)
import Data.Tree.AVL.List(asTreeLenL,toListL)
import Data.Tree.AVL.Internals.DelUtils(popHLN,popHLZ,popHLP)
import Data.Tree.AVL.Height(height,addHeight)
import Data.Tree.AVL.Internals.HJoin(joinH',spliceH)
import Data.Foldable (foldl', foldr)

import GHC.Exts
#include "ghcdefs.h"

-- | Join two AVL trees. This is the AVL equivalent of (++).
--
-- > asListL (l `join` r) = asListL l ++ asListL r
--
-- Complexity: O(log n), where n is the size of the larger of the two trees.
join :: AVL e -> AVL e -> AVL e
join :: forall e. AVL e -> AVL e -> AVL e
join AVL e
l AVL e
r = AVL e -> Int# -> AVL e -> Int# -> AVL e
forall e. AVL e -> Int# -> AVL e -> Int# -> AVL e
joinH' AVL e
l (AVL e -> Int#
forall e. AVL e -> Int#
height AVL e
l) AVL e
r (AVL e -> Int#
forall e. AVL e -> Int#
height AVL e
r)

-- Specialised list of AVL trees of known height, with leftmost element popped.
-- (used by concatAVL).
data HAVLS e = HE | H e (AVL e) UINT (HAVLS e)

-- | Concatenate a /finite/ list of AVL trees. During construction of the resulting tree the
-- input list is consumed lazily, but it will be consumed entirely before the result is returned.
--
-- > asListL (concatAVL avls) = concatMap asListL avls
--
-- Complexity: Umm..Dunno. Uses a divide and conquer approach to splice adjacent pairs of
-- trees in the list recursively, until only one tree remains. The complexity of each splice
-- is proportional to the difference in tree heights.
concatAVL :: [AVL e] -> AVL e
concatAVL :: forall e. [AVL e] -> AVL e
concatAVL []               = AVL e
forall e. AVL e
E
concatAVL (   AVL e
E       :[AVL e]
ts) = [AVL e] -> AVL e
forall e. [AVL e] -> AVL e
concatAVL [AVL e]
ts
concatAVL (t :: AVL e
t@(N AVL e
l e
_ AVL e
_):[AVL e]
ts) = AVL e -> Int# -> HAVLS e -> AVL e
forall e. AVL e -> Int# -> HAVLS e -> AVL e
concatHAVLS AVL e
t (Int# -> AVL e -> Int#
forall e. Int# -> AVL e -> Int#
addHeight L(2AVL e
) l) (mkHAVLS ts)
concatAVL (t :: AVL e
t@(Z AVL e
l e
_ AVL e
_):[AVL e]
ts) = AVL e -> Int# -> HAVLS e -> AVL e
forall e. AVL e -> Int# -> HAVLS e -> AVL e
concatHAVLS AVL e
t (Int# -> AVL e -> Int#
forall e. Int# -> AVL e -> Int#
addHeight L(1AVL e
) l) (mkHAVLS ts)
concatAVL (t :: AVL e
t@(P AVL e
_ e
_ AVL e
r):[AVL e]
ts) = AVL e -> Int# -> HAVLS e -> AVL e
forall e. AVL e -> Int# -> HAVLS e -> AVL e
concatHAVLS AVL e
t (Int# -> AVL e -> Int#
forall e. Int# -> AVL e -> Int#
addHeight L(2AVL e
) r) (mkHAVLS ts)

-- Recursively call mergePairs until only one tree remains.
-- The head of the current list has to be treated specially becuase it has no associated
-- bridging element.
concatHAVLS :: AVL e -> UINT -> HAVLS e -> AVL e
concatHAVLS :: forall e. AVL e -> Int# -> HAVLS e -> AVL e
concatHAVLS AVL e
l Int#
_   HAVLS e
HE               = AVL e
l
concatHAVLS AVL e
l Int#
hl (H e
e AVL e
r Int#
hr HAVLS e
hs) = case AVL e
-> Int#
-> e
-> AVL e
-> Int#
-> HAVLS e
-> (# AVL e, Int#, HAVLS e #)
forall e.
AVL e
-> Int#
-> e
-> AVL e
-> Int#
-> HAVLS e
-> (# AVL e, Int#, HAVLS e #)
mergePairs AVL e
l Int#
hl e
e AVL e
r Int#
hr HAVLS e
hs of
                                 UBT3(t,ht,hs_) -> AVL e -> Int# -> HAVLS e -> AVL e
forall e. AVL e -> Int# -> HAVLS e -> AVL e
concatHAVLS AVL e
t Int#
ht HAVLS e
hs_


-- Merge adjacent pairs in the current list.
-- The head of the current list has to be treated specially becuase it has no associated
-- bridging element.
-- This function is strict in both elements of the result pair.
{-# INLINE mergePairs #-}
mergePairs :: AVL e -> UINT -> e -> AVL e -> UINT -> HAVLS e -> UBT3(AVL e,UINT,HAVLS e)
mergePairs :: forall e.
AVL e
-> Int#
-> e
-> AVL e
-> Int#
-> HAVLS e
-> (# AVL e, Int#, HAVLS e #)
mergePairs AVL e
l Int#
hl e
e AVL e
r Int#
hr HAVLS e
hs = case AVL e -> Int# -> e -> AVL e -> Int# -> (# AVL e, Int# #)
forall e. AVL e -> Int# -> e -> AVL e -> Int# -> (# AVL e, Int# #)
spliceH AVL e
l Int#
hl e
e AVL e
r Int#
hr of
                            UBT2(t,ht) -> case HAVLS e
hs of
                               HAVLS e
HE              -> UBT3(t,ht,HE)
                               H e
e_ AVL e
t_ Int#
ht_ HAVLS e
hs_ -> let hs__ :: HAVLS e
hs__ = e -> AVL e -> Int# -> HAVLS e -> HAVLS e
forall e. e -> AVL e -> Int# -> HAVLS e -> HAVLS e
mergePairs_ e
e_ AVL e
t_ Int#
ht_ HAVLS e
hs_
                                                  in  HAVLS e
hs__ HAVLS e -> (# AVL e, Int#, HAVLS e #) -> (# AVL e, Int#, HAVLS e #)
forall a b. a -> b -> b
`seq` UBT3(t,ht,hs__)

-- Deals with the rest of mergePairs after the head of the current list has been dealt with.
-- This function is strict in the resulting list head and lazy in the tail.
mergePairs_ :: e -> AVL e -> UINT -> HAVLS e -> HAVLS e
mergePairs_ :: forall e. e -> AVL e -> Int# -> HAVLS e -> HAVLS e
mergePairs_ e
e AVL e
l Int#
hl  HAVLS e
HE            = e -> AVL e -> Int# -> HAVLS e -> HAVLS e
forall e. e -> AVL e -> Int# -> HAVLS e -> HAVLS e
H e
e AVL e
l Int#
hl HAVLS e
forall e. HAVLS e
HE
mergePairs_ e
e AVL e
l Int#
hl (H e
e_ AVL e
r Int#
hr HAVLS e
hs) = case AVL e -> Int# -> e -> AVL e -> Int# -> (# AVL e, Int# #)
forall e. AVL e -> Int# -> e -> AVL e -> Int# -> (# AVL e, Int# #)
spliceH AVL e
l Int#
hl e
e_ AVL e
r Int#
hr of
                                    UBT2(t,ht) -> case HAVLS e
hs of
                                       HAVLS e
HE               -> e -> AVL e -> Int# -> HAVLS e -> HAVLS e
forall e. e -> AVL e -> Int# -> HAVLS e -> HAVLS e
H e
e AVL e
t Int#
ht HAVLS e
forall e. HAVLS e
HE
                                       H e
e__ AVL e
r_ Int#
hr_ HAVLS e
hs_ -> e -> AVL e -> Int# -> HAVLS e -> HAVLS e
forall e. e -> AVL e -> Int# -> HAVLS e -> HAVLS e
H e
e AVL e
t Int#
ht (e -> AVL e -> Int# -> HAVLS e -> HAVLS e
forall e. e -> AVL e -> Int# -> HAVLS e -> HAVLS e
mergePairs_ e
e__ AVL e
r_ Int#
hr_ HAVLS e
hs_)

-- Uses popHL to get the leftmost element from each tree and calculate the (popped) tree height.
-- The popped element is used as a bridging element for splicing purposes.
-- Empty and singleton trees get special treatment.
-- This function is strict in the resulting list head and lazy in the tail.
mkHAVLS :: [AVL e] -> HAVLS e
mkHAVLS :: forall e. [AVL e] -> HAVLS e
mkHAVLS []             = HAVLS e
forall e. HAVLS e
HE
mkHAVLS ( AVL e
E       :[AVL e]
ts) = [AVL e] -> HAVLS e
forall e. [AVL e] -> HAVLS e
mkHAVLS [AVL e]
ts                -- Discard empty trees
mkHAVLS ((N AVL e
l e
e AVL e
r):[AVL e]
ts) = case AVL e -> e -> AVL e -> (# e, AVL e, Int# #)
forall e. AVL e -> e -> AVL e -> (# e, AVL e, Int# #)
popHLN AVL e
l e
e AVL e
r of      -- Never a singlton with N
                         UBT3(e_,t,ht) -> e -> AVL e -> Int# -> HAVLS e -> HAVLS e
forall e. e -> AVL e -> Int# -> HAVLS e -> HAVLS e
H e
e_ AVL e
t Int#
ht ([AVL e] -> HAVLS e
forall e. [AVL e] -> HAVLS e
mkHAVLS [AVL e]
ts)
mkHAVLS ((Z AVL e
l e
e AVL e
r):[AVL e]
ts) = case AVL e -> e -> AVL e -> (# e, AVL e, Int# #)
forall e. AVL e -> e -> AVL e -> (# e, AVL e, Int# #)
popHLZ AVL e
l e
e AVL e
r of
                         UBT3(e_,t,ht) -> if Int# -> Bool
isTrue# (Int#
ht Int# -> Int# -> Int#
EQL L(0))
                                          then e -> [AVL e] -> HAVLS e
forall e. e -> [AVL e] -> HAVLS e
mkHAVLS_ e
e_ [AVL e]
ts                -- Deal with singleton
                                          else e -> AVL e -> Int# -> HAVLS e -> HAVLS e
forall e. e -> AVL e -> Int# -> HAVLS e -> HAVLS e
H e
e_ AVL e
t Int#
ht ([AVL e] -> HAVLS e
forall e. [AVL e] -> HAVLS e
mkHAVLS [AVL e]
ts)        -- Otherwise treat as normal
mkHAVLS ((P AVL e
l e
e AVL e
r):[AVL e]
ts) = case AVL e -> e -> AVL e -> (# e, AVL e, Int# #)
forall e. AVL e -> e -> AVL e -> (# e, AVL e, Int# #)
popHLP AVL e
l e
e AVL e
r of      -- Never a singlton with P
                         UBT3(e_,t,ht) -> e -> AVL e -> Int# -> HAVLS e -> HAVLS e
forall e. e -> AVL e -> Int# -> HAVLS e -> HAVLS e
H e
e_ AVL e
t Int#
ht ([AVL e] -> HAVLS e
forall e. [AVL e] -> HAVLS e
mkHAVLS [AVL e]
ts)
-- Deals with singletons (avoids unnecessary popHL in next in list)
mkHAVLS_ :: e -> [AVL e] -> HAVLS e
mkHAVLS_ :: forall e. e -> [AVL e] -> HAVLS e
mkHAVLS_ e
e []               = e -> AVL e -> Int# -> HAVLS e -> HAVLS e
forall e. e -> AVL e -> Int# -> HAVLS e -> HAVLS e
H e
e AVL e
forall e. AVL e
E L(0) HE    -- End of list reached anyway
mkHAVLS_ e
e (   AVL e
E       :[AVL e]
ts) = e -> [AVL e] -> HAVLS e
forall e. e -> [AVL e] -> HAVLS e
mkHAVLS_ e
e [AVL e]
ts    -- Discard empty trees
mkHAVLS_ e
e (t :: AVL e
t@(N AVL e
l e
_ AVL e
_):[AVL e]
ts) = e -> AVL e -> Int# -> HAVLS e -> HAVLS e
forall e. e -> AVL e -> Int# -> HAVLS e -> HAVLS e
H e
e AVL e
t (Int# -> AVL e -> Int#
forall e. Int# -> AVL e -> Int#
addHeight L(2AVL e
) l) (mkHAVLS ts)
mkHAVLS_ e
e (t :: AVL e
t@(Z AVL e
l e
_ AVL e
_):[AVL e]
ts) = e -> AVL e -> Int# -> HAVLS e -> HAVLS e
forall e. e -> AVL e -> Int# -> HAVLS e -> HAVLS e
H e
e AVL e
t (Int# -> AVL e -> Int#
forall e. Int# -> AVL e -> Int#
addHeight L(1AVL e
) l) (mkHAVLS ts)
mkHAVLS_ e
e (t :: AVL e
t@(P AVL e
_ e
_ AVL e
r):[AVL e]
ts) = e -> AVL e -> Int# -> HAVLS e -> HAVLS e
forall e. e -> AVL e -> Int# -> HAVLS e -> HAVLS e
H e
e AVL e
t (Int# -> AVL e -> Int#
forall e. Int# -> AVL e -> Int#
addHeight L(2AVL e
) r) (mkHAVLS ts)

-- | Similar to 'concatAVL', except the resulting tree is flat.
-- This function evaluates the entire list of trees before constructing the result.
--
-- Complexity: O(n), where n is the total number of elements in the resulting tree.
flatConcat :: [AVL e] -> AVL e
flatConcat :: forall e. [AVL e] -> AVL e
flatConcat [AVL e]
avls = Int -> [e] -> AVL e
forall e. Int -> [e] -> AVL e
asTreeLenL ((Int -> AVL e -> Int) -> Int -> [AVL e] -> Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Int -> AVL e -> Int
forall e. Int -> AVL e -> Int
addSize Int
0 [AVL e]
avls) ((AVL e -> [e] -> [e]) -> [e] -> [AVL e] -> [e]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr AVL e -> [e] -> [e]
forall e. AVL e -> [e] -> [e]
toListL [] [AVL e]
avls)