module Test.Falsify.Internal.Generator.Definition (
    -- * Definition
    Gen(..)
  , bindWithoutShortcut
    -- * 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.Integer (Bit(..), encIntegerEliasG)
import Test.Falsify.Internal.SampleTree (SampleTree(..), Sample (..), pattern Inf)
import Test.Falsify.Internal.Search

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: 'Gen' is /NOT/ an instance of '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 Gen either.
newtype Gen a = Gen { forall a. Gen a -> SampleTree -> (a, [SampleTree])
runGen :: SampleTree -> (a, [SampleTree]) }
  deriving stock (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
<$ :: forall a b. a -> Gen b -> Gen a
$c<$ :: forall a b. a -> Gen b -> Gen a
fmap :: forall a b. (a -> b) -> Gen a -> Gen b
$cfmap :: forall a b. (a -> b) -> Gen a -> Gen b
Functor)

instance Applicative Gen where
  pure :: forall a. a -> Gen a
pure a
x = forall a. (SampleTree -> (a, [SampleTree])) -> Gen a
Gen forall a b. (a -> b) -> a -> b
$ \SampleTree
_st -> (a
x, [])
  <*> :: forall a 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  = 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 = forall a. (SampleTree -> (a, [SampleTree])) -> Gen a
Gen forall a b. (a -> b) -> a -> b
$ \(Inf Sample
s SampleTree
l SampleTree
r) ->
      let (a
a, [SampleTree]
ls) = forall a. Gen a -> SampleTree -> (a, [SampleTree])
runGen Gen a
x SampleTree
l
          (b
b, [SampleTree]
rs) = 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 forall a. a -> [a] -> NonEmpty a
:| [SampleTree]
ls) (SampleTree
r 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 = forall a. (SampleTree -> (a, [SampleTree])) -> Gen a
Gen forall a b. (a -> b) -> a -> b
$ \(Inf Sample
s SampleTree
l SampleTree
r) -> do
      let (Either a b
ma, [SampleTree]
ls) = 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) = 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 forall a. a -> [a] -> NonEmpty a
:| [SampleTree]
ls) (SampleTree
r forall a. a -> [a] -> NonEmpty a
:| [SampleTree]
rs))
        Right b
b ->
          (b
b, Sample
-> NonEmpty SampleTree -> NonEmpty SampleTree -> [SampleTree]
combineShrunk Sample
s (SampleTree
l forall a. a -> [a] -> NonEmpty a
:| [SampleTree]
ls) (SampleTree
r 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 forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [
      [Sample -> SampleTree -> SampleTree -> SampleTree
SampleTree Sample
s SampleTree
l' SampleTree
r  | SampleTree
l' <- forall a. SampleTree -> [a] -> [a]
unlessMinimal SampleTree
l [SampleTree]
ls]
    , [Sample -> SampleTree -> SampleTree -> SampleTree
SampleTree Sample
s SampleTree
l  SampleTree
r' | SampleTree
r' <- 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 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 = forall a. (SampleTree -> (a, [SampleTree])) -> Gen a
Gen forall a b. (a -> b) -> a -> b
$ \(Inf Sample
s SampleTree
l SampleTree
r) ->
    let (a
a, [SampleTree]
ls) = forall a. Gen a -> SampleTree -> (a, [SampleTree])
runGen Gen a
x SampleTree
l
        (b
b, [SampleTree]
rs) = 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 forall a. a -> [a] -> NonEmpty a
:| [SampleTree]
ls) (SampleTree
r 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) = 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]
        ]

{-------------------------------------------------------------------------------
  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 forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
a -> 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 = forall a. (SampleTree -> (a, [SampleTree])) -> Gen a
Gen forall a b. (a -> b) -> a -> b
$ \SampleTree
st ->
    let (b
b, [SampleTree]
shrunk) = forall a. Gen a -> SampleTree -> (a, [SampleTree])
runGen Gen b
g (forall k (is :: IxList) s a.
Is k A_Getter =>
Optic' k is s a -> s -> a
Optics.view Lens' SampleTree SampleTree
lens SampleTree
st)
    in (b
b, forall a b. (a -> b) -> [a] -> [b]
map (\SampleTree
st' -> forall k (is :: IxList) s t a b.
Is k A_Setter =>
Optic k is s t a b -> b -> s -> t
Optics.set Lens' SampleTree SampleTree
lens SampleTree
st' SampleTree
st) [SampleTree]
shrunk)
  where
    lens :: Lens' SampleTree SampleTree
    lens :: Lens' SampleTree SampleTree
lens = [Bit] -> Lens' SampleTree SampleTree
computeLens (Integer -> [Bit]
encIntegerEliasG forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral a
a)

    computeLens :: [Bit] -> Lens' SampleTree SampleTree
    computeLens :: [Bit] -> Lens' SampleTree SampleTree
computeLens []       = 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 forall a. Iso' a a
Optics.simple
    computeLens (Bit
O : [Bit]
bs) = Lens' SampleTree SampleTree
SampleTree.left  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] -> Lens' SampleTree SampleTree
computeLens [Bit]
bs
    computeLens (Bit
I : [Bit]
bs) = Lens' SampleTree SampleTree
SampleTree.right 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] -> Lens' 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 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
      (Sample -> [Word64]) -> Gen Sample
primWith (Word64 -> [Word64]
binarySearch 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 = forall a. (SampleTree -> (a, [SampleTree])) -> Gen a
Gen 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) 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 =
    forall a. Ord a => a -> a -> a
min Word64
n forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sample -> Word64
SampleTree.sampleValue forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
      (Sample -> [Word64]) -> Gen Sample
primWith (Word64 -> [Word64]
completeSearch 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 = forall a. (a -> Bool) -> [a] -> [a]
takeWhile (forall a. Ord a => a -> a -> Bool
<= Word64
n) [Word64
0 .. 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 = forall a. (SampleTree -> (a, [SampleTree])) -> Gen a
Gen 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) = forall a. (SampleTree -> (a, [SampleTree])) -> Gen a
Gen forall a b. (a -> b) -> a -> b
$ forall a. (a, [SampleTree]) -> (a, [SampleTree])
aux 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, [])