-- | Compound generators
module Test.Falsify.Reexported.Generator.Compound (
    -- * Taking advantage of 'Selective'
    choose
  , oneof
    -- * Lists
  , list
  , elem
  , pick
  , pickBiased
    -- ** Shuffling
  , shuffle
  , permutation
    -- * Tweak test data distribution
  , frequency
    -- * Trees
    -- ** Binary trees
  , tree
  , bst
    -- ** Shrink trees
  , IsValidShrink(..)
  , ShrinkTree
  , path
  , pathAny
    -- * Auxiliary
  , shrinkToNothing
  , mark
  ) where

import Prelude hiding (either, elem)

import Control.Monad
import Control.Selective
import Data.Either (either)
import Data.List.NonEmpty (NonEmpty(..))
import Data.Maybe (catMaybes)
import Data.Void

import qualified Data.List.NonEmpty as NE
import qualified Data.Tree          as Rose

import Data.Falsify.List (Permutation)
import Data.Falsify.Marked
import Data.Falsify.Tree (Tree(..), Interval(..), Endpoint(..))
import Test.Falsify.Internal.Generator
import Test.Falsify.Internal.Generator.Shrinking (IsValidShrink(..))
import Test.Falsify.Range (Range)
import Test.Falsify.Reexported.Generator.Shrinking
import Test.Falsify.Reexported.Generator.Simple

import qualified Data.Falsify.List  as List
import qualified Data.Falsify.Tree  as Tree
import qualified Test.Falsify.Range as Range

{-------------------------------------------------------------------------------
  Taking advantage of 'Selective'
-------------------------------------------------------------------------------}

-- | Generate a value with one of two generators
--
-- Shrinks towards the first generator;the two generators can shrink
-- independently from each other.
--
-- === Background
--
-- In the remainder of this docstring we give some background to this function,
-- which may be useful for general understanding of the @falsify@ library.
--
-- The implementation takes advantage of the that 'Gen' is a selective functor
-- to ensure that the two generators can shrink independently: if the initial
-- value of the generator is some @y@ produced by the second generator, later
-- shrunk to some @y'@, then if the generator can shrink to @x@ at some point,
-- produced by the /first/ generator, then shrinking effectively "starts over":
-- the value of @x@ is independent of @y'@.
--
-- That is different from doing this:
--
-- > do b <- bool
-- >    if b then l else r
--
-- In this case, @l@ and @r@ will be generated from the /same/ sample tree,
-- and so cannot shrink independently.
--
-- It is /also/ different from
--
-- > do x <- l
-- >    y <- r
-- >    b <- bool
-- >    return $ if b then x else y
--
-- In this case, @l@ and @r@ are run against /different/ sample trees, like we
-- do here, /but/ in this case if the current value produced by the generator is
-- produced by the right generator, then the sample tree used for the left
-- generator will always shrink to 'Minimal' (this /must/ be possible because
-- we're not currently using it); this means that we would then only be able to
-- shrink to a value from the left generator if the /minimal/ value produced by
-- that generator happens to work.
--
-- To rephrase that last point: generating values that are not actually used
-- will lead to poor shrinking, since those values can always be shrunk to their
-- minimal value, independently from whatever property is being tested: the
-- shrinker does not know that the value is not being used. The correct way to
-- conditionally use a value is to use the selective interface, as we do here.
choose :: Gen a -> Gen a -> Gen a
choose :: forall a. Gen a -> Gen a -> Gen a
choose = forall (f :: * -> *) a. Selective f => f Bool -> f a -> f a -> f a
ifS (Bool -> Gen Bool
bool Bool
True)

-- | Generate a value with one of many generators
--
-- Uniformly selects a generator and shrinks towards the first one.
oneof :: NonEmpty (Gen a) -> Gen a
oneof :: forall a. NonEmpty (Gen a) -> Gen a
oneof NonEmpty (Gen a)
gens = forall a. [(Word, Gen a)] -> Gen a
frequency forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (Word
1,) forall a b. (a -> b) -> a -> b
$ forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Gen a)
gens

{-------------------------------------------------------------------------------
  Auxiliary: marking elements
-------------------------------------------------------------------------------}

