-- | Sample tree
--
-- Intended for qualified import.
--
-- import Test.Falsify.Internal.SampleTree (SampleTree(..))
-- import qualified Test.Falsify.Internal.SampleTree as SampleTree
module Test.Falsify.Internal.SampleTree (
    -- * Definition
    SampleTree(..)
  , Sample(..)
  , pattern Inf
  , sampleValue
    -- * Lenses
  , next
  , left
  , right
    -- * Construction
  , fromPRNG
  , fromSeed
  , minimal
  , constant
    -- * Combinators
  , map
  , mod
  ) where

import Prelude hiding (map, mod)
import qualified Prelude

import Data.Word
import Optics.Core (Lens')
import System.Random.SplitMix

import qualified Optics.Core as Optics

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

-- | Sample tree
--
-- A sample tree is a (conceptually and sometimes actually) infinite tree
-- representing drawing values from and splitting a PRNG.
data SampleTree =
    -- | Default constructor
    --
    -- The type of ST is really
    --
    -- > ST :: Word64 & (SampleTree * SampleTree) -> SampleTree
    --
    -- where `(&)` is the additive conjunction from linear logic. In other
    -- words, the intention is that /either/ the @Word64@ is used, /or/
    -- the pair of subtrees; put another way, we /either/ draw a value from the
    -- PRNG, /or/ split it into two new PRNGs. See 'next' and 'split'.
    SampleTree Sample SampleTree SampleTree

    -- | Minimal tree (0 everywhere)
    --
    -- This constructor allows us to represent an infinite tree in a finite way
    -- and, importantly, /recognize/ a tree that is minimal everywhere. This is
    -- necessary when shrinking in the context of generators that generate
    -- infinitely large values.
  | Minimal
  deriving (Int -> SampleTree -> ShowS
[SampleTree] -> ShowS
SampleTree -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [SampleTree] -> ShowS
$cshowList :: [SampleTree] -> ShowS
show :: SampleTree -> String
$cshow :: SampleTree -> String
showsPrec :: Int -> SampleTree -> ShowS
$cshowsPrec :: Int -> SampleTree -> ShowS
Show)

{-------------------------------------------------------------------------------
  Samples
-------------------------------------------------------------------------------}

-- | Sample
--
-- The samples in the 'SampleTree' record if they were the originally produced
-- sample, or whether they have been shrunk.
data Sample =
    NotShrunk Word64
  | Shrunk    Word64
  deriving (Int -> Sample -> ShowS
[Sample] -> ShowS
Sample -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Sample] -> ShowS
$cshowList :: [Sample] -> ShowS
show :: Sample -> String
$cshow :: Sample -> String
showsPrec :: Int -> Sample -> ShowS
$cshowsPrec :: Int -> Sample -> ShowS
Show, Sample -> Sample -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Sample -> Sample -> Bool
$c/= :: Sample -> Sample -> Bool
== :: Sample -> Sample -> Bool
$c== :: Sample -> Sample -> Bool
Eq, Eq Sample
Sample -> Sample -> Bool
Sample -> Sample -> Ordering
Sample -> Sample -> Sample
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Sample -> Sample -> Sample
$cmin :: Sample -> Sample -> Sample
max :: Sample -> Sample -> Sample
$cmax :: Sample -> Sample -> Sample
>= :: Sample -> Sample -> Bool
$c>= :: Sample -> Sample -> Bool
> :: Sample -> Sample -> Bool
$c> :: Sample -> Sample -> Bool
<= :: Sample -> Sample -> Bool
$c<= :: Sample -> Sample -> Bool
< :: Sample -> Sample -> Bool
$c< :: Sample -> Sample -> Bool
compare :: Sample -> Sample -> Ordering
$ccompare :: Sample -> Sample -> Ordering
Ord)

sampleValue :: Sample -> Word64
sampleValue :: Sample -> Word64
sampleValue (NotShrunk Word64
s) = Word64
s
sampleValue (Shrunk    Word64
s) = Word64
s

{-------------------------------------------------------------------------------
  Views
-------------------------------------------------------------------------------}

view :: SampleTree -> (Sample, SampleTree, SampleTree)
view :: SampleTree -> (Sample, SampleTree, SampleTree)
view SampleTree
Minimal            = (Word64 -> Sample
Shrunk Word64
0, SampleTree
Minimal, SampleTree
Minimal)
view (SampleTree Sample
s SampleTree
l SampleTree
r) = (Sample
s, SampleTree
l, SampleTree
r)

