-- | Utilities for working with trees with marked elements
--
-- Intended for qualified import.
--
-- > import qualified Test.Falsify.Internal.Marked.Tree as MarkedTree
module Test.Falsify.Internal.Marked.Tree (
    propagate
  , apply
  , keepAtLeast
  ) where

import Prelude hiding (drop)

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

import Data.Falsify.Tree (Tree(..))
import Test.Falsify.Marked (Mark(..), Marked(..))

import qualified Test.Falsify.Marked as Marked
import qualified Data.Falsify.Tree   as Tree

{-------------------------------------------------------------------------------
  Utilities for working with marked trees
-------------------------------------------------------------------------------}

-- | Propagate 'Drop' marker down the tree
--
-- This is useful in conjunction with 'apply', 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 = Tree (Marked f a) -> Tree (Marked f a)
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                         = Tree (Marked f a)
forall a. Tree a
Leaf
    keep (Branch (Marked Mark
Keep f a
x) Tree (Marked f a)
l Tree (Marked f a)
r) = Marked f a
-> Tree (Marked f a) -> Tree (Marked f a) -> Tree (Marked f a)
forall a. a -> Tree a -> Tree a -> Tree a
Branch (Mark -> f a -> Marked f a
forall (f :: * -> *) a. Mark -> f a -> Marked f a
Marked Mark
Keep f a
x) (Tree (Marked f a) -> Tree (Marked f a)
forall (f :: * -> *) a. Tree (Marked f a) -> Tree (Marked f a)
keep Tree (Marked f a)
l) (Tree (Marked f a) -> Tree (Marked f a)
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) = Marked f a
-> Tree (Marked f a) -> Tree (Marked f a) -> Tree (Marked f a)
forall a. a -> Tree a -> Tree a -> Tree a
Branch (Mark -> f a -> Marked f a
forall (f :: * -> *) a. Mark -> f a -> Marked f a
Marked Mark
Drop f a
x) (Tree (Marked f a) -> Tree (Marked f a)
forall (f :: * -> *) a. Tree (Marked f a) -> Tree (Marked f a)
drop Tree (Marked f a)
l) (Tree (Marked f a) -> Tree (Marked f a)
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 = (Marked f a -> Marked f a)
-> Tree (Marked f a) -> Tree (Marked f a)
forall a b. (a -> b) -> Tree a -> Tree b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Marked f a -> Marked f a)
 -> Tree (Marked f a) -> Tree (Marked f a))