-- | Start with @Just x@ for some @x@, then shrink to @Nothing@
shrinkToNothing :: Gen a -> Gen (Maybe a)
shrinkToNothing :: forall a. Gen a -> Gen (Maybe a)
shrinkToNothing Gen a
g = forall a. a -> a -> Gen a
firstThen forall a. a -> Maybe a
Just (forall a b. a -> b -> a
const forall a. Maybe a
Nothing) forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen a
g

-- | Mark an element, shrinking towards 'Drop'
--
-- This is similar to 'shrinkToNothing', except that 'Marked' still has a value
-- in the 'Drop' case: marks are merely hints, that we may or may not use.
mark :: Gen a -> Gen (Marked Gen a)
mark :: forall a. Gen a -> Gen (Marked Gen a)
mark Gen a
x = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall (f :: * -> *) a. Mark -> f a -> Marked f a
Marked Gen a
x forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. a -> a -> Gen a
firstThen Mark
Keep Mark
Drop

{-------------------------------------------------------------------------------
  Lists
-------------------------------------------------------------------------------}

-- | Generate list of specified length
--
-- Shrinking behaviour:
--
-- * The length of the list will shrink as specified by the given range.
-- * We can drop random elements from the list, but prefer to drop them
--   from near the /end/ of the list.
--
-- Note on shrinking predictability: in the case that the specified 'Range' has
-- an origin which is neither the lower bound nor the upper bound (and only in
-- that case), 'list' can have confusing shrinking behaviour. For example,
-- suppose we have a range @(0, 10)@ with origin 5. Then we could start by
-- generating an intermediate list of length of 10 and then subsequently drop 5
-- elements from that, resulting in an optimal list length. However, we can now
-- shrink that length from 10 to 2 (which is closer to 5, after all), but now we
-- only have 2 elements to work with, and hence the generated list will now drop
-- from 5 elements to 2. This is not necessarily a problem, because that length
-- 2 can now subsequently shrink further towards closer to the origin (5), but
-- nonetheless it might result in confusing intermediate shrinking steps.
list :: Range Word -> Gen a -> Gen [a]
list :: forall a. Range Word -> Gen a -> Gen [a]
list Range Word
len Gen a
gen = do
    -- We do /NOT/ mark this call to 'inRange' as 'withoutShrinking': it could
    -- shrink towards larger values, in which case we really need to generate
    -- more elements. This doesn't really have any downsides: it merely means
    -- that we would prefer to shrink towards a prefix of the list first, before
    -- we try to drop random other elements from the list.
    --
    -- If we have an expression such as @(,) <$> list .. <*> list@, the two
    -- lists will be shrunk independently from each other due to the branching
    -- point above them. Hence, it doesn't matter if first generator uses "fewer
    -- samples" as it shrinks.
    Word
n <- forall a. Range a -> Gen a
inRange Range Word
len

    -- Generate @n@ marks, indicating for each element if we want to keep that
    -- element or not, so that we can drop elements from the middle of the list.
    --
    -- Due to the left-biased nature of shrinking, this will shrink towards
    -- dropped elements (@False@ values) near the start, but we want them near
    -- the /end/, so we reverse the list.
    [Marked Gen a]
marks <- forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a. Word -> [Marked f a] -> [Marked f a]
List.keepAtLeast (forall a. Range a -> a
Range.origin Range Word
len) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a] -> [a]
reverse) forall a b. (a -> b) -> a -> b
$
               forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
n) forall a b. (a -> b) -> a -> b
$ forall a. Gen a -> Gen (Marked Gen a)
mark Gen a
gen

    -- Finally, generate the elements we want to keep
    forall a. [Maybe a] -> [a]
catMaybes forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Selective f) =>
t (Marked f a) -> f (t (Maybe a))
selectAllKept [Marked Gen a]
marks

-- | Choose random element
--
-- Shrinks towards earlier elements.
--
-- NOTE: Does not work on infinite lists (it computes the length of the list).
elem :: NonEmpty a -> Gen a
elem :: forall a. NonEmpty a -> Gen a
elem = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\([a]
_before, a
x, [a]
_after) -> a
x) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. NonEmpty a -> Gen ([a], a, [a])
pick

