-- | Binary trees
--
-- Intended for qualified import.
--
-- > import Data.Falsify.Tree (Tree(..))
-- > import qualified Data.Falsify.Tree as Tree
module Data.Falsify.Tree (
    Tree(Leaf, Branch)
    -- * Properties
  , size
  , weight
  , height
    -- * BST
  , lookup
    -- * Balancing
  , isWeightBalanced
  , isHeightBalanced
    -- * Debugging
  , render
  ) where

import Prelude hiding (lookup)

import GHC.Show

import qualified Data.Tree as Rose

{-------------------------------------------------------------------------------
  Definition
-------------------------------------------------------------------------------}

-- | Binary tree
--
-- Each branch caches the size of the subtree, so that 'size' can be @O(1)@.
data Tree a =
    Leaf

    -- 'Branch_' caches the size of the tree
  | Branch_ {-# UNPACK #-} !Word a (Tree a) (Tree a)
  deriving stock (Tree a -> Tree a -> Bool
(Tree a -> Tree a -> Bool)
-> (Tree a -> Tree a -> Bool) -> Eq (Tree a)
forall a. Eq a => Tree a -> Tree a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Tree a -> Tree a -> Bool
== :: Tree a -> Tree a -> Bool
$c/= :: forall a. Eq a => Tree a -> Tree a -> Bool
/= :: Tree a -> Tree a -> Bool
Eq, (forall a b. (a -> b) -> Tree a -> Tree b)
-> (forall a b. a -> Tree b -> Tree a) -> Functor Tree
forall a b. a -> Tree b -> Tree a
forall a b. (a -> b) -> Tree a -> Tree b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> Tree a -> Tree b
fmap :: forall a b. (a -> b) -> Tree a -> Tree b
$c<$ :: forall a b. a -> Tree b -> Tree a
<$ :: forall a b. a -> Tree b -> Tree a
Functor, (forall m. Monoid m => Tree m -> m)
-> (forall m a. Monoid m => (a -> m) -> Tree a -> m)
-> (forall m a. Monoid m => (a -> m) -> Tree a -> m)
-> (forall a b. (a -> b -> b) -> b -> Tree a -> b)
-> (forall a b. (a -> b -> b) -> b -> Tree a -> b)
-> (forall b a. (b -> a -> b) -> b -> Tree a -> b)
-> (forall b a. (b -> a -> b) -> b -> Tree a -> b)
-> (forall a. (a -> a -> a) -> Tree a -> a)
-> (forall a. (a -> a -> a) -> Tree a -> a)
-> (forall a. Tree a -> [a])
-> (forall a. Tree a -> Bool)
-> (forall a. Tree a -> Int)
-> (forall a. Eq a => a -> Tree a -> Bool)
-> (forall a. Ord a => Tree a -> a)
-> (forall a. Ord a => Tree a -> a)
-> (forall a. Num a => Tree a -> a)
-> (forall a. Num a => Tree a -> a)
-> Foldable Tree
forall a. Eq a => a -> Tree a -> Bool
forall a. Num a => Tree a -> a
forall a. Ord a => Tree a -> a
forall m. Monoid m => Tree m -> m
forall a. Tree a -> Bool
forall a. Tree a -> Int
forall a. Tree a -> [a]
forall a. (a -> a -> a) -> Tree a -> a
forall m a. Monoid m => (a -> m) -> Tree a -> m
forall b a. (b -> a -> b) -> b -> Tree a -> b
forall a b. (a -> b -> b) -> b -> Tree a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
$cfold :: forall m. Monoid m => Tree m -> m
fold :: forall m. Monoid m => Tree m -> m
$cfoldMap :: forall m a. Monoid m => (a -> m) -> Tree a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> Tree a -> m
$cfoldMap' :: forall m a. Monoid m => (a -> m) -> Tree a -> m
foldMap' :: forall m a. Monoid m => (a -> m) -> Tree a -> m
$cfoldr :: forall a b. (a -> b -> b) -> b -> Tree a -> b
foldr :: forall a b. (a -> b -> b) -> b -> Tree a -> b
$cfoldr' :: forall a b. (a -> b -> b) -> b -> Tree a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> Tree a -> b
$cfoldl :: forall b a. (b -> a -> b) -> b -> Tree a -> b
foldl :: forall b a. (b -> a -> b) -> b -> Tree a -> b
$cfoldl' :: forall b a. (b -> a -> b) -> b -> Tree a -> b
foldl' :: forall b a. (b -> a -> b) -> b -> Tree a -> b
$cfoldr1 :: forall a. (a -> a -> a) -> Tree a -> a
foldr1 :: forall a. (a -> a -> a) -> Tree a -> a
$cfoldl1 :: forall a. (a -> a -> a) -> Tree a -> a
foldl1 :: forall a. (a -> a -> a) -> Tree a -> a
$ctoList :: forall a. Tree a -> [a]
toList :: forall a. Tree a -> [a]
$cnull :: forall a. Tree a -> Bool
null :: forall a. Tree a -> Bool
$clength :: forall a. Tree a -> Int
length :: forall a. Tree a -> Int
$celem :: forall a. Eq a => a -> Tree a -> Bool
elem :: forall a. Eq a => a -> Tree a -> Bool
$cmaximum :: forall a. Ord a => Tree a -> a
maximum :: forall a. Ord a => Tree a -> a
$cminimum :: forall a. Ord a => Tree a -> a
minimum :: forall a. Ord a => Tree a -> a
$csum :: forall a. Num a => Tree a -> a
sum :: forall a. Num a => Tree a -> a
$cproduct :: forall a. Num a => Tree a -> a
product :: forall a. Num a => Tree a -> a
Foldable, Functor Tree
Foldable Tree
(Functor Tree, Foldable Tree) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> Tree a -> f (Tree b))
-> (forall (f :: * -> *) a.
    Applicative f =>
    Tree (f a) -> f (Tree a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> Tree a -> m (Tree b))
-> (forall (m :: * -> *) a. Monad m => Tree (m a) -> m (Tree a))
-> Traversable Tree
forall (t :: * -> *).
(Functor t, Foldable t) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (m :: * -> *) a. Monad m => Tree (m a) -> m (Tree a)
forall (f :: * -> *) a. Applicative f => Tree (f a) -> f (Tree a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Tree a -> m (Tree b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Tree a -> f (Tree b)
$ctraverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Tree a -> f (Tree b)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Tree a -> f (Tree b)
$csequenceA :: forall (f :: * -> *) a. Applicative f => Tree (f a) -> f (Tree a)
sequenceA :: forall (f :: * -> *) a. Applicative f => Tree (f a) -> f (Tree a)
$cmapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Tree a -> m (Tree b)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Tree a -> m (Tree b)
$csequence :: forall (m :: * -> *) a. Monad m => Tree (m a) -> m (Tree a)
sequence :: forall (m :: * -> *) a. Monad m => Tree (m a) -> m (Tree a)
Traversable)

{-------------------------------------------------------------------------------
  Properties
-------------------------------------------------------------------------------}

-- | Size of the tree
--
-- @O(1)@
size :: Tree a -> Word
size :: forall a. Tree a -> Word
size Tree a
Leaf              = Word
0
size (Branch_ Word
s a
_ Tree a
_ Tree a
_) = Word
s

-- | Weight of the tree
--
-- The weight of a tree is simply its size plus one.
--
-- @O(1)@
weight :: Tree a -> Word
weight :: forall a. Tree a -> Word
weight = Word -> Word
forall a. Enum a => a -> a
succ (Word -> Word) -> (Tree a -> Word) -> Tree a -> Word
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree a -> Word
forall a. Tree a -> Word
size

-- | Height of the tree
--
-- The height of a tree is the maximum length from the root to any of the leafs.
--
-- @O(1)@
height :: Tree a -> Word
height :: forall a. Tree a -> Word
height Tree a
Leaf           = Word
0
height (Branch a
_ Tree a
l Tree a
r) = Word
1 Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word -> Word -> Word
forall a. Ord a => a -> a -> a
max (Tree a -> Word
forall a. Tree a -> Word
height Tree a
l) (Tree a -> Word
forall a. Tree a -> Word
height Tree a
r)


{-------------------------------------------------------------------------------
  Pattern synonyms that hide the size argument
-------------------------------------------------------------------------------}

viewBranch :: Tree a -> Maybe (a, Tree a, Tree a)
viewBranch :: forall a. Tree a -> Maybe (a, Tree a, Tree a)
viewBranch Tree a
Leaf              = Maybe (a, Tree a, Tree a)
forall a. Maybe a
Nothing
viewBranch (Branch_ Word
_ a
x Tree a
l Tree a
r) = (a, Tree a, Tree a) -> Maybe (a, Tree a, Tree a)
forall a. a -> Maybe a
Just (a
x, Tree a
l, Tree a
r)

branch :: a -> Tree a -> Tree a -> Tree a
branch :: forall a. a -> Tree a -> Tree a -> Tree a
branch a
x Tree a
l Tree a
r = Word -> a -> Tree a -> Tree a -> Tree a
forall a. Word -> a -> Tree a -> Tree a -> Tree a
Branch_ (Word
1 Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Tree a -> Word
forall a. Tree a -> Word
size Tree a
l Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Tree a -> Word
forall a. Tree a -> Word
size Tree a
r) a
x Tree a
l Tree a
r

pattern Branch :: a -> Tree a -> Tree a -> Tree a
pattern $mBranch :: forall {r} {a}.
Tree a -> (a -> Tree a -> Tree a -> r) -> ((# #) -> r) -> r
$bBranch :: forall a. a -> Tree a -> Tree a -> Tree a
Branch x l r <- (viewBranch -> Just (x, l, r))
  where
    Branch = a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
branch

{-# COMPLETE Leaf, Branch #-}

{-------------------------------------------------------------------------------
  'Show' instance that depends on the pattern synonyms
-------------------------------------------------------------------------------}

instance Show a => Show (Tree a) where
  showsPrec :: Int -> Tree a -> ShowS
showsPrec Int
_ Tree a
Leaf           = String -> ShowS
showString String
"Leaf"
  showsPrec Int
a (Branch a
x Tree a
l Tree a
r) = Bool -> ShowS -> ShowS
showParen (Int
a Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
appPrec) (ShowS -> ShowS) -> ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$
        String -> ShowS
showString String
"Branch "
      ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> a -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
appPrec1 a
x
      ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS
showSpace
      ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Tree a -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
appPrec1 Tree a
l
      ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS
showSpace
      ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Tree a -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
appPrec1 Tree a
r

{-------------------------------------------------------------------------------
  BST
-------------------------------------------------------------------------------}

-- | Look value up in BST
--
-- NOTE: The 'Tree' datatype itself does /NOT/ guarantee that the tree is in
-- fact a BST. It is the responsibility of the caller to ensure this.
lookup :: Ord a => a -> Tree (a, b) -> Maybe b
lookup :: forall a b. Ord a => a -> Tree (a, b) -> Maybe b
lookup a
a' (Branch (a
a, b
b) Tree (a, b)
l Tree (a, b)
r)
  | a
a' a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
a    = a -> Tree (a, b) -> Maybe b
forall a b. Ord a => a -> Tree (a, b) -> Maybe b
lookup a
a' Tree (a, b)
l
  | a
a' a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
a    = a -> Tree (a, b) -> Maybe b
forall a b. Ord a => a -> Tree (a, b) -> Maybe b
lookup a
a' Tree (a, b)
r
  | Bool
otherwise = b -> Maybe b
forall a. a -> Maybe a
Just b
b
lookup a
_ Tree (a, b)
Leaf = Maybe b
forall a. Maybe a
Nothing

{-------------------------------------------------------------------------------
  Balancing
-------------------------------------------------------------------------------}

-- | Check if the tree is weight-balanced
--
-- A tree is weight-balanced if the weights of the subtrees does not differ
-- by more than a factor 3.
--
-- See "Balancing weight-balanced trees", Hirai and Yamamoto, JFP 21(3), 2011.
isWeightBalanced :: Tree a -> Bool
isWeightBalanced :: forall a. Tree a -> Bool
isWeightBalanced = (Tree a -> Tree a -> Bool) -> Tree a -> Bool
forall a. (Tree a -> Tree a -> Bool) -> Tree a -> Bool
checkBalanceCondition Tree a -> Tree a -> Bool
forall a. Tree a -> Tree a -> Bool
isBalanced
  where
    delta :: Word
    delta :: Word
delta = Word
3

    isBalanced :: Tree a -> Tree a -> Bool
    isBalanced :: forall a. Tree a -> Tree a -> Bool
isBalanced Tree a
a Tree a
b = [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and [
          Word
delta Word -> Word -> Word
forall a. Num a => a -> a -> a
* Tree a -> Word
forall a. Tree a -> Word
weight Tree a
a Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
>= Tree a -> Word
forall a. Tree a -> Word
weight Tree a
b
        , Word
delta Word -> Word -> Word
forall a. Num a => a -> a -> a
* Tree a -> Word
forall a. Tree a -> Word
weight Tree a
b Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
>= Tree a -> Word
forall a. Tree a -> Word
weight Tree a
a
        ]

-- | Check if a tree is height-balanced
--
-- A tree is height balanced if the heights of its subtrees do not differ
-- by more than one.
isHeightBalanced :: Tree a -> Bool
isHeightBalanced :: forall a. Tree a -> Bool
isHeightBalanced = (Tree a -> Tree a -> Bool) -> Tree a -> Bool
forall a. (Tree a -> Tree a -> Bool) -> Tree a -> Bool
checkBalanceCondition Tree a -> Tree a -> Bool
forall a. Tree a -> Tree a -> Bool
isBalanced
  where
    isBalanced :: Tree a -> Tree a -> Bool
    isBalanced :: forall a. Tree a -> Tree a -> Bool
isBalanced Tree a
a Tree a
b = [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or [
          (Tree a -> Word
forall a. Tree a -> Word
height Tree a
a Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
<= Tree a -> Word
forall a. Tree a -> Word
height Tree a
b) Bool -> Bool -> Bool
&& (Tree a -> Word
forall a. Tree a -> Word
height Tree a
b Word -> Word -> Word
forall a. Num a => a -> a -> a
- Tree a -> Word
forall a. Tree a -> Word
height Tree a
a Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
<= Word
1)
        , (Tree a -> Word
forall a. Tree a -> Word
height Tree a
b Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
<= Tree a -> Word
forall a. Tree a -> Word
height Tree a
a) Bool -> Bool -> Bool
&& (Tree a -> Word
forall a. Tree a -> Word
height Tree a
a Word -> Word -> Word
forall a. Num a => a -> a -> a
- Tree a -> Word
forall a. Tree a -> Word
height Tree a
b Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
<= Word
1)
        ]

-- | Internal auxiliary: check given tree balance condition
--
-- Property @p l r@ will be checked at every branch in the tree.
checkBalanceCondition :: forall a. (Tree a -> Tree a -> Bool) -> Tree a -> Bool
checkBalanceCondition :: forall a. (Tree a -> Tree a -> Bool) -> Tree a -> Bool
checkBalanceCondition Tree a -> Tree a -> Bool
p = Tree a -> Bool
go
  where
    go :: Tree a -> Bool
    go :: Tree a -> Bool
go Tree a
Leaf           = Bool
True
    go (Branch a
_ Tree a
l Tree a
r) = [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and [Tree a -> Tree a -> Bool
p Tree a
l Tree a
r, Tree a -> Bool
go Tree a
l, Tree a -> Bool
go Tree a
r]

{-------------------------------------------------------------------------------
  Debugging
-------------------------------------------------------------------------------}

-- | Render tree
--
-- This is intended for debugging only.
render :: Tree String -> String
render :: Tree String -> String
render = Tree String -> String
Rose.drawTree (Tree String -> String)
-> (Tree String -> Tree String) -> Tree String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree String -> Tree String
conv
  where
    conv :: Tree String -> Rose.Tree String
    conv :: Tree String -> Tree String
conv Tree String
Leaf           = String -> [Tree String] -> Tree String
forall a. a -> [Tree a] -> Tree a
Rose.Node String
"*" []
    conv (Branch String
x Tree String
l Tree String
r) = String -> [Tree String] -> Tree String
forall a. a -> [Tree a] -> Tree a
Rose.Node String
x [Tree String -> Tree String
conv Tree String
l, Tree String -> Tree String
conv Tree String
r]