-- |
-- Copyright   :  (c) Adrian Hey 2004,2005
-- License     :  BSD3
--
-- AVL Tree size related utilities.

{-# OPTIONS_HADDOCK hide #-}
module Data.Tree.AVL.Size
        (-- * AVL tree size utilities
         size,addSize,clipSize,
         addSize#,size#,
        ) where

import Data.Tree.AVL.Internals.Types(AVL(..))
import Data.Tree.AVL.Height(addHeight)

import GHC.Base
#include "ghcdefs.h"

-- | A convenience wrapper for 'addSize#'.
size :: AVL e -> Int
size :: forall e. AVL e -> Int
size AVL e
t = ASINT(addSize# L(0) t)
{-# INLINE size #-}

-- | A convenience wrapper for 'addSize#'.
size# :: AVL e -> UINT
size# :: forall e. AVL e -> Int#
size# AVL e
t = Int# -> AVL e -> Int#
forall e. Int# -> AVL e -> Int#
addSize# L(0AVL e
) t
{-# INLINE size# #-}

-- | See 'addSize#'.
addSize :: Int -> AVL e -> Int
addSize :: forall e. Int -> AVL e -> Int
addSize ASINT(n) AVL e
t = ASINT(addSize# n t)
{-# INLINE addSize #-}

#define AddSize addSize#

{-----------------------------------------
Notes for fast size calculation.
 case (h,avl)
      (0,_      ) -> 0            -- Must be E
      (1,_      ) -> 1            -- Must be (Z  E        _  E       )
      (2,N _ _ _) -> 2            -- Must be (N  E        _ (Z E _ E))
      (2,Z _ _ _) -> 3            -- Must be (Z (Z E _ E) _ (Z E _ E))
      (2,P _ _ _) -> 2            -- Must be (P (Z E _ E) _  E       )
      (3,N _ _ r) -> 2 + size 2 r -- Must be (N (Z E _ E) _  r       )
      (3,P l _ _) -> 2 + size 2 l -- Must be (P  l        _ (Z E _ E))
------------------------------------------}

-- | Fast algorithm to add the size of a tree to the first argument. This avoids visiting about 50% of tree nodes
-- by using fact that trees with small heights can only have particular shapes.
-- So it's still O(n), but with substantial saving in constant factors.
--
-- Complexity: O(n)
AddSize :: UINT -> AVL e -> UINT
AddSize n E         = n
AddSize n (N l _ rAVL e
) = case addHeight L(2) l of
                      L(2) -> INCINT2(n)
                      L(3) -> fas2 INCINT2(n) r
                      h    -> fasNP n h l r
AddSize n (Z l _ rAVL e
) = case addHeight L(1) l of
                      L(1) -> INCINT1(n)
                      L(2) -> INCINT3(n)
                      L(3) -> fas2 (fas2 INCINT1(n) l) r
                      h    -> fasZ n h l r
AddSize n (P l _ rAVL e
) = case addHeight L(2) r of
                      L(2) -> INCINT2(n)
                      L(3) -> fas2 INCINT2(n) l
                      h    -> fasNP n h r l
-- Parent Height (h) >= 4 !!
fasNP,fasZ :: UINT -> UINT -> AVL e -> AVL e -> UINT
fasNP :: forall e. Int# -> Int# -> AVL e -> AVL e -> Int#
fasNP Int#
n Int#
h AVL e
l AVL e
r = Int# -> Int# -> AVL e -> Int#
forall e. Int# -> Int# -> AVL e -> Int#
fasG3 (Int# -> Int# -> AVL e -> Int#
forall e. Int# -> Int# -> AVL e -> Int#
fasG2 INCINT1(n) DECINT2(hAVL e
) l) DECINT1(AVL e
h) r
fasZ :: forall e. Int# -> Int# -> AVL e -> AVL e -> Int#
fasZ  Int#
n Int#
h AVL e
l AVL e
r = Int# -> Int# -> AVL e -> Int#
forall e. Int# -> Int# -> AVL e -> Int#
fasG3 (Int# -> Int# -> AVL e -> Int#
forall e. Int# -> Int# -> AVL e -> Int#
fasG3 INCINT1(n) DECINT1(hAVL e
) l) DECINT1(AVL e
h) r
-- h>=2 !!
fasG2 :: UINT -> UINT -> AVL e -> UINT
fasG2 :: forall e. Int# -> Int# -> AVL e -> Int#
fasG2 Int#
n L(2)  t        = fas2  n   t
fasG2 Int#
n Int#
h     AVL e
t        = Int# -> Int# -> AVL e -> Int#
forall e. Int# -> Int# -> AVL e -> Int#
fasG3 Int#
n Int#
h AVL e
t
{-# INLINE fasG2 #-}
-- h>=3 !!
fasG3 :: UINT -> UINT -> AVL e -> UINT
fasG3 :: forall e. Int# -> Int# -> AVL e -> Int#
fasG3 Int#
n L(3) (AVL e
N e
_ AVL e
_ r) = fas2 INCINT2(AVL e
n) r
fasG3 Int#
n L(3) (AVL e
Z e
l AVL e
_ r) = fas2 (fas2 INCINT1(n) AVL e
l) r
fasG3 Int#
n L(3) (AVL e
P e
l AVL e
_ _) = fas2 INCINT2(AVL e
n) l
fasG3 Int#
n Int#
h    (N AVL e
l e
_ AVL e
r) = Int# -> Int# -> AVL e -> AVL e -> Int#
forall e. Int# -> Int# -> AVL e -> AVL e -> Int#
fasNP Int#
n Int#
h AVL e
l AVL e
r -- h>=4
fasG3 Int#
n Int#
h    (Z AVL e
l e
_ AVL e
r) = Int# -> Int# -> AVL e -> AVL e -> Int#
forall e. Int# -> Int# -> AVL e -> AVL e -> Int#
fasZ  Int#
n Int#
h AVL e
l AVL e
r -- h>=4
fasG3 Int#
n Int#
h    (P AVL e
l e
_ AVL e
r) = Int# -> Int# -> AVL e -> AVL e -> Int#
forall e. Int# -> Int# -> AVL e -> AVL e -> Int#
fasNP Int#
n Int#
h AVL e
r AVL e
l -- h>=4
fasG3 Int#
_ Int#
_     AVL e
E        = [Char] -> Int#
forall a. HasCallStack => [Char] -> a
error [Char]
"AddSize: Bad Tree." -- impossible
-- h=2 !!
fas2 :: UINT -> AVL e -> UINT
fas2 :: forall e. Int# -> AVL e -> Int#
fas2 Int#
n (N AVL e
_ e
_ AVL e
_) = INCINT2(n)
fas2 Int#
n (Z AVL e
_ e
_ AVL e
_) = INCINT3(n)
fas2 Int#
n (P AVL e
_ e
_ AVL e
_) = INCINT2(n)
fas2 Int#
_  AVL e
E        = [Char] -> Int#
forall a. HasCallStack => [Char] -> a
error [Char]
"AddSize: Bad Tree." -- impossible
{-# INLINE fas2 #-}

-- | Returns the exact tree size in the form @('Just' n)@ if this is less than or
-- equal to the input clip value. Returns @'Nothing'@ of the size is greater than
-- the clip value. This function exploits the same optimisation as 'addSize'.
--
-- Complexity: O(min n c) where n is tree size and c is clip value.
clipSize ::  Int -> AVL e -> Maybe Int
clipSize :: forall e. Int -> AVL e -> Maybe Int
clipSize ASINT(c) AVL e
t = let c_ :: Int#
c_ = Int# -> AVL e -> Int#
forall e. Int# -> AVL e -> Int#
cSzh Int#
c AVL e
t in if   Int# -> Bool
isTrue# (Int#
c_ LTN L(0))
                                           then Maybe Int
forall a. Maybe a
Nothing
                                           else Int -> Maybe Int
forall a. a -> Maybe a
Just ASINT(SUBINT(c,c_))
-- First entry calculates initial height
cSzh :: UINT -> AVL e -> UINT
cSzh :: forall e. Int# -> AVL e -> Int#
cSzh Int#
c  AVL e
E        = Int#
c
cSzh Int#
c (N AVL e
l e
_ AVL e
r) = case Int# -> AVL e -> Int#
forall e. Int# -> AVL e -> Int#
addHeight L(2AVL e
) l of
                   L(2) -> DECINT2(c)
                   L(3) -> cSzNP3 c     r
                   Int#
h    -> Int# -> Int# -> AVL e -> AVL e -> Int#
forall e. Int# -> Int# -> AVL e -> AVL e -> Int#
cSzNP  Int#
c Int#
h AVL e
l AVL e
r
cSzh Int#
c (Z AVL e
l e
_ AVL e
r) = case Int# -> AVL e -> Int#
forall e. Int# -> AVL e -> Int#
addHeight L(1AVL e
) l of
                   L(1) -> DECINT1(c)
                   L(2) -> DECINT3(c)
                   L(3) -> cSzZ3 c   l r
                   Int#
h    -> Int# -> Int# -> AVL e -> AVL e -> Int#
forall e. Int# -> Int# -> AVL e -> AVL e -> Int#
cSzZ  Int#
c Int#
h AVL e
l AVL e
r
cSzh Int#
c (P AVL e
l e
_ AVL e
r) = case Int# -> AVL e -> Int#
forall e. Int# -> AVL e -> Int#
addHeight L(2AVL e
) r of
                   L(2) -> DECINT2(c)
                   L(3) -> cSzNP3 c     l
                   Int#
h    -> Int# -> Int# -> AVL e -> AVL e -> Int#
forall e. Int# -> Int# -> AVL e -> AVL e -> Int#
cSzNP  Int#
c Int#
h AVL e
r AVL e
l
-- Parent Height = 3 !!
cSzNP3 :: UINT -> AVL e -> UINT
cSzNP3 :: forall e. Int# -> AVL e -> Int#
cSzNP3 Int#
c AVL e
t = if Int# -> Bool
isTrue# (Int#
c LTN L(4)) then L(-1) Int# -> AVL e -> Int#
forall e. Int# -> AVL e -> Int#
else cSz2 DECINT2(c) t
cSzZ3  :: UINT -> AVL e -> AVL e -> UINT
cSzZ3 :: forall e. Int# -> AVL e -> AVL e -> Int#
cSzZ3  Int#
c AVL e
l AVL e
r = if Int# -> Bool
isTrue# (Int#
c LTN L(5)) then L(-1)
                             else let c_ :: Int#
c_ = Int# -> AVL e -> Int#
forall e. Int# -> AVL e -> Int#
cSz2 DECINT1(c) l
                                  in if Int# -> Bool
isTrue# (Int#
c_ LTN L(2)) then L(-1)
                                                    else Int# -> AVL e -> Int#
forall e. Int# -> AVL e -> Int#
cSz2 Int#
c_ AVL e
r
-- Parent Height (h) >= 4 !!
cSzNP,cSzZ :: UINT -> UINT -> AVL e -> AVL e -> UINT
cSzNP :: forall e. Int# -> Int# -> AVL e -> AVL e -> Int#
cSzNP Int#
c Int#
h AVL e
l AVL e
r = if Int# -> Bool
isTrue# (Int#
c LTN L(7)) then L(-1)
                              else let c_ :: Int#
c_ = Int# -> Int# -> AVL e -> Int#
forall e. Int# -> Int# -> AVL e -> Int#
cSzG2 DECINT1(c) DECINT2(hAVL e
) l       -- (h-2) >= 2
                                   in if Int# -> Bool
isTrue# (Int#
c_ LTN L(4)) then L(-1)
                                                     else Int# -> Int# -> AVL e -> Int#
forall e. Int# -> Int# -> AVL e -> Int#
cSzG3 Int#
c_ DECINT1(h) r -- (h-1) >= 3
cSzZ :: forall e. Int# -> Int# -> AVL e -> AVL e -> Int#
cSzZ Int#
c Int#
h AVL e
l AVL e
r = if Int# -> Bool
isTrue# (Int#
c LTN L(9)) then L(-1)
                             else let c_ :: Int#