-- | Generalization of 'elem' that additionally returns the parts of the list
-- before and after the element
pick :: NonEmpty a -> Gen ([a], a, [a])
pick :: forall a. NonEmpty a -> Gen ([a], a, [a])
pick = \NonEmpty a
xs ->
    forall a. [a] -> [a] -> Int -> ([a], a, [a])
aux [] (forall a. NonEmpty a -> [a]
NE.toList NonEmpty a
xs) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
      forall a. Range a -> Gen a
inRange (forall a. (Integral a, FiniteBits a) => (a, a) -> Range a
Range.between (Int
0, forall (t :: * -> *) a. Foldable t => t a -> Int
length NonEmpty a
xs forall a. Num a => a -> a -> a
- Int
1))
  where
    aux :: [a] -> [a] -> Int -> ([a], a, [a])
    aux :: forall a. [a] -> [a] -> Int -> ([a], a, [a])
aux [a]
_    []     Int
_ = forall a. HasCallStack => [Char] -> a
error [Char]
"pick: impossible"
    aux [a]
prev (a
x:[a]
xs) Int
0 = (forall a. [a] -> [a]
reverse [a]
prev, a
x, [a]
xs)
    aux [a]
prev (a
x:[a]
xs) Int
i = forall a. [a] -> [a] -> Int -> ([a], a, [a])
aux (a
xforall a. a -> [a] -> [a]
:[a]
prev) [a]
xs (Int
i forall a. Num a => a -> a -> a
- Int
1)

-- | Choose random element from a list
--
-- This is different from 'elem': it avoids first computing the length of the
-- list, and is biased towards elements earlier in the list. The advantage is
-- that this works for infinite lists, too.
--
-- Also returns the elements from the list before and after the chosen element.
pickBiased :: NonEmpty a -> Gen ([a], a, [a])
pickBiased :: forall a. NonEmpty a -> Gen ([a], a, [a])
pickBiased = \NonEmpty a
xs -> forall a.
[NonEmpty a] -> NonEmpty (NonEmpty a) -> Gen ([a], a, [a])
pickChunk [] (forall a. Word -> NonEmpty a -> NonEmpty (NonEmpty a)
List.chunksOfNonEmpty Word
chunkSize NonEmpty a
xs)
  where
    chunkSize :: Word
    chunkSize :: Word
chunkSize = Word
1_000

    -- We want to avoid computing the length of the list, but equally we don't
    -- want to skew /too/ heavily towards the start of the list. Therefore we
    -- chunk the list (this is lazy), then flip a coin for each chunk, and once
    -- we find a chunk, do an unbiased choice within that chunk.
    pickChunk :: [NonEmpty a] -> NonEmpty (NonEmpty a) -> Gen ([a], a, [a])
    pickChunk :: forall a.
[NonEmpty a] -> NonEmpty (NonEmpty a) -> Gen ([a], a, [a])
pickChunk [NonEmpty a]
prev (NonEmpty a
chunk :| []) = do
        -- No choice left: we must generate use this chunk
        forall a.
[NonEmpty a] -> NonEmpty a -> [NonEmpty a] -> Gen ([a], a, [a])
withChunk [NonEmpty a]
prev NonEmpty a
chunk []
    pickChunk [NonEmpty a]
prev (NonEmpty a
chunk :| next :: [NonEmpty a]
next@(NonEmpty a
n:[NonEmpty a]
ns)) = do
        Bool
useChunk <- Bool -> Gen Bool
bool Bool
True
        if Bool
useChunk
          then forall a.
[NonEmpty a] -> NonEmpty a -> [NonEmpty a] -> Gen ([a], a, [a])
withChunk [NonEmpty a]
prev NonEmpty a
chunk [NonEmpty a]
next
          else forall a.
