module Data.Falsify.Tree (
    Tree(Leaf, Branch)
    -- * Dealing with marks
  , propagate
  , genKept
  , keepAtLeast
    -- * Binary search trees
  , Interval(..)
  , Endpoint(..)
  , inclusiveBounds
  , lookup
    -- * Debugging
  , drawTree
  ) where

import Prelude hiding (drop, lookup)

import Control.Selective (Selective, ifS)
import Control.Monad.State
import GHC.Show

import qualified Data.Tree as Rose

import Data.Falsify.Marked

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

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
forall a. Eq a => Tree a -> Tree a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Tree a -> Tree a -> Bool
$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
Eq, 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
<$ :: forall a b. a -> Tree b -> Tree a
$c<$ :: forall a b. a -> Tree b -> Tree a
fmap :: forall a b. (a -> b) -> Tree a -> Tree b
$cfmap :: forall a b. (a -> b) -> Tree a -> Tree b
Functor, 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
product :: forall a. Num a => Tree a -> a
$cproduct :: forall a. Num a => Tree a -> a
sum :: forall a. Num a => Tree a -> a
$csum :: forall a. Num a => Tree a -> a
minimum :: forall a. Ord a => Tree a -> a
$cminimum :: forall a. Ord a => Tree a -> a
maximum :: forall a. Ord a => Tree a -> a
$cmaximum :: forall a. Ord a => Tree a -> a
elem :: forall a. Eq a => a -> Tree a -> Bool
$celem :: forall a. Eq a => a -> Tree a -> Bool
length :: forall a. Tree a -> Int
$clength :: forall a. Tree a -> Int
null :: forall a. Tree a -> Bool
$cnull :: forall a. Tree a -> Bool
toList :: forall a. Tree a -> [a]
$ctoList :: forall a. Tree a -> [a]
foldl1 :: forall a. (a -> a -> a) -> Tree a -> a
$cfoldl1 :: forall a. (a -> a -> a) -> Tree a -> a
foldr1 :: forall a. (a -> a -> a) -> Tree a -> a
$cfoldr1 :: forall a. (a -> a -> a) -> Tree a -> a
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
$cfoldl :: forall b a. (b -> a -> 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
$cfoldr :: forall a b. (a -> b -> b) -> b -> Tree a -> b
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
$cfoldMap :: forall m a. Monoid m => (a -> m) -> Tree a -> m
fold :: forall m. Monoid m => Tree m -> m
$cfold :: forall m. Monoid m => Tree m -> m
Foldable, Functor Tree
Foldable 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)
sequence :: forall (m :: * -> *) a. Monad m => Tree (m a) -> m (Tree a)
$csequence :: forall (m :: * -> *) a. Monad m => Tree (m a) -> m (Tree a)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Tree a -> m (Tree b)
$cmapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Tree a -> m (Tree b)
sequenceA :: forall (f :: * -> *) a. Applicative f => Tree (f a) -> f (Tree a)
$csequenceA :: forall (f :: * -> *) a. Applicative f => Tree (f a) -> f (Tree a)
traverse :: 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)
Traversable)

{-------------------------------------------------------------------------------
  Tree stats
-------------------------------------------------------------------------------}

-- | 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

