module Test.Falsify.Internal.Generator.Shrinking (
    -- * User-specified shrinking
    shrinkToOneOf
  , firstThen
  , shrinkWith
    -- * Support for shrink trees
  , fromShrinkTree
  , toShrinkTree
  , toShrinkTreeWithContext
  ) where

import Prelude hiding (properFraction)

import Data.Word

import qualified Data.Tree as Rose

import Test.Falsify.Internal.Generator
import Test.Falsify.SampleTree (SampleTree(..), Sample(..))
import Test.Falsify.ShrinkTree (ShrinkTree(..))

import qualified Test.Falsify.Context    as Context
import qualified Test.Falsify.ShrinkTree as ShrinkTree

{-------------------------------------------------------------------------------
  Specialized shrinking behaviour
-------------------------------------------------------------------------------}

-- | Start with @x@, then shrink to one of the @xs@
--
-- Once shrunk, will not shrink again.
--
-- Minimal value is the first shrunk value, if it exists, and the original
-- otherwise.
shrinkToOneOf :: forall a. a -> [a] -> Gen a
shrinkToOneOf :: forall a. a -> [a] -> Gen a
shrinkToOneOf a
x [a]
xs =
    Sample -> a
aux (Sample -> a) -> Gen Sample -> Gen a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Sample -> [Word64]) -> Gen Sample
primWith Sample -> [Word64]
shrinker
  where
    aux :: Sample -> a
    aux :: Sample -> a
aux (NotShrunk Word64
_) = a
x
    aux (Shrunk    Word64
i) = Word64 -> [a] -> a
index Word64
i [a]
xs

    -- When we shrink, we will try a bunch of new sample trees; we must ensure
    -- that we can try /any/ of the possible shrunk values.
    --
    -- We use this to implement 'fromShrinkTree'. Here, we explore a rose tree
    -- of possibilities; at every level in the tree, once we make a choice,
    -- we should commit to that choice and not consider it over and over again.
    -- Thus, once shrunk, we should not shrink any further.
    shrinker :: Sample -> [Word64]
    shrinker :: Sample -> [Word64]
shrinker (Shrunk Word64
_)    = []
    shrinker (NotShrunk Word64
_) = (Word64 -> a -> Word64) -> [Word64] -> [a] -> [Word64]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Word64 -> a -> Word64
forall a b. a -> b -> a
const [Word64
0..] [a]
xs

    -- Index the list of possible shrunk values. This is a bit like @(!!)@ from
    -- the prelude, but with some edge cases.
    --
    -- - If the list is empty, we return the unshrunk value.
    -- - Otherwise, if the index exceeds the bounds, we return the last element.
    --
    -- These two special cases can arise in one of two circumstances:
    --
    -- - When we run the generator against the 'Minimal' tree. This will give us
    --   a @Shrunk 0@ value, independent of what the specified shrinking
    --   function does, and it is important that we produce the right value.
    -- - When the generator is run against a sample tree that was shrunk wrt to
    --   a /different/ generator. In this case the value could be anything;
    --   we return the final ("least preferred") element, and then rely on
    --   later shrinking to replace this with a more preferred element.
    index :: Word64 -> [a] -> a
    index :: Word64 -> [a] -> a
index Word64
_ []     = a
x
    index Word64
_ [a
y]    = a
y
    index Word64
0 (a
y:[a]
_)  = a
y
    index Word64
n (a
_:[a]
ys) = Word64 -> [a] -> a
index (Word64
n Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1) [a]
ys

-- | Generator that always produces @x@ as initial value, and shrinks to @y@
firstThen :: forall a. a -> a -> Gen a
firstThen :: forall a. a -> a -> Gen a
firstThen a
x a
y = a
x a -> [a] -> Gen a
forall a. a -> [a] -> Gen a
`shrinkToOneOf` [a
y]

-- | Shrink with provided shrinker
--
-- This provides compatibility with QuickCheck-style manual shrinking.
--
-- Defined in terms of 'fromShrinkTree'; see discussion there for some
-- notes on performance.
shrinkWith :: forall a. (a -> [a]) -> Gen a -> Gen a
shrinkWith :: forall a. (a -> [a]) -> Gen a -> Gen a
shrinkWith a -> [a]
f Gen a
gen = do
    -- It is critical that we do not apply normal shrinking of the 'SampleTree'
    -- here (not even to 'Minimal'). If we did, then the resulting shrink tree
    -- would change, and we would be unable to iteratively construct a path
    -- through the shrink tree.
    --
    -- Of course, it can still happen that the generator gets reapplied in a
    -- different context; we must take this case into account in
    -- 'shrinkToOneOf'.
    a