[NonEmpty a] -> NonEmpty (NonEmpty a) -> Gen ([a], a, [a])
pickChunk (NonEmpty a
chunkforall a. a -> [a] -> [a]
:[NonEmpty a]
prev) (NonEmpty a
n forall a. a -> [a] -> NonEmpty a
:| [NonEmpty a]
ns)

    withChunk :: [NonEmpty a] -> NonEmpty a -> [NonEmpty a] -> Gen ([a], a, [a])
    withChunk :: forall a.
[NonEmpty a] -> NonEmpty a -> [NonEmpty a] -> Gen ([a], a, [a])
withChunk [NonEmpty a]
prev NonEmpty a
chunk [NonEmpty a]
next = do
        ([a]
chunkBefore, a
chunkElem, [a]
chunkAfter) <- forall a. NonEmpty a -> Gen ([a], a, [a])
pick NonEmpty a
chunk
        forall (m :: * -> *) a. Monad m => a -> m a
return (
            forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall a b. (a -> b) -> a -> b
$ forall a. [a] -> [a]
reverse forall a b. (a -> b) -> a -> b
$ [a]
chunkBefore forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map forall a. NonEmpty a -> [a]
NE.toList [NonEmpty a]
prev
          , a
chunkElem
          , [a]
chunkAfter forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a. NonEmpty a -> [a]
NE.toList [NonEmpty a]
next
          )

{-------------------------------------------------------------------------------
  Tweak test data distribution
-------------------------------------------------------------------------------}

-- | Choose generator with the given frequency
--
-- For example,
--
-- > frequency [
-- >     (1, genA)
-- >   , (2, genB)
-- >   ]
--
-- will use @genA@ 1/3rd of the time, and @genB@ 2/3rds.
--
-- Shrinks towards generators earlier in the list; the generators themselves
-- are independent from each other (shrinking of @genB@ does not affect
-- shrinking of @genA@).
--
-- Precondition: there should at least one generator with non-zero frequency.
frequency :: forall a. [(Word, Gen a)] -> Gen a
frequency :: forall a. [(Word, Gen a)] -> Gen a
frequency [(Word, Gen a)]
gens =
    case forall a. (a -> Bool) -> [a] -> [a]
filter ((forall a. Eq a => a -> a -> Bool
/= Word
0) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(Word, (Gen a, Word))]
indexedGens of
      []    -> forall a. HasCallStack => [Char] -> a
error [Char]
"frequency: no generators with non-zero frequency"
      [(Word, (Gen a, Word))]
gens' -> do
        let r :: Range Word
            r :: Range Word
r = forall a. (Integral a, FiniteBits a) => (a, a) -> Range a
Range.between (Word
0, forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(Word, (Gen a, Word))]
gens') forall a. Num a => a -> a -> a
- Word
1)
        (Gen a
gen, Word
genIx) <- (\Word
i -> forall x. Word -> [(Word, x)] -> x
frequencyLookup Word
i [(Word, (Gen a, Word))]
gens') forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Range a -> Gen a
inRange Range Word
r
        forall a b. Integral a => a -> Gen b -> Gen b
perturb Word
genIx Gen a
gen
  where
    -- We need to be careful: we don't want to perturb the generator by the
    -- value generated by 'inRange', because many different values could
    -- correspond to the /same/ generator. Instead, we assign each generator its
    -- own index, and use that instead.
    indexedGens :: [(Word, (Gen a, Word))]
    indexedGens :: [(Word, (Gen a, Word))]
indexedGens = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(Word
f, Gen a
g) Word
i -> (Word
f, (Gen a
g, Word
i))) [(Word, Gen a)]
gens [Word
0..]

-- | Internal auxiliary to 'frequency'
frequencyLookup :: Word -> [(Word, x)] -> x
frequencyLookup :: forall x. Word -> [(Word, x)] -> x
frequencyLookup = \Word
i [(Word, x)]
xs ->
    case forall x. Word -> [(Word, x)] -> Maybe x
go Word
i [(Word, x)]
xs of
      Just x
x  -> x
x
      Maybe x
Nothing ->
        forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [
           [Char]
"frequencyLookup: index "
         , forall a. Show a => a -> [Char]
show Word
i
         , [Char]
" out of range of "
         , forall a. Show a => a -> [Char]
show (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(Word, x)]
xs)
         ]
  where
    go :: Word -> [(Word, x)] -> Maybe x
    go :: forall x. Word -> [(Word, x)] -> Maybe x