{-------------------------------------------------------------------------------
  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              = forall a. Maybe a
Nothing
viewBranch (Branch_ Word
_ a
x Tree a
l Tree a
r) = 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 = forall a. Word -> a -> Tree a -> Tree a -> Tree a
Branch_ (Word
1 forall a. Num a => a -> a -> a
+ forall a. Tree a -> Word
size Tree a
l forall a. Num a => a -> a -> a
+ 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 $bBranch :: forall a. a -> Tree a -> Tree a -> Tree a
$mBranch :: forall {r} {a}.
Tree a -> (a -> Tree a -> Tree a -> r) -> ((# #) -> r) -> r
Branch x l r <- (viewBranch -> Just (x, l, r))
  where
    Branch = 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 forall a. Ord a => a -> a -> Bool
> Int
appPrec) forall a b. (a -> b) -> a -> b
$
        String -> ShowS
showString String
"Branch "
      forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => Int -> a -> ShowS
showsPrec Int
appPrec1 a
x
      forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS
showSpace
      forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => Int -> a -> ShowS
showsPrec Int
appPrec1 Tree a
l
      forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS
showSpace
      forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => Int -> a -> ShowS
showsPrec Int
appPrec1 Tree a
r

{-------------------------------------------------------------------------------
  Dealing with marks
-------------------------------------------------------------------------------}

-- | Propagate 'Drop' marker down the tree
--
-- This is useful in conjunction with 'genKept', which truncates entire
-- subtrees.
propagate :: Tree (Marked f a) -> Tree (Marked f a)
propagate :: forall (f :: * -> *) a. Tree (Marked f a) -> Tree (Marked f a)
propagate = forall (f :: * -> *) a. Tree (Marked f a) -> Tree (Marked f a)
keep
  where
    keep :: Tree (Marked f a) -> Tree (Marked f a)
    keep :: forall (f :: * -> *) a. Tree (Marked f a) -> Tree (Marked f a)
keep Tree (Marked f a)
Leaf                         = forall a. Tree a
Leaf
    keep (Branch (Marked Mark
Keep f a
x) Tree (Marked f a)
l Tree (Marked f a)
r) = forall a. a -> Tree a -> Tree a -> Tree a
Branch (forall (f :: * -> *) a. Mark -> f a -> Marked f a
Marked Mark
Keep f a
x) (forall (f :: * -> *) a. Tree (Marked f a) -> Tree (Marked f a)
keep Tree (Marked f a)
l) (forall (f :: * -> *) a. Tree (Marked f a) -> Tree (Marked f a)
keep Tree (Marked f a)
r)
    keep (Branch (Marked Mark
Drop f a
x) Tree (Marked f a)
l Tree (Marked f a)
r) = forall a. a -> Tree a -> Tree a -> Tree a
Branch (forall (f :: * -> *) a. Mark -> f a -> Marked f a
Marked Mark
Drop f a
x) (forall (f :: * -> *) a. Tree (Marked f a) -> Tree (Marked f a)
drop Tree (Marked f a)
l) (forall (f :: * -> *) a. Tree (Marked f a) -> Tree (Marked f a)
drop Tree (Marked f a)
r)

    drop :: Tree (Marked f a) -> Tree (Marked f a)
    drop :: forall (f :: * -> *) a. Tree (Marked f a) -> Tree (Marked f a)
drop = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a -> b) -> a -> b
$ \(Marked Mark
_ f a
x) -> forall (f :: * -> *) a. Mark -> f a -> Marked f a
Marked Mark
Drop f a
x

-- | Generate those values we want to keep
--
-- Whenever we meet an element marked 'Drop', that entire subtree is dropped.
genKept :: forall f a. Selective f => Tree (Marked f a) -> f (Tree a)
genKept :: forall (f :: * -> *) a.
Selective f =>
Tree (Marked f a) -> f (Tree a)
genKept = Tree (Marked f a) -> f (Tree a)
go
  where
    go :: Tree (Marked f a) -> f (Tree a)
    go :: Tree (Marked f a) -> f (Tree a)
go Tree (Marked f a)
Leaf                      = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Tree a
Leaf
    go (Branch (Marked Mark
m f a
g) Tree (Marked f a)
l Tree (Marked f a)
r) = forall (f :: * -> *) a. Selective f => f Bool -> f a -> f a -> f a
ifS (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Mark
m forall a. Eq a => a -> a -> Bool
== Mark
Keep)
                                     (forall a. a -> Tree a -> Tree a -> Tree a
Branch forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a
g forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Tree (Marked f a) -> f (Tree a)
go Tree (Marked f a)
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Tree (Marked f a) -> f (Tree a)
go Tree (Marked f a)
r)
                                     (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Tree a
Leaf)

-- | Change enough nodes currently marked as 'Drop' to 'Keep' to ensure at
-- least @n@ nodes are marked 'Keep'.
--
-- Precondition: any 'Drop' marks must have been propagated; see 'propagate'.
-- Postcondition: this property is preserved.
keepAtLeast :: Word -> Tree (Marked f a) -> Tree (Marked f a)
keepAtLeast :: forall (f :: * -> *) a.
Word -> Tree (Marked f a) -> Tree (Marked f a)
keepAtLeast = \Word
n Tree (Marked f a)
t ->
    let kept :: Word
kept = forall (t :: * -> *) (f :: * -> *) a.
Foldable t =>
t (Marked f a) -> Word
countKept Tree (Marked f a)
t
    in if Word
kept forall a. Ord a => a -> a -> Bool
>= Word
n
         then Tree (Marked f a)
t
         else forall s a. State s a -> s -> a
evalState (forall (f :: * -> *) a.
Tree (Marked f a) -> State Word (Tree (Marked f a))
go Tree (Marked f a)
t) (Word
n forall a. Num a => a -> a -> a
- Word
kept)
  where
    go :: Tree (Marked f a) -> State Word (Tree (Marked f a))
    go :: forall (f :: * -> *) a.
Tree (Marked f a) -> State Word (Tree (Marked f a))
go   Tree (Marked f a)
Leaf                         = forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Tree a
Leaf
    go   (Branch (Marked Mark
Keep f a
x) Tree (Marked f a)
l Tree (Marked f a)
r) = forall a. a -> Tree a -> Tree a -> Tree a
Branch (forall (f :: * -> *) a. Mark -> f a -> Marked f a
Marked Mark
Keep f a
x) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) a.
Tree (Marked f a) -> State Word (Tree (Marked f a))
go Tree (Marked f a)
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a.
Tree (Marked f a) -> State Word (Tree (Marked f a))
go Tree (Marked f a)
r
    go t :: Tree (Marked f a)
t@(Branch (Marked Mark
Drop f a
x) Tree (Marked f a)
l Tree (Marked f a)
r) = forall s (m :: * -> *). MonadState s m => m s
get forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
         Word
0 ->
           -- Nothing left to drop
           forall (m :: * -> *) a. Monad m => a -> m a
return Tree (Marked f a)
t
         Word
n | forall a. Tree a -> Word
size Tree (Marked f a)
t forall a. Ord a => a -> a -> Bool
<= Word
n -> do
          -- We can keep the entire subtree
          forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ Word
n forall a. Num a => a -> a -> a
- forall a. Tree a -> Word
size Tree (Marked f a)
t
          forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a. Mark -> f a -> Marked f a
Marked Mark
Keep forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a. Marked f a -> f a
unmark) Tree (Marked f a)
t
         Word