c_ = Int# -> Int# -> AVL e -> Int#
forall e. Int# -> Int# -> AVL e -> Int#
cSzG3 DECINT1(c) DECINT1(hAVL e
) l        -- (h-1) >= 3
                                  in if Int# -> Bool
isTrue# (Int#
c_ LTN L(4)) then L(-1)
                                                    else Int# -> Int# -> AVL e -> Int#
forall e. Int# -> Int# -> AVL e -> Int#
cSzG3 Int#
c_ DECINT1(h) r  -- (h-1) >= 3
-- h>=2 !!
cSzG2 :: UINT -> UINT -> AVL e -> UINT
cSzG2 :: forall e. Int# -> Int# -> AVL e -> Int#
cSzG2 Int#
c L(2)  t        = cSz2  c   t
cSzG2 Int#
c Int#
h     AVL e
t        = Int# -> Int# -> AVL e -> Int#
forall e. Int# -> Int# -> AVL e -> Int#
cSzG3 Int#
c Int#
h AVL e
t
{-# INLINE cSzG2 #-}
-- h>=3 !!
cSzG3 :: UINT -> UINT -> AVL e -> UINT
cSzG3 :: forall e. Int# -> Int# -> AVL e -> Int#
cSzG3 Int#
c L(3) (AVL e
N e
_ AVL e
_ r) = cSzNP3 c   r
cSzG3 Int#
c L(3) (AVL e
Z e
l AVL e
_ r) = cSzZ3  AVL e
c AVL e
l r
cSzG3 Int#
c L(3) (AVL e
P e
l AVL e
_ _) = cSzNP3 AVL e
c l
cSzG3 Int#
c Int#
h    (N AVL e
l e
_ AVL e
r) = Int# -> Int# -> AVL e -> AVL e -> Int#
forall e. Int# -> Int# -> AVL e -> AVL e -> Int#
cSzNP Int#
c Int#
h AVL e
l AVL e
r -- h>=4
cSzG3 Int#
c Int#
h    (Z AVL e
l e
_ AVL e
r) = Int# -> Int# -> AVL e -> AVL e -> Int#
forall e. Int# -> Int# -> AVL e -> AVL e -> Int#
cSzZ  Int#
c Int#
h AVL e
l AVL e
r -- h>=4
cSzG3 Int#
c Int#
h    (P AVL e
l e
_ AVL e
r) = Int# -> Int# -> AVL e -> AVL e -> Int#
forall e. Int# -> Int# -> AVL e -> AVL e -> Int#
cSzNP Int#
c Int#
h AVL e
r AVL e
l -- h>=4
cSzG3 Int#
_ Int#
_     AVL e
E        = [Char] -> Int#
forall a. HasCallStack => [Char] -> a
error [Char]
"clipSize: Bad Tree." -- impossible
-- h=2 !!
cSz2 :: UINT -> AVL e -> UINT
cSz2 :: forall e. Int# -> AVL e -> Int#
cSz2 Int#
c (N AVL e
_ e
_ AVL e
_) = DECINT2(c)
cSz2 Int#
c (Z AVL e
_ e
_ AVL e
_) = DECINT3(c)
cSz2 Int#
c (P AVL e
_ e
_ AVL e
_) = DECINT2(c)
cSz2 Int#
_  AVL e
E        = [Char] -> Int#
forall a. HasCallStack => [Char] -> a
error [Char]
"clipSize: Bad Tree." -- impossible
{-# INLINE cSz2 #-}