go Word
_ []       = forall a. Maybe a
Nothing
    go Word
i ((Word
n, x
x):[(Word, x)]
xs)
      | Word
i forall a. Ord a => a -> a -> Bool
< Word
n     = forall a. a -> Maybe a
Just x
x
      | Bool
otherwise = forall x. Word -> [(Word, x)] -> Maybe x
go (Word
i forall a. Num a => a -> a -> a
- Word
n) [(Word, x)]
xs

{-------------------------------------------------------------------------------
  Shuffling
-------------------------------------------------------------------------------}

-- | Shuffle list (construct a permutation)
--
-- Shrinking behaviour: 'shuffle' is defined in terms of 'permutation', which
-- provides some guarantees: it shrinks towards making changes near the /start/
-- of the list, and towards swapping /fewer/ elements of the list.
--
-- It is difficult to define precisely how this affects the resulting list, but
-- we /can/ say that if for a particular counter-example it suffices if two
-- lists are different in /one/ element, then the shuffled list will in fact
-- only be different in /one/ place from the original, and that one element will
-- have been swapped with an immediate neighbour.
shuffle :: [a] -> Gen [a]
shuffle :: forall a. [a] -> Gen [a]
shuffle [a]
xs =
    forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Permutation -> [a] -> [a]
List.applyPermutation [a]
xs forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
      Word -> Gen Permutation
permutation (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs)

-- | Generate permutation for a list of length @n@
--
-- This is essentially an implemention of Fisher-Yates, in that we generate a
-- series of swaps (i, j), with 1 <= i <= n - 1 and @0 <= j <= i@, except that
--
-- * We can shrink a choice of @i@ (towards 1).
-- * We can drop arbitrary swaps.
--
-- This ensures that we shrink towards making swaps nearer the /start/ of the
-- list, as well as towards /fewer/ swaps.
--
-- We make no attempt to make the permutation canonical; doing so makes it
-- extremely difficult to get predicable shrinking behaviour.
permutation :: Word -> Gen Permutation
permutation :: Word -> Gen Permutation
permutation Word
0 = forall (m :: * -> *) a. Monad m => a -> m a
return []
permutation Word
1 = forall (m :: * -> *) a. Monad m => a -> m a
return []
permutation Word
n = do
    [Marked Gen (Word, Word)]
swaps <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall a. Gen a -> Gen (Marked Gen a)
mark forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Gen (Word, Word)
genSwap) [Word
n forall a. Num a => a -> a -> a
- Word
1, Word
n forall a. Num a => a -> a -> a
- Word
2 .. Word
1]
    forall a. [Maybe a] -> [a]
catMaybes forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Selective f) =>
t (Marked f a) -> f (t (Maybe a))
selectAllKept [Marked Gen (Word, Word)]
swaps
  where
    genSwap :: Word -> Gen (Word, Word)
    genSwap :: Word -> Gen (Word, Word)
genSwap Word
i = do
        Word
i' <- forall a. Range a -> Gen a
inRange forall a b. (a -> b) -> a -> b
$ forall a. (Integral a, FiniteBits a) => (a, a) -> Range a
Range.between (Word
1, Word
i)
        Word
j  <- forall a. Range a -> Gen a
inRange forall a b. (a -> b) -> a -> b
$ forall a. (Integral a, FiniteBits a) => (a, a) -> Range a
Range.between (Word
i, Word
0)
        forall (m :: * -> *) a. Monad m => a -> m a
return (Word
i', forall a. Ord a => a -> a -> a
min Word
i' Word
j)

{-------------------------------------------------------------------------------
  Binary trees
-------------------------------------------------------------------------------}

-- | Generate binary tree
tree :: forall a. Range Word -> Gen a -> Gen (Tree a)
tree :: forall a. Range Word -> Gen a -> Gen (Tree a)
tree Range Word
size Gen a
gen = do
    Word
n <- forall a. Range a -> Gen a
inRange Range Word
size
    Tree (Marked Gen a)
