module Test.Falsify.Internal.Generator (
    -- * Definition
    Gen(..)
  , bindWithoutShortcut
  , minimalValue
    -- * Primitive generators
  , prim
  , primWith
  , exhaustive
  , captureLocalTree
    -- * Generator independence
  , bindIntegral
  , perturb
    -- * Combinators
  , withoutShrinking
  ) where

import Control.Monad
import Control.Selective
import Data.List.NonEmpty (NonEmpty((:|)))
import Data.Word
import Optics.Core (Lens', (%))

import qualified Optics.Core as Optics

import Data.Falsify.Internal.Integer (Bit(..), encIntegerEliasG)
import Test.Falsify.Internal.Search
import Test.Falsify.SampleTree (SampleTree(..), pattern Inf, Sample(..))

import qualified Test.Falsify.SampleTree          as SampleTree
import qualified Test.Falsify.Internal.SampleTree as SampleTree

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

-- | Generator of a random value
--
-- Generators can be combined through their 'Functor', 'Applicative' and 'Monad'
-- interfaces. The primitive generator is 'prim', but most users will probably
-- want to construct their generators using the predefined from
-- "Test.Falsify.Generator" as building blocks.
--
-- Generators support \"internal integrated shrinking\". Shrinking is
-- /integrated/ in the sense of Hedgehog, meaning that we don't write a separate
-- shrinker at all, but the shrink behaviour is implied by the generator. For
-- example, if you have a generator @genList@ for a list of numbers, then
--
-- > filter even <$> genList
--
-- will only generate even numbers, and that property is automatically preserved
-- during shrinking. Shrinking is /internal/ in the sense of Hypothesis, meaning
-- that unlike in Hedgehog, shrinking works correctly even in the context of
-- monadic bind. For example, if you do
--
-- > do n <- genListLength
-- >    replicateM n someOtherGen
--
-- then we can shrink @n@ and the results from @someOtherGen@ in any order (that
-- said, users may prefer to use the dedicated
-- 'Test.Falsify.Generator.Compound.list' generator for this purpose, which
-- improves on this in a few ways).
--
-- NOTE: t'Gen' is /NOT/ an instance of 'Control.Applicative.Alternative'; this
-- would not be compatible with the generation of infinite data structures. For
-- the same reason, we do not have a monad transformer version of t'Gen' either.
newtype Gen a = Gen { forall a. Gen a -> SampleTree -> (a, [SampleTree])
runGen :: SampleTree -> (a, [SampleTree]) }
  deriving stock ((forall a b. (a -> b) -> Gen a -> Gen b)
-> (forall a b. a -> Gen b -> Gen a) -> Functor Gen
forall a b. a -> Gen b -> Gen a
forall a b. (a -> b) -> Gen a -> Gen b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> Gen a -> Gen b
fmap :: forall a b. (a -> b) -> Gen a -> Gen b
$c<$ :: forall a b. a -> Gen b -> Gen a
<$ :: forall a b. a -> Gen b -> Gen a
Functor)

instance Applicative Gen where
  pure :: forall a. a -> Gen a
pure a
x = (SampleTree -> (a, [SampleTree])) -> Gen a
forall a. (SampleTree -> (a, [SampleTree])) -> Gen a
Gen ((SampleTree -> (a, [SampleTree])) -> Gen a)
-> (SampleTree -> (a, [SampleTree])) -> Gen a
forall a b. (a -> b) -> a -> b
$ \SampleTree
_st -> (a
x, [])
  <*> :: forall a b. Gen (a -> b) -> Gen a -> Gen b
(<*>)  = Gen (a -> b) -> Gen a -> Gen b
forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

instance Monad Gen where
  return :: forall a. a -> Gen a
return  = a -> Gen a
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  Gen a
x >>= :: forall a b. Gen a -> (a -> Gen b) -> Gen b
>>= a -> Gen b
f = (SampleTree -> (b, [SampleTree])) -> Gen b
forall a. (SampleTree -> (a, [SampleTree])) -> Gen a
Gen ((SampleTree -> (b, [SampleTree])) -> Gen b)
-> (SampleTree -> (b, [SampleTree])) -> Gen b
forall a b. (a -> b) -> a -> b
$ \(Inf Sample
s SampleTree
l SampleTree
r) ->
      let (a
a, [SampleTree]
ls) = Gen a -> SampleTree -> (a, [SampleTree])
forall a. Gen a -> SampleTree -> (a, [SampleTree])
runGen Gen a
x SampleTree
l
          (b
b, [SampleTree]
rs) = Gen b -> SampleTree -> (b, [SampleTree])
forall a. Gen a -> SampleTree -> (a, [SampleTree])
runGen (a -> Gen b
f a
a) SampleTree
r
      in (b
b, Sample
-> NonEmpty SampleTree -> NonEmpty SampleTree -> [SampleTree]
combineShrunk Sample
s (SampleTree
l SampleTree -> [SampleTree] -> NonEmpty SampleTree
forall a. a -> [a] -> NonEmpty a
:| [SampleTree]
ls) (SampleTree
r SampleTree -> [SampleTree] -> NonEmpty SampleTree
forall a. a -> [a] -> NonEmpty a
:| [SampleTree]
rs))

instance Selective Gen where
  select :: forall a b. Gen (Either a b) -> Gen (a -> b) -> Gen b
select Gen (Either a b)
e Gen (a -> b)
f = (SampleTree -> (b, [SampleTree])) -> Gen b
forall a. (SampleTree -> (a, [SampleTree])) -> Gen a
Gen ((SampleTree -> (b, [SampleTree])) -> Gen b)
-> (SampleTree -> (b, [SampleTree])) -> Gen b
forall a b. (a -> b) -> a -> b
$ \(Inf Sample
s SampleTree
l SampleTree
r) -> do
      let (Either a b
ma, [SampleTree]
ls) = Gen (Either a b) -> SampleTree -> (Either a b, [SampleTree])
forall a. Gen a -> SampleTree -> (a, [SampleTree])
runGen Gen (Either a b)
e SampleTree
l
      case Either a b
ma of
        Left a
a ->
          let (a -> b
f', [SampleTree]
rs) = Gen (a -> b) -> SampleTree -> (a -> b, [SampleTree])
forall a. Gen a -> SampleTree -> (a, [SampleTree])
runGen Gen (a -> b)
f SampleTree
r
          in (a -> b
f' a
a, Sample
-> NonEmpty SampleTree -> NonEmpty SampleTree -> [SampleTree]
combineShrunk Sample
s (SampleTree
l SampleTree -> [SampleTree] -> NonEmpty SampleTree
forall a. a -> [a] -> NonEmpty a
:| [SampleTree]
ls) (SampleTree
r SampleTree -> [SampleTree] -> NonEmpty SampleTree
forall a. a -> [a] -> NonEmpty a
:| [SampleTree]
rs))
        Right b
b ->
          (b
b, Sample
-> NonEmpty SampleTree -> NonEmpty SampleTree -> [SampleTree]
combineShrunk Sample
s (SampleTree
l SampleTree -> [SampleTree] -> NonEmpty SampleTree
forall a. a -> [a] -> NonEmpty a
:| [SampleTree]
ls) (SampleTree
r SampleTree -> [SampleTree] -> NonEmpty SampleTree
forall a. a -> [a] -> NonEmpty a
:| []))

-- | Combine shrunk left and right sample trees
--
-- This is an internal function only.
combineShrunk ::
     Sample
  -> NonEmpty SampleTree -- ^ Original and shrunk left  trees
  -> NonEmpty SampleTree -- ^ Original and shrunk right trees
  -> [SampleTree]
combineShrunk :: Sample
-> NonEmpty SampleTree -> NonEmpty SampleTree -> [SampleTree]
combineShrunk Sample
s (SampleTree
l :| [SampleTree]
ls) (SampleTree
r :| [SampleTree]
rs) = [SampleTree] -> [SampleTree]
shortcut ([SampleTree] -> [SampleTree]) -> [SampleTree] -> [SampleTree]
forall a b. (a -> b) -> a -> b
$ [[SampleTree]] -> [SampleTree]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [
      [Sample -> SampleTree -> SampleTree -> SampleTree
SampleTree Sample
s SampleTree
l' SampleTree
r  | SampleTree
l' <- SampleTree -> [SampleTree] -> [SampleTree]
forall a. SampleTree -> [a] -> [a]
unlessMinimal SampleTree
l [SampleTree]
ls]
    , [Sample -> SampleTree -> SampleTree -> SampleTree
SampleTree Sample
s SampleTree
l  SampleTree
r' | SampleTree
r' <- SampleTree -> [SampleTree] -> [SampleTree]
forall a. SampleTree -> [a] -> [a]
unlessMinimal SampleTree
r [SampleTree]
rs]
    ]
  where
    -- We must be careful not to force @ls@/@rs@ if the tree is already minimal.
    unlessMinimal :: SampleTree -> [a] -> [a]
    unlessMinimal :: forall a. SampleTree -> [a] -> [a]
unlessMinimal SampleTree
Minimal [a]
_  = []
    unlessMinimal SampleTree
_       [a]
xs = [a]
xs

    shortcut :: [SampleTree] -> [SampleTree]
    shortcut :: [SampleTree] -> [SampleTree]
shortcut [] = []
    shortcut [SampleTree]
ts = SampleTree
Minimal SampleTree -> [SampleTree] -> [SampleTree]
forall a. a -> [a] -> [a]
: [SampleTree]
ts

-- | Varation on @(>>=)@ that doesn't apply the shortcut to 'Minimal'
--
-- This function is primarily useful for debugging @falsify@ itself; users
-- will probably never need it.
bindWithoutShortcut :: Gen a -> (a -> Gen b) -> Gen b
bindWithoutShortcut :: forall a b. Gen a -> (a -> Gen b) -> Gen b
bindWithoutShortcut Gen a
x a -> Gen b
f = (SampleTree -> (b, [SampleTree])) -> Gen b
forall a. (SampleTree -> (a, [SampleTree])) -> Gen a
Gen ((SampleTree -> (b, [SampleTree])) -> Gen b)
-> (SampleTree -> (b, [SampleTree])) -> Gen b
forall a b. (a -> b) -> a -> b
$ \(Inf Sample
s SampleTree
l SampleTree
r) ->
    let (a
a, [SampleTree]
ls) = Gen a -> SampleTree -> (a, [SampleTree])
forall a. Gen a -> SampleTree -> (a, [SampleTree])
runGen Gen a
x SampleTree
l
        (b
b, [SampleTree]
rs) = Gen b -> SampleTree -> (b, [SampleTree])
forall a. Gen a -> SampleTree -> (a, [SampleTree])
runGen (a -> Gen b
f a
a) SampleTree
r
    in (b
b, Sample
-> NonEmpty SampleTree -> NonEmpty SampleTree -> [SampleTree]
combine Sample
s (SampleTree
l SampleTree -> [SampleTree] -> NonEmpty SampleTree
forall a. a -> [a] -> NonEmpty a
:| [SampleTree]
ls) (SampleTree
r SampleTree -> [SampleTree] -> NonEmpty SampleTree
forall a. a -> [a] -> NonEmpty a
:| [SampleTree]
rs))
  where
    -- Variation on 'combineShrunk' that doesn't apply the shortcut
    combine ::
         Sample
      -> NonEmpty SampleTree -- ^ Original and shrunk left  trees
      -> NonEmpty SampleTree -- ^ Original and shrunk right trees
      -> [SampleTree]
    combine :: Sample
-> NonEmpty SampleTree -> NonEmpty SampleTree -> [SampleTree]
combine Sample
s (SampleTree
l :| [SampleTree]
ls) (SampleTree
r :| [SampleTree]
rs) = [[SampleTree]] -> [SampleTree]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [
          [Sample -> SampleTree -> SampleTree -> SampleTree
SampleTree Sample
s SampleTree
l' SampleTree
r  | SampleTree
l' <- [SampleTree]
ls]
        , [Sample -> SampleTree -> SampleTree -> SampleTree
SampleTree Sample
s SampleTree
l  SampleTree
r' | SampleTree
r' <- [SampleTree]
rs]
        ]

-- | Get the value produced by the generator on the minimal sample tree.
--
-- Having @Gen a@ is a proof that @a@ is inhabited, so this function
-- gives access to a witness.
minimalValue :: Gen a -> a
minimalValue :: forall a. Gen a -> a
minimalValue Gen a
g = (a, [SampleTree]) -> a
forall a b. (a, b) -> a
fst (Gen a -> SampleTree -> (a, [SampleTree])
forall a. Gen a -> SampleTree -> (a, [SampleTree])
runGen Gen a
g SampleTree
Minimal)

{-------------------------------------------------------------------------------
  Generator independence
-------------------------------------------------------------------------------}

-- | Selective bind
--
-- Unlike monadic bind, the RHS is generated and shrunk completely independently
-- for each different value of @a@ produced by the LHS.
--
-- This is a generalization of 'bindS' to arbitrary integral values; it is also
-- much more efficient than 'bindS'.
--
-- NOTE: This is only one way to make a generator independent. See 'perturb'
-- for more primitive combinator.
bindIntegral :: Integral a => Gen a -> (a -> Gen b) -> Gen b
bindIntegral :: forall a b. Integral a => Gen a -> (a -> Gen b) -> Gen b
bindIntegral Gen a
x a -> Gen b
f = Gen a
x Gen a -> (a -> Gen b) -> Gen b
forall a b. Gen a -> (a -> Gen b) -> Gen b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
a -> a -> Gen b -> Gen b
forall a b. Integral a => a -> Gen b -> Gen b
perturb a
a (a -> Gen b
f a
a)

-- | Run generator on different part of the sample tree depending on @a@
perturb :: Integral a => a -> Gen b -> Gen b
perturb :: forall a b. Integral a => a -> Gen b -> Gen b
perturb a
a Gen b
g = (SampleTree -> (b, [SampleTree])) -> Gen b
forall a. (SampleTree -> (a, [SampleTree])) -> Gen a
Gen ((SampleTree -> (b, [SampleTree])) -> Gen b)
-> (SampleTree -> (b, [SampleTree])) -> Gen b
forall a b. (a -> b) -> a -> b
$ \SampleTree
st ->
    let (b
b, [SampleTree]
shrunk) = Gen b -> SampleTree -> (b, [SampleTree])
forall a. Gen a -> SampleTree -> (a, [SampleTree])
runGen Gen b
g (Optic' A_Lens NoIx SampleTree SampleTree
-> SampleTree -> SampleTree
forall k (is :: IxList) s a.
Is k A_Getter =>
Optic' k is s a -> s -> a
Optics.view Optic' A_Lens NoIx SampleTree SampleTree
lens SampleTree
st)
    in (b
b, (SampleTree -> SampleTree) -> [SampleTree] -> [SampleTree]
forall a b. (a -> b) -> [a] -> [b]
map (\SampleTree
st' -> Optic' A_Lens NoIx SampleTree SampleTree
-> SampleTree -> SampleTree -> SampleTree
forall k (is :: IxList) s t a b.
Is k A_Setter =>
Optic k is s t a b -> b -> s -> t
Optics.set Optic' A_Lens NoIx SampleTree SampleTree
lens SampleTree
st' SampleTree
st) [SampleTree]
shrunk)
  where
    lens :: Lens' SampleTree SampleTree
    lens :: Optic' A_Lens NoIx SampleTree SampleTree
lens = [Bit] -> Optic' A_Lens NoIx SampleTree SampleTree
computeLens (Integer -> [Bit]
encIntegerEliasG (Integer -> [Bit]) -> Integer -> [Bit]
forall a b. (a -> b) -> a -> b
$ a -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
a)

    computeLens :: [Bit] -> Lens' SampleTree SampleTree
    computeLens :: [Bit] -> Optic' A_Lens NoIx SampleTree SampleTree
computeLens []       = Optic An_Iso NoIx SampleTree SampleTree SampleTree SampleTree
-> Optic' A_Lens NoIx SampleTree SampleTree
forall destKind srcKind (is :: IxList) s t a b.
Is srcKind destKind =>
Optic srcKind is s t a b -> Optic destKind is s t a b
Optics.castOptic Optic An_Iso NoIx SampleTree SampleTree SampleTree SampleTree
forall a. Iso' a a
Optics.simple
    computeLens (Bit
O : [Bit]
bs) = Optic' A_Lens NoIx SampleTree SampleTree
SampleTree.left  Optic' A_Lens NoIx SampleTree SampleTree
-> Optic' A_Lens NoIx SampleTree SampleTree
-> Optic' A_Lens NoIx SampleTree SampleTree
forall k l m (is :: IxList) (js :: IxList) (ks :: IxList) s t u v a
       b.
(JoinKinds k l m, AppendIndices is js ks) =>
Optic k is s t u v -> Optic l js u v a b -> Optic m ks s t a b
% [Bit] -> Optic' A_Lens NoIx SampleTree SampleTree
computeLens [Bit]
bs
    computeLens (Bit
I : [Bit]
bs) = Optic' A_Lens NoIx SampleTree SampleTree
SampleTree.right Optic' A_Lens NoIx SampleTree SampleTree
-> Optic' A_Lens NoIx SampleTree SampleTree
-> Optic' A_Lens NoIx SampleTree SampleTree
forall k l m (is :: IxList) (js :: IxList) (ks :: IxList) s t u v a
       b.
(JoinKinds k l m, AppendIndices is js ks) =>
Optic k is s t u v -> Optic l js u v a b -> Optic m ks s t a b
% [Bit] -> Optic' A_Lens NoIx SampleTree SampleTree
computeLens [Bit]
bs

{-------------------------------------------------------------------------------
  Primitive generators
-------------------------------------------------------------------------------}

-- | Uniform selection of 'Word64', shrinking towards 0, using binary search
--
-- This is a primitive generator; most users will probably not want to use this
-- generator directly.
prim :: Gen Word64
prim :: Gen Word64
prim =
    Sample -> Word64
SampleTree.sampleValue (Sample -> Word64) -> Gen Sample -> Gen Word64
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
      (Sample -> [Word64]) -> Gen Sample
primWith (Word64 -> [Word64]
binarySearch (Word64 -> [Word64]) -> (Sample -> Word64) -> Sample -> [Word64]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sample -> Word64
SampleTree.sampleValue)

-- | Generalization of 'prim' that allows to override the shrink behaviour
--
-- This is only required in rare circumstances. Most users will probably never
-- need to use this generator.
primWith :: (Sample -> [Word64]) -> Gen Sample
primWith :: (Sample -> [Word64]) -> Gen Sample
primWith Sample -> [Word64]
f = (SampleTree -> (Sample, [SampleTree])) -> Gen Sample
forall a. (SampleTree -> (a, [SampleTree])) -> Gen a
Gen ((SampleTree -> (Sample, [SampleTree])) -> Gen Sample)
-> (SampleTree -> (Sample, [SampleTree])) -> Gen Sample
forall a b. (a -> b) -> a -> b
$ \(Inf Sample
s SampleTree
l SampleTree
r) -> (
      Sample
s
    , (\Word64
s' -> Sample -> SampleTree -> SampleTree -> SampleTree
SampleTree (Word64 -> Sample
Shrunk Word64
s') SampleTree
l SampleTree
r) (Word64 -> SampleTree) -> [Word64] -> [SampleTree]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sample -> [Word64]
f Sample
s
    )

-- | Generate arbitrary value @x <= n@
--
-- Unlike 'prim', 'exhaustive' does not execute binary search. Instead, /all/
-- smaller values are considered. This is potentially very expensive; the
-- primary use case for this generator is testing shrinking behaviour, where
-- binary search can lead to some unpredicatable results.
--
-- This does /NOT/ do uniform selection: for small @n@, the generator will with
-- overwhelming probability produce @n@ itself as initial value.
--
-- This is a primitive generator; most users will probably not want to use this
-- generator directly.
exhaustive :: Word64 -> Gen Word64
exhaustive :: Word64 -> Gen Word64
exhaustive Word64
n =
    Word64 -> Word64 -> Word64
forall a. Ord a => a -> a -> a
min Word64
n (Word64 -> Word64) -> (Sample -> Word64) -> Sample -> Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sample -> Word64
SampleTree.sampleValue (Sample -> Word64) -> Gen Sample -> Gen Word64
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
      (Sample -> [Word64]) -> Gen Sample
primWith (Word64 -> [Word64]
completeSearch (Word64 -> [Word64]) -> (Sample -> Word64) -> Sample -> [Word64]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sample -> Word64
SampleTree.sampleValue)
  where
    completeSearch :: Word64 -> [Word64]
    completeSearch :: Word64 -> [Word64]
completeSearch Word64
0 = []
    completeSearch Word64
x = (Word64 -> Bool) -> [Word64] -> [Word64]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word64
n) [Word64
0 .. Word64 -> Word64
forall a. Enum a => a -> a
pred Word64
x]

-- | Capture the local sample tree
--
-- This generator does not shrink.
captureLocalTree :: Gen SampleTree
captureLocalTree :: Gen SampleTree
captureLocalTree = (SampleTree -> (SampleTree, [SampleTree])) -> Gen SampleTree
forall a. (SampleTree -> (a, [SampleTree])) -> Gen a
Gen ((SampleTree -> (SampleTree, [SampleTree])) -> Gen SampleTree)
-> (SampleTree -> (SampleTree, [SampleTree])) -> Gen SampleTree
forall a b. (a -> b) -> a -> b
$ \SampleTree
st -> (SampleTree
st, [])

{-------------------------------------------------------------------------------
  Shrinking combinators
-------------------------------------------------------------------------------}

-- | Disable shrinking in the given generator
--
-- Due to the nature of internal shrinking, it is always possible that a
-- generator gets reapplied to samples that were shrunk wrt to a /different/
-- generator. In this sense, 'withoutShrinking' should be considered to be a
-- hint only.
--
-- This function is only occassionally necessary; most users will probably not
-- need to use it.
withoutShrinking :: Gen a -> Gen a
withoutShrinking :: forall a. Gen a -> Gen a
withoutShrinking (Gen SampleTree -> (a, [SampleTree])
g) = (SampleTree -> (a, [SampleTree])) -> Gen a
forall a. (SampleTree -> (a, [SampleTree])) -> Gen a
Gen ((SampleTree -> (a, [SampleTree])) -> Gen a)
-> (SampleTree -> (a, [SampleTree])) -> Gen a
forall a b. (a -> b) -> a -> b
$ (a, [SampleTree]) -> (a, [SampleTree])
forall a. (a, [SampleTree]) -> (a, [SampleTree])
aux ((a, [SampleTree]) -> (a, [SampleTree]))
-> (SampleTree -> (a, [SampleTree]))
-> SampleTree
-> (a, [SampleTree])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SampleTree -> (a, [SampleTree])
g
  where
    aux :: (a, [SampleTree]) -> (a, [SampleTree])
    aux :: forall a. (a, [SampleTree]) -> (a, [SampleTree])
aux (a
outcome, [SampleTree]
_) = (a
outcome, [])