-> (Marked f a -> Marked f a)
-> Tree (Marked f a)
-> Tree (Marked f a)
forall a b. (a -> b) -> a -> b
$ \(Marked Mark
_ f a
x) -> Mark -> f a -> Marked f a
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.
apply :: forall f a. Selective f => Tree (Marked f a) -> f (Tree a)
apply :: forall (f :: * -> *) a.
Selective f =>
Tree (Marked f a) -> f (Tree a)
apply = 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                      = Tree a -> f (Tree a)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Tree a
forall a. Tree a
Leaf
    go (Branch (Marked Mark
m f a
g) Tree (Marked f a)
l Tree (Marked f a)
r) = f Bool -> f (Tree a) -> f (Tree a) -> f (Tree a)
forall (f :: * -> *) a. Selective f => f Bool -> f a -> f a -> f a
ifS (Bool -> f Bool
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> f Bool) -> Bool -> f Bool
forall a b. (a -> b) -> a -> b
$ Mark
m Mark -> Mark -> Bool
forall a. Eq a => a -> a -> Bool
== Mark
Keep)
                                     (a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
Branch (a -> Tree a -> Tree a -> Tree a)
-> f a -> f (Tree a -> Tree a -> Tree a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a
g f (Tree a -> Tree a -> Tree a)
-> f (Tree a) -> f (Tree a -> Tree a)
forall a b. f (a -> b) -> f a -> f b
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 f (Tree a -> Tree a) -> f (Tree a) -> f (Tree a)
forall a b. f (a -> b) -> f a -> f b
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)
                                     (Tree a -> f (Tree a)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Tree a
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 = Tree (Marked f a) -> Word
forall (t :: * -> *) (f :: * -> *) a.
Foldable t =>
t (Marked f a) -> Word
Marked.countKept Tree (Marked f a)
t
    in if Word
kept Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
>= Word
n
         then Tree (Marked f a)
t
         else State Word (Tree (Marked f a)) -> Word -> Tree (Marked f a)
forall s a. State s a -> s -> a
evalState (Tree (Marked f a) -> State Word (Tree (Marked f a))
forall (f :: * -> *) a.
Tree (Marked f a) -> State Word (Tree (Marked f a))
go Tree (Marked f a)
t) (Word
n Word -> Word -> Word
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                         = Tree (Marked f a) -> StateT Word Identity (Tree (Marked f a))
forall a. a -> StateT Word Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return Tree (Marked f a)
forall a. Tree a
Leaf
    go   (Branch (Marked Mark
Keep f a
x) Tree (Marked f a)
l Tree (Marked f a)
r) = Marked f a
-> Tree (Marked f a) -> Tree (Marked f a) -> Tree (Marked f a)
forall a. a -> Tree a -> Tree a -> Tree a
Branch (Mark -> f a -> Marked f a
forall (f :: * -> *) a. Mark -> f a -> Marked f a
Marked Mark
Keep f a
x) (Tree (Marked f a) -> Tree (Marked f a) -> Tree (Marked f a))
-> StateT Word Identity (Tree (Marked f a))
-> StateT Word Identity (Tree (Marked f a) -> Tree (Marked f a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Tree (Marked f a) -> StateT Word Identity (Tree (Marked f a))
forall (f :: * -> *) a.
Tree (Marked f a) -> State Word (Tree (Marked f a))
go Tree (Marked f a)
l StateT Word Identity (Tree (Marked f a) -> Tree (Marked f a))
-> StateT Word Identity (Tree (Marked f a))
-> StateT Word Identity (Tree (Marked f a))
forall a b.
StateT Word Identity (a -> b)
-> StateT Word Identity a -> StateT Word Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Tree (Marked f a) -> StateT Word Identity (Tree (Marked f a))
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) = StateT Word Identity Word
forall s (m :: * -> *). MonadState s m => m s
get StateT Word Identity Word
-> (Word -> StateT Word Identity (Tree (Marked f a)))
-> StateT Word Identity (Tree (Marked f a))
forall a b.
StateT Word Identity a
-> (a -> StateT Word Identity b) -> StateT Word Identity b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
         Word
0 ->
           -- Nothing left to drop
           Tree (Marked f a) -> StateT Word Identity (Tree (Marked f a))
forall a. a -> StateT Word Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return Tree (Marked f a)
t
         Word
n | Tree (Marked f a) -> Word
forall a. Tree a -> Word
Tree.size Tree (Marked f a)
t Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
<= Word
n -> do
          -- We can keep the entire subtree
          Word -> StateT Word Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Word -> StateT Word Identity ())
-> Word -> StateT Word Identity ()
forall a b. (a -> b) -> a -> b
$ Word
n Word -> Word -> Word
forall a. Num a => a -> a -> a
- Tree (Marked f a) -> Word
forall a. Tree a -> Word
Tree.size Tree (Marked f a)
t
          Tree (Marked f a) -> StateT Word Identity (Tree (Marked f a))
forall a. a -> StateT Word Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (Tree (Marked f a) -> StateT Word Identity (Tree (Marked f a)))
-> Tree (Marked f a) -> StateT Word Identity (Tree (Marked f a))
forall a b. (a -> b) -> a -> b
$ (Marked f a -> Marked f a)
-> Tree (Marked f a) -> Tree (Marked f a)
forall a b. (a -> b) -> Tree a -> Tree b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Mark -> f a -> Marked f a
forall (f :: * -> *) a. Mark -> f a -> Marked f a
Marked Mark
Keep (f a -> Marked f a)
-> (Marked f a -> f a) -> Marked f a -> Marked f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Marked f a -> f a
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'
          Word -> StateT Word Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Word -> StateT Word Identity ())
-> Word -> StateT Word Identity ()
forall a b. (a -> b) -> a -> b
$ Word
n Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
1
          Marked f a
-> Tree (Marked f a) -> Tree (Marked f a) -> Tree (Marked f a)
forall a. a -> Tree a -> Tree a -> Tree a
Branch (Mark -> f a -> Marked f a
forall (f :: * -> *) a. Mark -> f a -> Marked f a
Marked Mark
Keep f a
x) (Tree (Marked f a) -> Tree (Marked f a) -> Tree (Marked f a))
-> StateT Word Identity (Tree (Marked f a))
-> StateT Word Identity (Tree (Marked f a) -> Tree (Marked f a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Tree (Marked f a) -> StateT Word Identity (Tree (Marked f a))
forall (f :: * -> *) a.
Tree (Marked f a) -> State Word (Tree (Marked f a))
go Tree (Marked f a)
l StateT Word Identity (Tree (Marked f a) -> Tree (Marked f a))
-> StateT Word Identity (Tree (Marked f a))
-> StateT Word Identity (Tree (Marked f a))
forall a b.
StateT Word Identity (a -> b)
-> StateT Word Identity a -> StateT Word Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Tree (Marked f a) -> StateT Word Identity (Tree (Marked f a))
forall (f :: * -> *) a.
Tree (Marked f a) -> State Word (Tree (Marked f a))
go Tree (Marked f a)
r