t <- forall (f :: * -> *) a.
Word -> Tree (Marked f a) -> Tree (Marked f a)
Tree.keepAtLeast (forall a. Range a -> a
Range.origin Range Word
size) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a. Tree (Marked f a) -> Tree (Marked f a)
Tree.propagate forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Word -> Gen (Tree (Marked Gen a))
go Word
n
    forall (f :: * -> *) a.
Selective f =>
Tree (Marked f a) -> f (Tree a)
Tree.genKept Tree (Marked Gen a)
t
  where
    go :: Word -> Gen (Tree (Marked Gen a))
    go :: Word -> Gen (Tree (Marked Gen a))
go Word
0 = forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Tree a
Leaf
    go Word
n = do
        -- Generate element at the root
        Marked Gen a
x <- forall a. Gen a -> Gen (Marked Gen a)
mark Gen a
gen

        -- Choose how many elements to put in the left subtree
        --
        -- This ranges from none (right-biased) to all (left-biased), shrinking
        -- towards half the number of elements: hence, towards a balanced tree.
        Word
inLeft <- forall a. Range a -> Gen a
inRange forall a b. (a -> b) -> a -> b
$ forall a. (Integral a, FiniteBits a) => (a, a) -> a -> Range a
Range.withOrigin (Word
0, Word
n forall a. Num a => a -> a -> a
- Word
1) ((Word
n forall a. Num a => a -> a -> a
- Word
1) forall a. Integral a => a -> a -> a
`div` Word
2)
        let inRight :: Word
inRight = (Word
n forall a. Num a => a -> a -> a
- Word
1) forall a. Num a => a -> a -> a
- Word
inLeft
        forall a. a -> Tree a -> Tree a -> Tree a
Branch Marked Gen a
x forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Word -> Gen (Tree (Marked Gen a))
go Word
inLeft forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Word -> Gen (Tree (Marked Gen a))
go Word
inRight