n ->  do
          -- We cannot delete the entire subtree. In order to preserve the
          -- "drop property", we /must/ mark this node as 'Keep'
          forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ Word
n forall a. Num a => a -> a -> a
- Word
1
          forall a. a -> Tree a -> Tree a -> Tree a
Branch (forall (f :: * -> *) a. Mark -> f a -> Marked f a
Marked Mark
Keep f a
x) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) a.
Tree (Marked f a) -> State Word (Tree (Marked f a))
go Tree (Marked f a)
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a.
Tree (Marked f a) -> State Word (Tree (Marked f a))
go Tree (Marked f a)
r

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

data Endpoint a = Inclusive a | Exclusive a
data Interval a = Interval (Endpoint a) (Endpoint a)

-- | Compute interval with inclusive bounds, without exceeding range
--
-- Returns 'Nothing' if the interval is empty, and @Just@ the inclusive
-- lower and upper bound otherwise.
inclusiveBounds :: forall a. (Ord a, Enum a) => Interval a -> Maybe (a, a)
inclusiveBounds :: forall a. (Ord a, Enum a) => Interval a -> Maybe (a, a)
inclusiveBounds = \(Interval Endpoint a
lo Endpoint a
hi) -> Endpoint a -> Endpoint a -> Maybe (a, a)
go Endpoint a
lo Endpoint a
hi
  where
    -- The inequality checks in @go@ justify the use of @pred@ or @succ@
    go :: Endpoint a -> Endpoint a -> Maybe (a, a)
    go :: Endpoint a -> Endpoint a -> Maybe (a, a)
go (Inclusive a
lo) (Inclusive a
hi)
      | a
lo forall a. Ord a => a -> a -> Bool
<= a
hi  = forall a. a -> Maybe a
Just (a
lo, a
hi)
      | Bool
otherwise = forall a. Maybe a
Nothing
    go (Exclusive a
lo) (Inclusive a
hi)
      | a
lo forall a. Ord a => a -> a -> Bool
< a
hi   = forall a. a -> Maybe a
Just (forall a. Enum a => a -> a
succ a
lo, a
hi)
      | Bool
otherwise = forall a. Maybe a
Nothing
    go (Inclusive a
lo) (Exclusive a
hi)
      | a
lo forall a. Ord a => a -> a -> Bool
< a
hi   = forall a. a -> Maybe a
Just (a
lo, forall a. Enum a => a -> a
pred a
hi)
      | Bool
otherwise = forall a. Maybe a
Nothing
    go (Exclusive a
lo) (Exclusive a
hi)
      | a
lo forall a. Ord a => a -> a -> Bool
< a
hi   = if forall a. Enum a => a -> a
succ a
lo forall a. Ord a => a -> a -> Bool
> forall a. Enum a => a -> a
pred a
hi
                      then forall a. Maybe a
Nothing
                      else forall a. a -> Maybe a
Just (forall a. Enum a => a -> a
succ a
lo, forall a. Enum a => a -> a
pred a
hi)
      | Bool
otherwise = forall a. Maybe a
Nothing


-- | 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' forall a. Ord a => a -> a -> Bool
< a
a    = forall a b. Ord a => a -> Tree (a, b) -> Maybe b
lookup a
a' Tree (a, b)
l
  | a
a' forall a. Ord a => a -> a -> Bool
> a
a    = forall a b. Ord a => a -> Tree (a, b) -> Maybe b
lookup a
a' Tree (a, b)
r
  | Bool
otherwise = forall a. a -> Maybe a
Just b
b
lookup a
_ Tree (a, b)
Leaf = forall a. Maybe a
Nothing

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

drawTree :: Tree String -> String
drawTree :: Tree String -> String
drawTree = Tree String -> String
Rose.drawTree 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           = forall a. a -> [Tree a] -> Tree a
Rose.Node String
"*" []
    conv (Branch String
x Tree String
l Tree String
r) = 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]