x <- Gen a -> Gen a
forall a. Gen a -> Gen a
withoutShrinking Gen a
gen
    ShrinkTree a -> Gen a
forall a. ShrinkTree a -> Gen a
fromShrinkTree (ShrinkTree a -> Gen a) -> ShrinkTree a -> Gen a
forall a b. (a -> b) -> a -> b
$ a -> (a -> [a]) -> ShrinkTree a
forall a. a -> (a -> [a]) -> ShrinkTree a
ShrinkTree.unfold a
x a -> [a]
f

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

-- | Construct generator from shrink tree
--
-- This provides compatibility with Hedgehog-style integrated shrinking.
--
-- This is O(n^2) in the number of shrink steps: as this shrinks, the generator
-- is growing a path of indices which locates a particular value in the shrink
-- tree (resulting from unfolding the provided shrinking function). At each
-- step during the shrinking process the shrink tree is re-evaluated and the
-- next value in the tree is located; since this path throws linearly, the
-- overall cost is O(n^2).
--
-- The O(n^2) cost is only incurred on /locating/ the next element to be tested;
-- the property is not reevaluated at already-shrunk values.
fromShrinkTree :: forall a. ShrinkTree a -> Gen a
fromShrinkTree :: forall a. ShrinkTree a -> Gen a
fromShrinkTree = Tree a -> Gen a
go (Tree a -> Gen a)
-> (ShrinkTree a -> Tree a) -> ShrinkTree a -> Gen a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShrinkTree a -> Tree a
forall a. ShrinkTree a -> Tree a
unwrapShrinkTree
  where
    go :: Rose.Tree a -> Gen a
    go :: Tree a -> Gen a
go (Rose.Node a
x [Tree a]
xs) = do
        Maybe (Tree a)
next <- Maybe (Tree a)
forall a. Maybe a
Nothing Maybe (Tree a) -> [Maybe (Tree a)] -> Gen (Maybe (Tree a))
forall a. a -> [a] -> Gen a
`shrinkToOneOf` (Tree a -> Maybe (Tree a)) -> [Tree a] -> [Maybe (Tree a)]
forall a b. (a -> b) -> [a] -> [b]
map Tree a -> Maybe (Tree a)
forall a. a -> Maybe a
Just [Tree a]
xs
        case Maybe (Tree a)
next of
          Maybe (Tree a)
Nothing -> a -> Gen a
forall a. a -> Gen a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x
          Just Tree a
x' -> Tree a -> Gen a
go Tree a
x'