-- | Pattern synonym for treating the sample tree as infinite
pattern Inf :: Sample -> SampleTree -> SampleTree -> SampleTree
pattern $mInf :: forall {r}.
SampleTree
-> (Sample -> SampleTree -> SampleTree -> r) -> ((# #) -> r) -> r
Inf s l r <- (view -> (s, l, r))

{-# COMPLETE Inf #-}

{-------------------------------------------------------------------------------
  Lenses

  NOTE: The setter part of these lenses leaves 'Minimal' sample tree unchanged.
-------------------------------------------------------------------------------}

next :: Lens' SampleTree Sample
next :: Lens' SampleTree Sample
next = forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
Optics.lens SampleTree -> Sample
getter SampleTree -> Sample -> SampleTree
setter
  where
    getter :: SampleTree -> Sample
    getter :: SampleTree -> Sample
getter (Inf Sample
s SampleTree
_ SampleTree
_) = Sample
s

    setter :: SampleTree -> Sample -> SampleTree
    setter :: SampleTree -> Sample -> SampleTree
setter SampleTree
Minimal Sample
_            = SampleTree
Minimal
    setter (SampleTree Sample
_ SampleTree
l SampleTree
r) Sample
s = Sample -> SampleTree -> SampleTree -> SampleTree
SampleTree Sample
s SampleTree
l SampleTree
r

left :: Lens' SampleTree SampleTree
left :: Lens' SampleTree SampleTree
left = forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
Optics.lens SampleTree -> SampleTree
getter SampleTree -> SampleTree -> SampleTree
setter
  where
    getter :: SampleTree -> SampleTree
    getter :: SampleTree -> SampleTree
getter (Inf Sample
_ SampleTree
l SampleTree
_) = SampleTree
l

    setter :: SampleTree -> SampleTree -> SampleTree
    setter :: SampleTree -> SampleTree -> SampleTree
setter SampleTree
Minimal            SampleTree
_ = SampleTree
Minimal
    setter (SampleTree Sample
s SampleTree
_ SampleTree
r) SampleTree
l = Sample -> SampleTree -> SampleTree -> SampleTree
SampleTree Sample
s SampleTree
l SampleTree
r

right :: Lens' SampleTree SampleTree
right :: Lens' SampleTree SampleTree
right = forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
Optics.lens SampleTree -> SampleTree
getter SampleTree -> SampleTree -> SampleTree
setter
  where
    getter :: SampleTree -> SampleTree
    getter :: SampleTree -> SampleTree
getter (Inf Sample
_ SampleTree
_ SampleTree
r) = SampleTree
r

    setter :: SampleTree -> SampleTree -> SampleTree
    setter :: SampleTree -> SampleTree -> SampleTree
setter SampleTree
Minimal            SampleTree
_ = SampleTree
Minimal
    setter (SampleTree Sample
s SampleTree
l SampleTree
_) SampleTree
r = Sample -> SampleTree -> SampleTree -> SampleTree
SampleTree Sample
s SampleTree
l SampleTree
r

{-------------------------------------------------------------------------------
  Construction
-------------------------------------------------------------------------------}

fromPRNG :: SMGen -> SampleTree
fromPRNG :: SMGen -> SampleTree
fromPRNG = SMGen -> SampleTree
go
  where
    go :: SMGen -> SampleTree
    go :: SMGen -> SampleTree
go SMGen
g =
        let (Word64
n, SMGen
_) = SMGen -> (Word64, SMGen)
nextWord64 SMGen
g
            (SMGen
l, SMGen
r) = SMGen -> (SMGen, SMGen)
splitSMGen SMGen
g
        in Sample -> SampleTree -> SampleTree -> SampleTree
SampleTree (Word64 -> Sample
NotShrunk Word64
n) (SMGen -> SampleTree
go SMGen
l) (SMGen -> SampleTree
go SMGen
r)

fromSeed :: Word64 -> SampleTree
fromSeed :: Word64 -> SampleTree
fromSeed = SMGen -> SampleTree
fromPRNG forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word64 -> SMGen
mkSMGen

-- | Minimal sample tree
--
-- Generators should produce the \"simplest\" value when given this tree,
-- for some suitable application-specific definition of \"simple\".
minimal :: SampleTree
minimal :: SampleTree
minimal = SampleTree
Minimal

-- | Sample tree that is the given value everywhere
--
-- This is primarily useful for debugging.
constant :: Word64 -> SampleTree
constant :: Word64 -> SampleTree
constant Word64
s = SampleTree
go
  where
    go :: SampleTree
    go :: SampleTree
go = Sample -> SampleTree -> SampleTree -> SampleTree
SampleTree (Word64 -> Sample
NotShrunk Word64
s) SampleTree
go SampleTree
go

{-------------------------------------------------------------------------------
  Combinators
-------------------------------------------------------------------------------}

-- | Map function over all random samples in the tree
--
-- Precondition: the function must preserve zeros:
--
-- > f 0 == 0
--
-- This means that we have
--
-- > map f M == M
--
-- This is primarily useful for debugging.
map :: (Word64 -> Word64) -> SampleTree -> SampleTree
map :: (Word64 -> Word64) -> SampleTree -> SampleTree
map Word64 -> Word64
f = SampleTree -> SampleTree
go
  where
    go :: SampleTree -> SampleTree
    go :: SampleTree -> SampleTree
go (SampleTree Sample
s SampleTree
l SampleTree
r) = Sample -> SampleTree -> SampleTree -> SampleTree
SampleTree (Sample -> Sample
mapSample Sample
s) (SampleTree -> SampleTree
go SampleTree
l) (SampleTree -> SampleTree
go SampleTree
r)
    go SampleTree
Minimal            = SampleTree
Minimal

    mapSample :: Sample -> Sample
    mapSample :: Sample -> Sample
mapSample (NotShrunk Word64
s) = Word64 -> Sample
NotShrunk (Word64 -> Word64
f Word64
s)
    mapSample (Shrunk    Word64
s) = Word64 -> Sample
Shrunk    (Word64 -> Word64
f Word64
s)

-- | Apply @mod m@ at every sample in the tree
--
-- This is primarily useful for debugging.
mod :: Word64 -> SampleTree -> SampleTree
mod :: Word64 -> SampleTree -> SampleTree
mod Word64
m = (Word64 -> Word64) -> SampleTree -> SampleTree
map (\Word64
s -> Word64
s forall a. Integral a => a -> a -> a
`Prelude.mod` Word64
m)