-- | Construct binary search tree
--
-- Shrinks by replacing entire subtrees by the empty tree.
bst :: forall a b. Integral a => (a -> Gen b) -> Interval a -> Gen (Tree (a, b))
bst :: forall a b.
Integral a =>
(a -> Gen b) -> Interval a -> Gen (Tree (a, b))
bst a -> Gen b
gen = Interval a -> Gen (Tree a)
go forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (\a
a -> (a
a,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> Gen b
gen a
a)
  where
    go :: Interval a -> Gen (Tree a)
    go :: Interval a -> Gen (Tree a)
go Interval a
i =
        case forall a. (Ord a, Enum a) => Interval a -> Maybe (a, a)
Tree.inclusiveBounds Interval a
i of
          Maybe (a, a)
Nothing       -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Tree a
Leaf
          Just (a
lo, a
hi) -> forall a. a -> a -> Gen a
firstThen forall a. a -> a
id (forall a b. a -> b -> a
const forall a. Tree a
Leaf) forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> a -> a -> Gen (Tree a)
go' a
lo a
hi

    -- inclusive bounds, lo <= hi
    go' :: a -> a -> Gen (Tree a)
    go' :: a -> a -> Gen (Tree a)
go' a
lo a
hi = forall a. a -> Tree a -> Tree a -> Tree a
Branch a
mid
            forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Interval a -> Gen (Tree a)
go (forall a. Endpoint a -> Endpoint a -> Interval a
Interval (forall a. a -> Endpoint a
Inclusive a
lo) (forall a. a -> Endpoint a
Exclusive a
mid))
            forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Interval a -> Gen (Tree a)
go (forall a. Endpoint a -> Endpoint a -> Interval a
Interval (forall a. a -> Endpoint a
Exclusive a
mid) (forall a. a -> Endpoint a
Inclusive a
hi))
      where
        mid :: a
        mid :: a
mid = a
lo forall a. Num a => a -> a -> a
+ ((a
hi forall a. Num a => a -> a -> a
- a
lo) forall a. Integral a => a -> a -> a
`div` a
2)

{-------------------------------------------------------------------------------
  Shrink trees
-------------------------------------------------------------------------------}

type ShrinkTree = Rose.Tree

-- | Generate semi-random path through the tree
--
-- Will only construct paths that satisfy the given predicate (typically, a
-- property that is being tested).
--
-- Shrinks towards shorter paths, and towards paths that use subtrees that
-- appear earlier in the list of subtrees at any node in the tree.
--
-- See also 'pathAny'.
path :: forall a p n.
     (a -> IsValidShrink p n) -- ^ Predicate
  -> ShrinkTree a
  -> Gen (Either n (NonEmpty p))
path :: forall a p n.
(a -> IsValidShrink p n)
-> ShrinkTree a -> Gen (Either n (NonEmpty p))
path a -> IsValidShrink p n
validShrink = \(Rose.Node a
a [ShrinkTree a]
as) ->
    case a -> IsValidShrink p n
validShrink a
a of
      InvalidShrink n
n -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a b. a -> Either a b
Left n
n
      ValidShrink   p
p -> forall a b. b -> Either a b
Right forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> p -> [ShrinkTree a] -> Gen (NonEmpty p)
go p
p [ShrinkTree a]
as
  where
    -- We only want to pick a shrunk value that matches the predicate, but we
    -- potentially waste a /lot/ of work if we first evaluate the predicate for
    -- /all/ potential shrunk values and then choose. So, instead we choose
    -- first, evaluate the predicate, and if it fails, choose again.
    go :: p -> [Rose.Tree a] -> Gen (NonEmpty p)
    go :: p -> [ShrinkTree a] -> Gen (NonEmpty p)
go p
p []     = forall (f :: * -> *) a. Applicative f => a -> f a
pure (p
p forall a. a -> [a] -> NonEmpty a
:| [])
    go p
p (ShrinkTree a
a:[ShrinkTree a]
as) = do
        ([ShrinkTree a]
before, ShrinkTree a
a', [ShrinkTree a]
after) <- forall a. NonEmpty a -> Gen ([a], a, [a])
pickBiased (ShrinkTree a
a forall a. a -> [a] -> NonEmpty a
:| [ShrinkTree a]
as)

        case ShrinkTree a -> Maybe (p, [ShrinkTree a])
checkPred ShrinkTree a
a' of
          Maybe (p, [ShrinkTree a])
Nothing ->
            -- Not a valid shrink step. Pick a different one.
            p -> [ShrinkTree a] -> Gen (NonEmpty p)
go p
p ([ShrinkTree a]
before forall a. [a] -> [a] -> [a]
++ [ShrinkTree a]
after)
          Just (p
p', [ShrinkTree a]
as') ->
            -- Found a valid shrink step.
            --
            -- We only call @choose@ once we found a valid shrink step,
            -- otherwise we would skew very heavily towards shorter paths.
            forall a. Gen a -> Gen a -> Gen a
choose
              (forall (f :: * -> *) a. Applicative f => a -> f a
pure (p
p forall a. a -> [a] -> NonEmpty a
:| []))
              (forall a. a -> NonEmpty a -> NonEmpty a
NE.cons p
p forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> p -> [ShrinkTree a] -> Gen (NonEmpty p)
go p
p' [ShrinkTree a]
as')

    checkPred :: Rose.Tree a -> Maybe (p, [Rose.Tree a])
    checkPred :: ShrinkTree a -> Maybe (p, [ShrinkTree a])
checkPred (Rose.Node a
a [ShrinkTree a]
as) =
       case a -> IsValidShrink p n
validShrink a
a of
         InvalidShrink n
_ -> forall a. Maybe a
Nothing
         ValidShrink   p
b -> forall a. a -> Maybe a
Just (p
b, [ShrinkTree a]
as)

-- | Variation on 'path' without a predicate.
pathAny :: ShrinkTree a -> Gen (NonEmpty a)
pathAny :: forall a. ShrinkTree a -> Gen (NonEmpty a)
pathAny = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either forall a. Void -> a
absurd forall a. a -> a
id) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a p n.
(a -> IsValidShrink p n)
-> ShrinkTree a -> Gen (Either n (NonEmpty p))
path forall p n. p -> IsValidShrink p n
ValidShrink