-- | Compound generators
module Test.Falsify.Internal.Generator.Compound (
    -- * Taking advantage of 'Control.Selective.Selective'
    choose
  , oneof
    -- * Lists
  , list
  , elem
  , pick
  , pickBiased
    -- ** Shuffling
  , shuffle
  , permutation
    -- * Tweak test data distribution
  , frequency
    -- * Trees
    -- ** Binary trees
  , tree
  , bst
    -- ** Shrink trees
  , IsValidShrink(..)
  , 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.Permutation (Permutation)
import Data.Falsify.Tree (Tree(..))
import Test.Falsify.Internal.Generator
import Test.Falsify.Internal.Generator.Shrinking
import Test.Falsify.Internal.Generator.Simple
import Test.Falsify.Internal.Range
import Test.Falsify.Internal.Shrinking (IsValidShrink(..))
import Test.Falsify.Marked (Mark(..), Marked(..))
import Test.Falsify.ShrinkTree (ShrinkTree(..))

import qualified Data.Falsify.Internal.List        as List
import qualified Data.Falsify.Permutation          as Permutation
import qualified Test.Falsify.Internal.Marked.Tree as MarkedTree
import qualified Test.Falsify.Marked               as Marked
import qualified Test.Falsify.Range                as Range

{-------------------------------------------------------------------------------
  Taking advantage of 'Control.Selective.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 t'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 'Test.Falsify.SampleTree.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 = Gen Bool -> Gen a -> Gen a -> Gen a
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 = [(Word, Gen a)] -> Gen a
forall a. [(Word, Gen a)] -> Gen a
frequency ([(Word, Gen a)] -> Gen a) -> [(Word, Gen a)] -> Gen a
forall a b. (a -> b) -> a -> b
$ (Gen a -> (Word, Gen a)) -> [Gen a] -> [(Word, Gen a)]
forall a b. (a -> b) -> [a] -> [b]
map (Word
1,) ([Gen a] -> [(Word, Gen a)]) -> [Gen a] -> [(Word, Gen a)]
forall a b. (a -> b) -> a -> b
$ NonEmpty (Gen a) -> [Gen a]
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 = (a -> Maybe a) -> (a -> Maybe a) -> Gen (a -> Maybe a)
forall a. a -> a -> Gen a
firstThen a -> Maybe a
forall a. a -> Maybe a
Just (Maybe a -> a -> Maybe a
forall a b. a -> b -> a
const Maybe a
forall a. Maybe a
Nothing) Gen (a -> Maybe a) -> Gen a -> Gen (Maybe a)
forall a b. Gen (a -> b) -> Gen a -> Gen b
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 t'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 = (Mark -> Gen a -> Marked Gen a) -> Gen a -> Mark -> Marked Gen a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Mark -> Gen a -> Marked Gen a
forall (f :: * -> *) a. Mark -> f a -> Marked f a
Marked Gen a
x (Mark -> Marked Gen a) -> Gen Mark -> Gen (Marked Gen a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mark -> Mark -> Gen Mark
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
--
-- The implementation of 'list' uses a combination of two principles to produce
-- a list of the desired length:
--
-- * We generate a random list /length/ in the specified 'Range', and produce an
--   initial length of that length
-- * We then /drop/ elements from the resulting list, whilst still respecting the
--   specified 'Range'.
--
-- This ensures that we will produce a list with a length that tends towards the
-- origin of the specified 'Range', but whilst still being able to drop elements
-- from anywhere within the list, rather than just shrinking towards a prefix
-- (or suffix) from the initial list.
--
-- 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), this combination can
-- have potentially 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 might now shrink
-- the /length/ from 10 to 2 (which is closer to 5, after all). Now we only have
-- 2 elements to work with, and hence the generated list will now drop from 5
-- elements to 2, even though we were already at the ideal list length. 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 <- Range Word -> Gen Word
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 <- ([Marked Gen a] -> [Marked Gen a])
-> Gen [Marked Gen a] -> Gen [Marked Gen a]
forall a b. (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Word -> [Marked Gen a] -> [Marked Gen a]
forall (f :: * -> *) a. Word -> [Marked f a] -> [Marked f a]
List.keepAtLeast (Range Word -> Word
forall a. Range a -> a
Range.origin Range Word
len) ([Marked Gen a] -> [Marked Gen a])
-> ([Marked Gen a] -> [Marked Gen a])
-> [Marked Gen a]
-> [Marked Gen a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Marked Gen a] -> [Marked Gen a]
forall a. [a] -> [a]
reverse) (Gen [Marked Gen a] -> Gen [Marked Gen a])
-> Gen [Marked Gen a] -> Gen [Marked Gen a]
forall a b. (a -> b) -> a -> b
$
               Int -> Gen (Marked Gen a) -> Gen [Marked Gen a]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
n) (Gen (Marked Gen a) -> Gen [Marked Gen a])
-> Gen (Marked Gen a) -> Gen [Marked Gen a]
forall a b. (a -> b) -> a -> b
$ Gen a -> Gen (Marked Gen a)
forall a. Gen a -> Gen (Marked Gen a)
mark Gen a
gen

    -- Finally, generate the elements we want to keep
    [Maybe a] -> [a]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe a] -> [a]) -> Gen [Maybe a] -> Gen [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Marked Gen a] -> Gen [Maybe a]
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Selective f) =>
t (Marked f a) -> f (t (Maybe a))
Marked.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 = (([a], a, [a]) -> a) -> Gen ([a], a, [a]) -> Gen a
forall a b. (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\([a]
_before, a
x, [a]
_after) -> a
x) (Gen ([a], a, [a]) -> Gen a)
-> (NonEmpty a -> Gen ([a], a, [a])) -> NonEmpty a -> Gen a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty a -> Gen ([a], a, [a])
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 ->
    [a] -> [a] -> Int -> ([a], a, [a])