-- | Expose the full shrink tree of a generator
--
-- The generator is passed 'Context.Initial' for the initial step, and then
-- 'Context.Shrinking' with the number of the shrink step after that. The
-- 'Test.Falsify.Context.Final' step is /not/ included.
--
-- This generator does not shrink.
toShrinkTree :: forall a. Gen a -> Gen (ShrinkTree a)
toShrinkTree :: forall a. Gen a -> Gen (ShrinkTree a)
toShrinkTree Gen a
gen = ((Execution, a) -> a) -> ShrinkTree (Execution, a) -> ShrinkTree a
forall a b. (a -> b) -> ShrinkTree a -> ShrinkTree b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Execution, a) -> a
forall a b. (a, b) -> b
snd (ShrinkTree (Execution, a) -> ShrinkTree a)
-> Gen (ShrinkTree (Execution, a)) -> Gen (ShrinkTree a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> (Execution -> Gen a) -> Gen (ShrinkTree (Execution, a))
forall a.
Bool -> (Execution -> Gen a) -> Gen (ShrinkTree (Execution, a))
toShrinkTreeWithContext Bool
False (Gen a -> Execution -> Gen a
forall a b. a -> b -> a
const Gen a
gen)

-- | Generalization of 'toShrinkTree'
toShrinkTreeWithContext :: forall a.
     Bool -- ^ Include 'Context.Final' step?
  -> (Context.Execution -> Gen a)
  -> Gen (ShrinkTree (Context.Execution, a))
toShrinkTreeWithContext :: forall a.
Bool -> (Execution -> Gen a) -> Gen (ShrinkTree (Execution, a))
toShrinkTreeWithContext Bool
includeFinal Execution -> Gen a
gen = do
    Seed
initSeed <- Execution -> SampleTree -> Seed
Seed Execution
Context.Initial (SampleTree -> Seed) -> Gen SampleTree -> Gen Seed
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen SampleTree
captureLocalTree
    ShrinkTree (Execution, a) -> Gen (ShrinkTree (Execution, a))
forall a. a -> Gen a
forall (m :: * -> *) a. Monad m => a -> m a
return (ShrinkTree (Execution, a) -> Gen (ShrinkTree (Execution, a)))
-> ShrinkTree (Execution, a) -> Gen (ShrinkTree (Execution, a))
forall a b. (a -> b) -> a -> b
$ Tree (Execution, a) -> ShrinkTree (Execution, a)
forall a. Tree a -> ShrinkTree a
WrapShrinkTree (Tree (Execution, a) -> ShrinkTree (Execution, a))
-> Tree (Execution, a) -> ShrinkTree (Execution, a)
forall a b. (a -> b) -> a -> b
$ (Seed -> ((Execution, a), [Seed])) -> Seed -> Tree (Execution, a)
forall b a. (b -> (a, [b])) -> b -> Tree a
Rose.unfoldTree Seed -> ((Execution, a), [Seed])
aux Seed
initSeed
  where
    aux :: Seed -> ((Context.Execution, a), [Seed])
    aux :: Seed -> ((Execution, a), [Seed])
aux seed :: Seed
seed@Seed{Execution
seedContext :: Execution
seedContext :: Seed -> Execution
seedContext, SampleTree
seedSampleTree :: SampleTree
seedSampleTree :: Seed -> SampleTree
seedSampleTree} =
        case Gen a -> SampleTree -> (a, [SampleTree])
forall a. Gen a -> SampleTree -> (a, [SampleTree])
runGen (Execution -> Gen a
gen Execution
seedContext) SampleTree
seedSampleTree of
          (a
a, [SampleTree]
shrunk) -> ((Execution
seedContext, a
a), Seed -> [SampleTree] -> [Seed]
nextSeed Seed
seed [SampleTree]
shrunk)

    nextSeed :: Seed -> [SampleTree] -> [Seed]
    nextSeed :: Seed -> [SampleTree] -> [Seed]
nextSeed Seed{Execution
seedContext :: Seed -> Execution
seedContext :: Execution
seedContext, SampleTree
seedSampleTree :: Seed -> SampleTree
seedSampleTree :: SampleTree
seedSampleTree} [SampleTree]
shrunk =
       case Execution
seedContext of
         Execution
Context.Initial ->
           case [SampleTree]
shrunk of
             [] | Bool
includeFinal -> [Execution -> SampleTree -> Seed
Seed (Word -> Execution
Context.Final Word
0) SampleTree
seedSampleTree]
             [SampleTree]
_  -> (SampleTree -> Seed) -> [SampleTree] -> [Seed]
forall a b. (a -> b) -> [a] -> [b]
map (Execution -> SampleTree -> Seed
Seed (Execution -> SampleTree -> Seed)
-> Execution -> SampleTree -> Seed
forall a b. (a -> b) -> a -> b
$ Word -> Execution
Context.Shrinking Word
0) [SampleTree]
shrunk
         Context.Shrinking Word
i ->
           case [SampleTree]
shrunk of
             [] | Bool
includeFinal -> [Execution -> SampleTree -> Seed
Seed (Word -> Execution
Context.Final (Word -> Execution) -> Word -> Execution
forall a b. (a -> b) -> a -> b
$ Word -> Word
forall a. Enum a => a -> a
succ Word
i) SampleTree
seedSampleTree]
             [SampleTree]
_  -> (SampleTree -> Seed) -> [SampleTree] -> [Seed]
forall a b. (a -> b) -> [a] -> [b]
map (Execution -> SampleTree -> Seed
Seed (Execution -> SampleTree -> Seed)
-> Execution -> SampleTree -> Seed
forall a b. (a -> b) -> a -> b
$ Word -> Execution
Context.Shrinking (Word -> Execution) -> Word -> Execution
forall a b. (a -> b) -> a -> b
$ Word -> Word
forall a. Enum a => a -> a
succ Word
i) [SampleTree]
shrunk
         Context.Final Word
_ ->
           []

-- | Internal: 'Rose.unfoldTree' seed, for constructing the shrink tree
data Seed = Seed{
      Seed -> Execution
seedContext    :: Context.Execution
    , Seed -> SampleTree
seedSampleTree :: SampleTree
    }