{-# OPTIONS_GHC -fglasgow-exts #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Tree.AVL.Size
-- Copyright   :  (c) Adrian Hey 2004,2005
-- License     :  BSD3
--
-- Maintainer  :  http://homepages.nildram.co.uk/~ahey/em.png
-- Stability   :  stable
-- Portability :  portable
--
-- AVL Tree size related utilities.
-----------------------------------------------------------------------------
module Data.Tree.AVL.Size
        (-- * AVL tree size utilities.
         size,addSize,fastAddSize,

        ) where

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

#ifdef __GLASGOW_HASKELL__
import GHC.Base
#include "ghcdefs.h"
#else
#include "h98defs.h"
#endif

-- | Counts the total number of elements in an AVL tree.
--
-- Complexity: O(n)
{-# INLINE size #-}
size :: AVL e -> Int
size = addSize 0

-- | Adds the size of a tree to the first argument.
--
-- Complexity: O(n)
{-# INLINE addSize #-}
addSize :: Int -> AVL e -> Int
addSize ASINT(n) t = ASINT(fastAddSize n t)

-- | Fast algorithm to calculate size. 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)
fastAddSize :: UINT -> AVL e -> UINT
fastAddSize n E         = n
fastAddSize n (N l _ r) = case addHeight L(2) l of
                          L(2) -> INCINT2(n)
                          h    -> fasN n h l r
fastAddSize n (Z l _ r) = case addHeight L(1) l of
                          L(1) -> INCINT1(n)
                          L(2) -> INCINT3(n)
                          h    -> fasZ n h l r
fastAddSize n (P l _ r) = case addHeight L(2) r of
                          L(2) -> INCINT2(n)
                          h    -> fasP n h l r

-- Local utilities used by fastAddSize, Only work if h >=3 !!
fasN,fasZ,fasP :: UINT -> UINT -> AVL e -> AVL e -> UINT
fasN n L(3) _ r = fas INCINT2(n)                    L(2)       r
fasN n h    l r = fas (fas INCINT1(n) DECINT2(h) l) DECINT1(h) r -- h>=4
fasZ n h    l r = fas (fas INCINT1(n) DECINT1(h) l) DECINT1(h) r
fasP n L(3) l _ = fas INCINT2(n)                    L(2)       l
fasP n h    l r = fas (fas INCINT1(n) DECINT2(h) r) DECINT1(h) l -- h>=4

-- Local Utility used by fasN,fasZ,fasP, Only works if h >= 2 !!
fas :: UINT -> UINT -> AVL e -> UINT
fas _ L(2)  E        = error "fas: Bug0"
fas n L(2) (N _ _ _) = INCINT2(n)
fas n L(2) (Z _ _ _) = INCINT3(n)
fas n L(2) (P _ _ _) = INCINT2(n)
-- So h must be >= 3 if we get here
fas n h    (N l _ r) = fasN n h l r
fas n h    (Z l _ r) = fasZ n h l r
fas n h    (P l _ r) = fasP n h l r
fas _ _     E        = error "fas: Bug1"

{-----------------------------------------
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))
------------------------------------------}