forall a. [a] -> [a] -> Int -> ([a], a, [a])
aux [] (NonEmpty a -> [a]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty a
xs) (Int -> ([a], a, [a])) -> Gen Int -> Gen ([a], a, [a])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
      Range Int -> Gen Int
forall a. Range a -> Gen a
inRange ((Int, Int) -> Range Int
forall a. (Integral a, FiniteBits a) => (a, a) -> Range a
Range.inclusive (Int
0, NonEmpty a -> Int
forall a. NonEmpty a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length NonEmpty a
xs Int -> Int -> Int
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
_ = [Char] -> ([a], a, [a])
forall a. HasCallStack => [Char] -> a
error [Char]
"pick: impossible"
    aux [a]
prev (a
x:[a]
xs) Int
0 = ([a] -> [a]
forall a. [a] -> [a]
reverse [a]
prev, a
x, [a]
xs)
    aux [a]
prev (a
x:[a]
xs) Int
i = [a] -> [a] -> Int -> ([a], a, [a])
forall a. [a] -> [a] -> Int -> ([a], a, [a])
aux (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
prev) [a]
xs (Int
i Int -> Int -> Int
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 -> [NonEmpty a] -> NonEmpty (NonEmpty a) -> Gen ([a], a, [a])
forall a.
[NonEmpty a] -> NonEmpty (NonEmpty a) -> Gen ([a], a, [a])
pickChunk [] (Word -> NonEmpty a -> NonEmpty (NonEmpty a)
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
        [NonEmpty a] -> NonEmpty a -> [NonEmpty a] -> Gen ([a], a, [a])
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 [NonEmpty a] -> NonEmpty a -> [NonEmpty a] -> Gen ([a], a, [a])
forall a.
[NonEmpty a] -> NonEmpty a -> [NonEmpty a] -> Gen ([a], a, [a])
withChunk [NonEmpty a]
prev NonEmpty a
chunk [NonEmpty a]
next
          else [NonEmpty a] -> NonEmpty (NonEmpty a) -> Gen ([a], a, [a])
forall a.
[NonEmpty a] -> NonEmpty (NonEmpty a) -> Gen ([a], a, [a])
pickChunk (NonEmpty a
chunkNonEmpty a -> [NonEmpty a] -> [NonEmpty a]
forall a. a -> [a] -> [a]
:[NonEmpty a]
prev) (NonEmpty a
n NonEmpty a -> [NonEmpty a] -> NonEmpty (NonEmpty a)
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) <- NonEmpty a -> Gen ([a], a, [a])
forall a. NonEmpty a -> Gen ([a], a, [a])
pick NonEmpty a
chunk
        ([a], a, [a]) -> Gen ([a], a, [a])
forall a. a -> Gen a
forall (m :: * -> *) a. Monad m => a -> m a
return (
            [[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[a]] -> [a]) -> [[a]] -> [a]
forall a b. (a -> b) -> a -> b
$ [[a]] -> [[a]]
forall a. [a] -> [a]
reverse ([[a]] -> [[a]]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> a -> b
$ [a]
chunkBefore [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: (NonEmpty a -> [a]) -> [NonEmpty a] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map NonEmpty a -> [a]
forall a. NonEmpty a -> [a]
NE.toList [NonEmpty a]
prev
          , a
chunkElem
          , [a]
chunkAfter [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ (NonEmpty a -> [a]) -> [NonEmpty a] -> [a]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap NonEmpty a -> [a]
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 ((Word, (Gen a, Word)) -> Bool)
-> [(Word, (Gen a, Word))] -> [(Word, (Gen a, Word))]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
/= Word
0) (Word -> Bool)
-> ((Word, (Gen a, Word)) -> Word) -> (Word, (Gen a, Word)) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word, (Gen a, Word)) -> Word
forall a b. (a, b) -> a
fst) [(Word, (Gen a, Word))]
indexedGens of
      []    -> [Char] -> Gen a
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 = (Word, Word) -> Range Word
forall a. (Integral a, FiniteBits a) => (a, a) -> Range a
Range.inclusive (Word
0, [Word] -> Word
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (((Word, (Gen a, Word)) -> Word)
-> [(Word, (Gen a, Word))] -> [Word]
forall a b. (a -> b) -> [a] -> [b]
map (Word, (Gen a, Word)) -> Word
forall a b. (a, b) -> a
fst [(Word, (Gen a, Word))]
gens') Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
1)
        (Gen a
gen, Word
genIx) <- (\Word
i -> Word -> [(Word, (Gen a, Word))] -> (Gen a, Word)
forall x. Word -> [(Word, x)] -> x
frequencyLookup Word
i [(Word, (Gen a, Word))]
gens') (Word -> (Gen a, Word)) -> Gen Word -> Gen (Gen a, Word)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Range Word -> Gen Word
forall a. Range a -> Gen a
inRange Range Word
r
        Word -> Gen a -> Gen a
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 = ((Word, Gen a) -> Word -> (Word, (Gen a, Word)))
-> [(Word, Gen a)] -> [Word] -> [(Word, (Gen a, Word))]
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 Word -> [(Word, x)] -> Maybe x
forall x. Word -> [(Word, x)] -> Maybe x
go Word
i [(Word, x)]
xs of
      Just x
x  -> x
x
      Maybe x
Nothing ->
        [Char] -> x
forall a. HasCallStack => [Char] -> a
error ([Char] -> x) -> [Char] -> x
forall a b. (a -> b) -> a -> b
$ [[Char]] -> [Char]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [
           [Char]
"frequencyLookup: index "
         , Word -> [Char]
forall a. Show a => a -> [Char]
show Word
i
         , [Char]
" out of range of "
         , [Word] -> [Char]
forall a. Show a => a -> [Char]
show (((Word, x) -> Word) -> [(Word, x)] -> [Word]
forall a b. (a -> b) -> [a] -> [b]
map (Word, x) -> Word
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
_ []       = Maybe x
forall a. Maybe a
Nothing
    go Word
i ((Word
n, x
x):[(Word, x)]
xs)
      | Word
i Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
< Word
n     = x -> Maybe x
forall a. a -> Maybe a
Just x
x
      | Bool
otherwise = Word -> [(Word, x)] -> Maybe x
forall x. Word -> [(Word, x)] -> Maybe x
go (Word
i Word -> Word -> Word
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 =
    (Permutation -> [a] -> [a]) -> [a] -> Permutation -> [a]
forall a b c. (a -> b -> c) -> b -> a -> c
flip Permutation -> [a] -> [a]
forall a. Permutation -> [a] -> [a]
Permutation.apply [a]
xs (Permutation -> [a]) -> Gen Permutation -> Gen [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
      Word -> Gen Permutation
permutation (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word) -> Int -> Word
forall a b. (a -> b) -> a -> b
$ [a] -> Int
forall a. [a] -> Int
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 = Permutation -> Gen Permutation
forall a. a -> Gen a
forall (m :: * -> *) a. Monad m => a -> m a
return Permutation
Permutation.identity
permutation Word
1 = Permutation -> Gen Permutation
forall a. a -> Gen a
forall (m :: * -> *) a. Monad m => a -> m a
return Permutation
Permutation.identity
permutation Word
n = do
    [Marked Gen (Word, Word)]
swaps <- (Word -> Gen (Marked Gen (Word, Word)))
-> [Word] -> Gen [Marked Gen (Word, Word)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Gen (Word, Word) -> Gen (Marked Gen (Word, Word))
forall a. Gen a -> Gen (Marked Gen a)
mark (Gen (Word, Word) -> Gen (Marked Gen (Word, Word)))
-> (Word -> Gen (Word, Word))
-> Word
-> Gen (Marked Gen (Word, Word))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Gen (Word, Word)
genSwap) [Word
n Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
1, Word
n Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
2 .. Word
1]
    [(Word, Word)] -> Permutation
Permutation.fromSwaps ([(Word, Word)] -> Permutation)
-> ([Maybe (Word, Word)] -> [(Word, Word)])
-> [Maybe (Word, Word)]
-> Permutation
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe (Word, Word)] -> [(Word, Word)]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe (Word, Word)] -> Permutation)
-> Gen [Maybe (Word, Word)] -> Gen Permutation
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Marked Gen (Word, Word)] -> Gen [Maybe (Word, Word)]
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Selective f) =>
t (Marked f a) -> f (t (Maybe a))
Marked.selectAllKept [Marked Gen (Word, Word)]
swaps
  where
    genSwap :: Word -> Gen (Word, Word)
    genSwap :: Word -> Gen (Word, Word)
genSwap Word
i = do
        Word
i' <- Range Word -> Gen Word
forall a. Range a -> Gen a
inRange (Range Word -> Gen Word) -> Range Word -> Gen Word
forall a b. (a -> b) -> a -> b
$ (Word, Word) -> Range Word
forall a. (Integral a, FiniteBits a) => (a, a) -> Range a
Range.inclusive (Word
1, Word
i)
        Word
j  <- Range Word -> Gen Word
forall a. Range a -> Gen a
inRange (Range Word -> Gen Word) -> Range Word -> Gen Word
forall a b. (a -> b) -> a -> b
$ (Word, Word) -> Range Word
forall a. (Integral a, FiniteBits a) => (a, a) -> Range a
Range.inclusive (Word
i, Word
0)
        (Word, Word) -> Gen (Word, Word)
forall a. a -> Gen a
forall (m :: * -> *) a. Monad m => a -> m a
return (Word
i', Word -> Word -> Word
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 <- Range Word -> Gen Word
forall a. Range a -> Gen a
inRange Range Word
size
    Tree (Marked Gen a)
t <- Word -> Tree (Marked Gen a) -> Tree (Marked Gen a)
forall (f :: * -> *) a.
Word -> Tree (Marked f a) -> Tree (Marked f a)
MarkedTree.keepAtLeast (Range Word -> Word
forall a. Range a -> a
Range.origin Range Word
size) (Tree (Marked Gen a) -> Tree (Marked Gen a))
-> (Tree (Marked Gen a) -> Tree (Marked Gen a))
-> Tree (Marked Gen a)
-> Tree (Marked Gen a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree (Marked Gen a) -> Tree (Marked Gen a)
forall (f :: * -> *) a. Tree (Marked f a) -> Tree (Marked f a)
MarkedTree.propagate (Tree (Marked Gen a) -> Tree (Marked Gen a))
-> Gen (Tree (Marked Gen a)) -> Gen (Tree (Marked Gen a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
           Word -> Gen (Tree (Marked Gen a))
go Word
n
    Tree (Marked Gen a) -> Gen (Tree a)
forall (f :: * -> *) a.
Selective f =>
Tree (Marked f a) -> f (Tree a)
MarkedTree.apply Tree (Marked Gen a)
t
  where
    go :: Word -> Gen (Tree (Marked Gen a))
    go :: Word -> Gen (Tree (Marked Gen a))
go Word
0 = Tree (Marked Gen a) -> Gen (Tree (Marked Gen a))
forall a. a -> Gen a
forall (m :: * -> *) a. Monad m => a -> m a
return Tree (Marked Gen a)
forall a. Tree a
Leaf
    go Word
n = do
        -- Generate element at the root
        Marked Gen a
x <- Gen a -> Gen (Marked Gen a)
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 <- Range Word -> Gen Word
forall a. Range a -> Gen a
inRange (Range Word -> Gen Word) -> Range Word -> Gen Word
forall a b. (a -> b) -> a -> b
$ (Word, Word) -> Word -> Range Word
forall a. (Integral a, FiniteBits a) => (a, a) -> a -> Range a
Range.withOrigin (Word
0, Word
n Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
1) ((Word
n Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
1) Word -> Word -> Word
forall a. Integral a => a -> a -> a
`div` Word
2)
        let inRight :: Word
inRight = (Word
n Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
1) Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
inLeft
        Marked Gen a
-> Tree (Marked Gen a)
-> Tree (Marked Gen a)
-> Tree (Marked Gen a)
forall a. a -> Tree a -> Tree a -> Tree a
Branch Marked Gen a
x (Tree (Marked Gen a) -> Tree (Marked Gen a) -> Tree (Marked Gen a))
-> Gen (Tree (Marked Gen a))
-> Gen (Tree (Marked Gen a) -> Tree (Marked Gen a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Word -> Gen (Tree (Marked Gen a))
go Word
inLeft Gen (Tree (Marked Gen a) -> Tree (Marked Gen a))
-> Gen (Tree (Marked Gen a)) -> Gen (Tree (Marked Gen a))
forall a b. Gen (a -> b) -> Gen a -> Gen b
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) -- ^ Generate value given a key
  -> (a, a)       -- ^ Inclusive range for the keys in the tree
  -> Gen (Tree (a, b))
bst :: forall a b.
Integral a =>
(a -> Gen b) -> (a, a) -> Gen (Tree (a, b))
bst a -> Gen b
gen = (a, a) -> Gen (Tree a)
go ((a, a) -> Gen (Tree a))
-> (Tree a -> Gen (Tree (a, b))) -> (a, a) -> Gen (Tree (a, b))
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> (a -> Gen (a, b)) -> Tree a -> Gen (Tree (a, b))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Tree a -> f (Tree b)
traverse (\a
a -> (a
a,) (b -> (a, b)) -> Gen b -> Gen (a, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> Gen b
gen a
a)
  where
    go :: (a, a) -> Gen (Tree a)
    go :: (a, a) -> Gen (Tree a)
go (a
lo, a
hi)
      | a
lo a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
hi  = Tree a -> Gen (Tree a)
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Tree a -> Gen (Tree a)) -> Tree a -> Gen (Tree a)
forall a b. (a -> b) -> a -> b
$ a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
Branch a
lo Tree a
forall a. Tree a
Leaf Tree a
forall a. Tree a
Leaf
      | a
lo a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
hi   = Tree a -> Gen (Tree a)
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Tree a
forall a. Tree a
Leaf
      | Bool
otherwise = (Tree a -> Tree a) -> (Tree a -> Tree a) -> Gen (Tree a -> Tree a)
forall a. a -> a -> Gen a
firstThen Tree a -> Tree a
forall a. a -> a
id (Tree a -> Tree a -> Tree a
forall a b. a -> b -> a
const Tree a
forall a. Tree a
Leaf) Gen (Tree a -> Tree a) -> Gen (Tree a) -> Gen (Tree a)
forall a b. Gen (a -> b) -> Gen a -> Gen b
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 =
        a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
Branch a
mid
          (Tree a -> Tree a -> Tree a)
-> Gen (Tree a) -> Gen (Tree a -> Tree a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a, a) -> Gen (Tree a)
go (a
lo, a -> a
forall a. Enum a => a -> a
pred a
mid)
          Gen (Tree a -> Tree a) -> Gen (Tree a) -> Gen (Tree a)
forall a b. Gen (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (a, a) -> Gen (Tree a)
go (a -> a
forall a. Enum a => a -> a
succ a
mid, a
hi)
      where
        -- Go through 'Integer' to avoid overflow
        mid' :: Integer
        mid' :: Integer
mid' = a -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
lo Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ ((a -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
hi Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- a -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
lo) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
2)

        mid :: a
        mid :: a
mid = Integer -> a
forall a. Num a => Integer -> a
fromInteger Integer
mid'

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

-- | 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 -> Either n p) -- ^ Predicate
  -> ShrinkTree a
  -> Gen (Either n (NonEmpty p))
path :: forall a p n.
(a -> Either n p) -> ShrinkTree a -> Gen (Either n (NonEmpty p))
path a -> Either n p
validShrink = \(WrapShrinkTree (Rose.Node a
a [Tree a]
as)) ->
    case a -> Either n p
validShrink a
a of
      Left  n
n -> Either n (NonEmpty p) -> Gen (Either n (NonEmpty p))
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either n (NonEmpty p) -> Gen (Either n (NonEmpty p)))
-> Either n (NonEmpty p) -> Gen (Either n (NonEmpty p))
forall a b. (a -> b) -> a -> b
$ n -> Either n (NonEmpty p)
forall a b. a -> Either a b
Left n
n
      Right p
p -> NonEmpty p -> Either n (NonEmpty p)
forall a b. b -> Either a b
Right (NonEmpty p -> Either n (NonEmpty p))
-> Gen (NonEmpty p) -> Gen (Either n (NonEmpty p))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> p -> [Tree a] -> Gen (NonEmpty p)
go p
p [Tree 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 -> [Tree a] -> Gen (NonEmpty p)
go p
p []     = NonEmpty p -> Gen (NonEmpty p)
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (p
p p -> [p] -> NonEmpty p
forall a. a -> [a] -> NonEmpty a
:| [])
    go p
p (Tree a
a:[Tree a]
as) = do
        ([Tree a]
before, Tree a
a', [Tree a]
after) <- NonEmpty (Tree a) -> Gen ([Tree a], Tree a, [Tree a])
forall a. NonEmpty a -> Gen ([a], a, [a])
pickBiased (Tree a
a Tree a -> [Tree a] -> NonEmpty (Tree a)
forall a. a -> [a] -> NonEmpty a
:| [Tree a]
as)

        case Tree a -> Maybe (p, [Tree a])
checkPred Tree a
a' of
          Maybe (p, [Tree a])
Nothing ->
            -- Not a valid shrink step. Pick a different one.
            p -> [Tree a] -> Gen (NonEmpty p)
go p
p ([Tree a]
before [Tree a] -> [Tree a] -> [Tree a]
forall a. [a] -> [a] -> [a]
++ [Tree a]
after)
          Just (p
p', [Tree 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.
            Gen (NonEmpty p) -> Gen (NonEmpty p) -> Gen (NonEmpty p)
forall a. Gen a -> Gen a -> Gen a
choose
              (NonEmpty p -> Gen (NonEmpty p)
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (p
p p -> [p] -> NonEmpty p
forall a. a -> [a] -> NonEmpty a
:| []))
              (p -> NonEmpty p -> NonEmpty p
forall a. a -> NonEmpty a -> NonEmpty a
NE.cons p
p (NonEmpty p -> NonEmpty p) -> Gen (NonEmpty p) -> Gen (NonEmpty p)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> p -> [Tree a] -> Gen (NonEmpty p)
go p
p' [Tree a]
as')

    checkPred :: Rose.Tree a -> Maybe (p, [Rose.Tree a])
    checkPred :: Tree a -> Maybe (p, [Tree a])
checkPred (Rose.Node a
a [Tree a]
as) =
       case a -> Either n p
validShrink a
a of
         Left  n
_ -> Maybe (p, [Tree a])
forall a. Maybe a
Nothing
         Right p
b -> (p, [Tree a]) -> Maybe (p, [Tree a])
forall a. a -> Maybe a
Just (p
b, [Tree a]
as)

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