{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Byron.Spec.Ledger.UTxO.Generators where

import Byron.Spec.Ledger.Core hiding (Range, range)
import Byron.Spec.Ledger.UTxO
import Control.Applicative (empty)
import Data.Bitraversable (bitraverse)
import qualified Data.Map.Strict as M
import Hedgehog (Gen, Property, Range, assert, forAll, property, (===))
import qualified Hedgehog.Gen as Gen
import Hedgehog.Internal.Gen (
  atLeast,
  ensure,
  mapGenT,
  runDiscardEffectT,
  toTree,
  toTreeMaybeT,
 )
import Hedgehog.Internal.Tree (NodeT (..), TreeT (..), treeValue)
import qualified Hedgehog.Internal.Tree as Tree
import qualified Hedgehog.Range as Range

--------------------------------------------------------------------------------
-- Initial TxOut generator
--------------------------------------------------------------------------------

-- | Generate a set of initial 'TxOut's from a set of 'Addr's
genInitialTxOuts :: [Addr] -> Gen [TxOut]
genInitialTxOuts :: [Addr] -> Gen [TxOut]
genInitialTxOuts =
  forall (m :: * -> *) a.
(MonadGen m, GenBase m ~ Identity) =>
(a -> Bool) -> m a -> m a
Gen.filter (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t a -> Bool
null)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> Gen b) -> [a] -> Gen [b]
genTraverseSubsequence (\Addr
a -> Addr -> Lovelace -> TxOut
TxOut Addr
a forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Lovelace
genLovelace)

genLovelace :: Gen Lovelace
genLovelace :: Gen Lovelace
genLovelace =
  Integer -> Lovelace
Lovelace forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). MonadGen m => Range Word32 -> m Word32
Gen.word32 (forall a. Integral a => a -> a -> a -> Range a
Range.linearFrom Word32
mid Word32
mn Word32
mx)
  where
    mn :: Word32
mn = Word32
1
    mx :: Word32
mx = Word32
10000
    mid :: Word32
mid = forall a b. (RealFrac a, Integral b) => a -> b
floor (forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32
mx forall a. Num a => a -> a -> a
- Word32
mn) forall a. Fractional a => a -> a -> a
/ Double
2 :: Double)

-- | Generate a subsequence of a list of values and traverse the subsequence
--   with a generator producer
genTraverseSubsequence :: (a -> Gen b) -> [a] -> Gen [b]
genTraverseSubsequence :: forall a b. (a -> Gen b) -> [a] -> Gen [b]
genTraverseSubsequence a -> Gen b
genA [a]
as =
  forall (m :: * -> *) a (n :: * -> *) b.
(TreeT (MaybeT m) a -> TreeT (MaybeT n) b) -> GenT m a -> GenT n b
mapGenT (forall (m :: * -> *) a. m (NodeT m a) -> TreeT m a
TreeT forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall (m :: * -> *) a. Monad m => [TreeT m a] -> m (NodeT m [a])
interleaveTreeT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. NodeT m a -> a
nodeValue forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. TreeT m a -> m (NodeT m a)
runTreeT) forall a b. (a -> b) -> a -> b
$ do
    [a]
sub <- forall (m :: * -> *) a. MonadGen m => [a] -> m [a]
Gen.subsequence [a]
as
    forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall (m :: * -> *) a.
MonadGen m =>
m a -> m (TreeT (MaybeT (GenBase m)) a)
toTreeMaybeT forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Gen b
genA) [a]
sub

-- | Generate a list using 'genTraverseSubsequence'
genList :: Range Int -> Gen a -> Gen [a]
genList :: forall a. Range Int -> Gen a -> Gen [a]
genList Range Int
range Gen a
gen = forall (m :: * -> *) a. MonadGen m => (Size -> m a) -> m a
Gen.sized forall a b. (a -> b) -> a -> b
$ \Size
gSize ->
  forall (m :: * -> *) a. MonadGen m => (a -> Bool) -> m a -> m a
ensure (forall a. Int -> [a] -> Bool
atLeast forall a b. (a -> b) -> a -> b
$ forall a. Ord a => Size -> Range a -> a
Range.lowerBound Size
gSize Range Int
range) forall a b. (a -> b) -> a -> b
$
    forall a b. (a -> Gen b) -> [a] -> Gen [b]
genTraverseSubsequence
      (forall a b. a -> b -> a
const Gen a
gen)
      (forall a. Int -> a -> [a]
replicate (forall a. Ord a => Size -> Range a -> a
Range.upperBound Size
gSize Range Int
range) ())

-- | Temporarily defined here until hedgehog exposes this function
interleaveTreeT :: Monad m => [TreeT m a] -> m (NodeT m [a])
interleaveTreeT :: forall (m :: * -> *) a. Monad m => [TreeT m a] -> m (NodeT m [a])
interleaveTreeT = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (m :: * -> *) a. Monad m => [NodeT m a] -> NodeT m [a]
Tree.interleave forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall (m :: * -> *) a. TreeT m a -> m (NodeT m a)
runTreeT

--------------------------------------------------------------------------------
-- Tx generator
--------------------------------------------------------------------------------

-- | Generate a valid transaction from a given 'UTxO'
genTxFromUTxO ::
  -- | List of addresses to choose from as recipients of the transaction
  -- outputs.
  [Addr] ->
  -- | Fee calculation
  (Tx -> Lovelace) ->
  -- | UTxO used to determine which unspent outputs can be used in the
  -- transaction.
  UTxO ->
  Gen Tx
genTxFromUTxO :: [Addr] -> (Tx -> Lovelace) -> UTxO -> Gen Tx
genTxFromUTxO [Addr]
addrs Tx -> Lovelace
txfee (UTxO Map TxIn TxOut
utxo) = do
  TxBody
txbody <- GenT Identity TxBody
genTxBody
  let wits :: [Wit]
wits = TxBody -> TxIn -> Wit
witnessForTxIn TxBody
txbody forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TxBody -> [TxIn]
inputs TxBody
txbody
  let tx :: Tx
tx = TxBody -> [Wit] -> Tx
Tx TxBody
txbody [Wit]
wits
  Gen Tx -> Gen Tx
subtractFees forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure Tx
tx
  where
    genTxBody :: GenT Identity TxBody
genTxBody =
      forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry [TxIn] -> [TxOut] -> TxBody
TxBody
        forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a.
(MonadGen m, GenBase m ~ Identity) =>
(a -> Bool) -> m a -> m a
Gen.filter
          (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t a -> Bool
null forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst)
          ( forall input output.
[input]
-> (input -> Integer)
-> Gen (Integer -> output)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> Gen ([input], [output])
genInputOutput
              (forall k a. Map k a -> [k]
M.keys Map TxIn TxOut
utxo)
              (forall b a. b -> (a -> b) -> Maybe a -> b
maybe Integer
0 (Lovelace -> Integer
unLovelace forall b c a. (b -> c) -> (a -> b) -> a -> c
. TxOut -> Lovelace
value) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Map TxIn TxOut
utxo)
              (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Lovelace
Lovelace) forall a b. (a -> b) -> a -> b
$ Addr -> Lovelace -> TxOut
TxOut forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) (m :: * -> *) a.
(Foldable f, MonadGen m) =>
f a -> m a
Gen.element [Addr]
addrs)
              (Lovelace -> Integer
unLovelace forall b c a. (b -> c) -> (a -> b) -> a -> c
. TxOut -> Lovelace
value)
              (\Integer -> Integer
f TxOut
out -> TxOut
out {value :: Lovelace
value = Integer -> Lovelace
Lovelace forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Integer
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lovelace -> Integer
unLovelace forall a b. (a -> b) -> a -> b
$ TxOut -> Lovelace
value TxOut
out})
          )

    witnessForTxIn :: TxBody -> TxIn -> Wit
    witnessForTxIn :: TxBody -> TxIn -> Wit
witnessForTxIn TxBody
tx TxIn
txin =
      case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup TxIn
txin Map TxIn TxOut
utxo of
        Just (TxOut (Addr VKey
pay) Lovelace
_) ->
          KeyPair -> TxBody -> Wit
witnessForTx (Owner -> KeyPair
keyPair forall a b. (a -> b) -> a -> b
$ forall a. HasOwner a => a -> Owner
owner VKey
pay) TxBody
tx
        Maybe TxOut
Nothing ->
          forall a. HasCallStack => [Char] -> a
error [Char]
"The generators must ensure that we are spending unspent inputs"

    witnessForTx :: KeyPair -> TxBody -> Wit
    witnessForTx :: KeyPair -> TxBody -> Wit
witnessForTx (KeyPair SKey
sk VKey
vk) TxBody
tx = VKey -> Sig TxBody -> Wit
Wit VKey
vk (forall a. SKey -> a -> Sig a
sign SKey
sk TxBody
tx)
    subtractFees :: Gen Tx -> Gen Tx
    subtractFees :: Gen Tx -> Gen Tx
subtractFees =
      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Tx -> Tx
subtractFees'
        -- In Byron, we must disallow empty outputs in transactions in order to
        -- maintain compatability with the `cardano-sl` implementation.
        -- In order to do this, while also potentially removing some outputs from
        -- the transaction to ensure that the transaction fee is covered, we only
        -- generate transactions whose sum of outputs is greater than the
        -- transaction fee. This way, we ensure that there will always remain at
        -- least 1 'Lovelace' in the outputs.
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a.
(MonadGen m, GenBase m ~ Identity) =>
(a -> Bool) -> m a -> m a
Gen.filter (\tx :: Tx
tx@(Tx TxBody
txb [Wit]
_) -> forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (TxOut -> Lovelace
value forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TxBody -> [TxOut]
outputs TxBody
txb) forall a. Ord a => a -> a -> Bool
> Tx -> Lovelace
txfee Tx
tx)
      where
        subtractFees' :: Tx -> Tx
subtractFees' tx :: Tx
tx@(Tx TxBody
txb [Wit]
_) =
          let newBody :: TxBody
newBody = TxBody
txb {outputs :: [TxOut]
outputs = forall n a.
(Num n, Ord n) =>
n -> (a -> n) -> ((n -> n) -> a -> a) -> [a] -> [a]
subFromList (Tx -> Lovelace
txfee Tx
tx) TxOut -> Lovelace
value (Lovelace -> Lovelace) -> TxOut -> TxOut
updateValue (TxBody -> [TxOut]
outputs TxBody
txb)}
           in TxBody -> [Wit] -> Tx
Tx TxBody
newBody (TxBody -> TxIn -> Wit
witnessForTxIn TxBody
newBody forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TxBody -> [TxIn]
inputs TxBody
newBody)

        updateValue :: (Lovelace -> Lovelace) -> TxOut -> TxOut
updateValue Lovelace -> Lovelace
f TxOut
out = TxOut
out {value :: Lovelace
value = Lovelace -> Lovelace
f (TxOut -> Lovelace
value TxOut
out)}

-- | A property to test that the entire shrink tree generated by
--   'genInputOutput' maintains the invariant that the inputs and outputs have
--   equal sums
--
--   NB: This uses small values for the list and values, because we force the
--   entire shrink tree, which can grow very large
propGenInputOutput :: Property
propGenInputOutput :: Property
propGenInputOutput = HasCallStack => PropertyT IO () -> Property
property forall a b. (a -> b) -> a -> b
$ do
  [Integer]
ins <- forall (m :: * -> *) a.
(Monad m, Show a, HasCallStack) =>
Gen a -> PropertyT m a
forAll forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MonadGen m => Range Int -> m a -> m [a]
Gen.list (forall a. Integral a => a -> a -> Range a
Range.linear Int
0 Int
5) (forall (m :: * -> *) a. (MonadGen m, Integral a) => Range a -> m a
Gen.integral (forall a. Integral a => a -> a -> Range a
Range.linear Integer
0 Integer
5))
  Tree ([Integer], [Integer])
insOutsTree <- forall (m :: * -> *) a.
(Monad m, Show a, HasCallStack) =>
Gen a -> PropertyT m a
forAll forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadGen m, GenBase m ~ Identity) =>
m a -> m (Tree a)
toTree (forall input output.
[input]
-> (input -> Integer)
-> Gen (Integer -> output)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> Gen ([input], [output])
genInputOutput [Integer]
ins forall a. a -> a
id (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. a -> a
id) forall a. a -> a
id forall a. a -> a
id)
  forall (m :: * -> *). (MonadTest m, HasCallStack) => Bool -> m ()
assert forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\([Integer]
ins', [Integer]
outs) -> forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Integer]
ins' forall a. Eq a => a -> a -> Bool
== forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Integer]
outs) Tree ([Integer], [Integer])
insOutsTree

-- | Generate a subsequence of the inputs and split the value into a number of
--   outputs
--
--   The shrink tree maintains the invariant that the sum of the inputs is equal
--   to the sum of the outputs. This is generalised to any type of input where
--   we can view the contained value, and type of output where we can generate
--   the type from a value, and view and modify the contained value
genInputOutput ::
  -- | List of input to take a subsequence of
  [input] ->
  -- | A view of the value in the input used to preserve the invariant
  (input -> Integer) ->
  -- | A applicative generator for an output given a value
  Gen (Integer -> output) ->
  -- | A view of the value in the output used to preserve the invariant
  (output -> Integer) ->
  -- | An update function for the output value used to preserve the invariant
  ((Integer -> Integer) -> output -> output) ->
  Gen ([input], [output])
genInputOutput :: forall input output.
[input]
-> (input -> Integer)
-> Gen (Integer -> output)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> Gen ([input], [output])
genInputOutput [input]
ins input -> Integer
inValue Gen (Integer -> output)
genOut output -> Integer
outValue (Integer -> Integer) -> output -> output
modifyOutValue =
  forall (m :: * -> *) a (n :: * -> *) b.
(TreeT (MaybeT m) a -> TreeT (MaybeT n) b) -> GenT m a -> GenT n b
mapGenT
    ( forall (m :: * -> *) a. m (NodeT m a) -> TreeT m a
TreeT
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. ( forall (m :: * -> *) input output.
Monad m =>
(input -> Integer)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> (TreeT m [input], TreeT m [output])
-> m (NodeT m ([input], [output]))
interleaveInputOutputTreeT input -> Integer
inValue output -> Integer
outValue (Integer -> Integer) -> output -> output
modifyOutValue
              forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. NodeT m a -> a
nodeValue
              forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<
          )
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. TreeT m a -> m (NodeT m a)
runTreeT
    )
    forall a b. (a -> b) -> a -> b
$ do
      TreeT (MaybeT Identity) [input]
insTree <- forall (m :: * -> *) a.
MonadGen m =>
m a -> m (TreeT (MaybeT (GenBase m)) a)
toTreeMaybeT (forall (m :: * -> *) a. MonadGen m => [a] -> m [a]
Gen.subsequence [input]
ins)
      case forall a. Tree a -> a
treeValue (forall (m :: * -> *) a.
Monad m =>
TreeT (MaybeT m) a -> TreeT m (Maybe a)
runDiscardEffectT TreeT (MaybeT Identity) [input]
insTree) of
        Maybe [input]
Nothing -> forall (f :: * -> *) a. Alternative f => f a
empty
        Just [input]
is ->
          (,)
            forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) a. Applicative f => a -> f a
pure TreeT (MaybeT Identity) [input]
insTree
            forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (m :: * -> *) a.
MonadGen m =>
m a -> m (TreeT (MaybeT (GenBase m)) a)
toTreeMaybeT
              ( forall a.
Integer
-> Gen (Integer -> a)
-> (a -> Integer)
-> ((Integer -> Integer) -> a -> a)
-> Gen [a]
genSplitValue
                  (forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum forall a b. (a -> b) -> a -> b
$ input -> Integer
inValue forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [input]
is)
                  Gen (Integer -> output)
genOut
                  output -> Integer
outValue
                  (Integer -> Integer) -> output -> output
modifyOutValue
              )

-- | Used as part of 'genInputOutput', so see there for details of the arguments
interleaveInputOutputTreeT ::
  Monad m =>
  (input -> Integer) ->
  (output -> Integer) ->
  ((Integer -> Integer) -> output -> output) ->
  (TreeT m [input], TreeT m [output]) ->
  m (NodeT m ([input], [output]))
interleaveInputOutputTreeT :: forall (m :: * -> *) input output.
Monad m =>
(input -> Integer)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> (TreeT m [input], TreeT m [output])
-> m (NodeT m ([input], [output]))
interleaveInputOutputTreeT input -> Integer
inValue output -> Integer
outValue (Integer -> Integer) -> output -> output
modifyOutValue =
  forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (m :: * -> *) input output.
Monad m =>
(input -> Integer)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> (NodeT m [input], NodeT m [output])
-> NodeT m ([input], [output])
interleaveInputOutput input -> Integer
inValue output -> Integer
outValue (Integer -> Integer) -> output -> output
modifyOutValue)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse forall (m :: * -> *) a. TreeT m a -> m (NodeT m a)
runTreeT forall (m :: * -> *) a. TreeT m a -> m (NodeT m a)
runTreeT

-- | Used as part of 'genInputOutput', so see there for details of the arguments
interleaveInputOutput ::
  Monad m =>
  (input -> Integer) ->
  (output -> Integer) ->
  ((Integer -> Integer) -> output -> output) ->
  (NodeT m [input], NodeT m [output]) ->
  NodeT m ([input], [output])
interleaveInputOutput :: forall (m :: * -> *) input output.
Monad m =>
(input -> Integer)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> (NodeT m [input], NodeT m [output])
-> NodeT m ([input], [output])
interleaveInputOutput input -> Integer
inValue output -> Integer
outValue (Integer -> Integer) -> output -> output
modifyOutValue (NodeT m [input]
as, NodeT m [output]
bs) =
  case forall (m :: * -> *) a. NodeT m a -> a
nodeValue NodeT m [input]
as of
    [] -> forall (m :: * -> *) a. a -> [TreeT m a] -> NodeT m a
NodeT ([], []) []
    [input]
_ ->
      forall (m :: * -> *) a. a -> [TreeT m a] -> NodeT m a
NodeT
        (forall (m :: * -> *) a. NodeT m a -> a
nodeValue NodeT m [input]
as, forall (m :: * -> *) a. NodeT m a -> a
nodeValue NodeT m [output]
bs)
        ( forall (m :: * -> *) input output.
Monad m =>
(input -> Integer)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> (NodeT m [input], NodeT m [output])
-> [TreeT m ([input], [output])]
shrinkLeftPreserving input -> Integer
inValue output -> Integer
outValue (Integer -> Integer) -> output -> output
modifyOutValue (NodeT m [input]
as, NodeT m [output]
bs)
            forall a. [a] -> [a] -> [a]
++ forall (m :: * -> *) input output.
Monad m =>
(input -> Integer)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> (NodeT m [input], NodeT m [output])
-> [TreeT m ([input], [output])]
shrinkRight input -> Integer
inValue output -> Integer
outValue (Integer -> Integer) -> output -> output
modifyOutValue (NodeT m [input]
as, NodeT m [output]
bs)
        )

-- | Used as part of 'genInputOutput', so see there for details of the arguments
shrinkRight ::
  Monad m =>
  (input -> Integer) ->
  (output -> Integer) ->
  ((Integer -> Integer) -> output -> output) ->
  (NodeT m [input], NodeT m [output]) ->
  [TreeT m ([input], [output])]
shrinkRight :: forall (m :: * -> *) input output.
Monad m =>
(input -> Integer)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> (NodeT m [input], NodeT m [output])
-> [TreeT m ([input], [output])]
shrinkRight input -> Integer
inValue output -> Integer
outValue (Integer -> Integer) -> output -> output
modifyOutValue (NodeT m [input]
xs, NodeT m [output]
ys1) = do
  TreeT m [output]
ys2 <- forall (m :: * -> *) a. NodeT m a -> [TreeT m a]
nodeChildren NodeT m [output]
ys1
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. m (NodeT m a) -> TreeT m a
TreeT forall a b. (a -> b) -> a -> b
$ do
    NodeT m [output]
ys3 <- forall (m :: * -> *) a. TreeT m a -> m (NodeT m a)
runTreeT TreeT m [output]
ys2
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) input output.
Monad m =>
(input -> Integer)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> (NodeT m [input], NodeT m [output])
-> NodeT m ([input], [output])
interleaveInputOutput input -> Integer
inValue output -> Integer
outValue (Integer -> Integer) -> output -> output
modifyOutValue (NodeT m [input]
xs, NodeT m [output]
ys3)

-- | Shrink the left value of a tuple, preserving the total value stored in the
--   right value
--
--   Used as part of 'genInputOutput', so see there for details of the arguments
shrinkLeftPreserving ::
  Monad m =>
  (input -> Integer) ->
  (output -> Integer) ->
  ((Integer -> Integer) -> output -> output) ->
  (NodeT m [input], NodeT m [output]) ->
  [TreeT m ([input], [output])]
shrinkLeftPreserving :: forall (m :: * -> *) input output.
Monad m =>
(input -> Integer)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> (NodeT m [input], NodeT m [output])
-> [TreeT m ([input], [output])]
shrinkLeftPreserving input -> Integer
inValue output -> Integer
outValue (Integer -> Integer) -> output -> output
modifyOutValue (NodeT m [input]
xs1, NodeT m [output]
ys1) = do
  TreeT m [input]
xs2 <- forall (m :: * -> *) a. NodeT m a -> [TreeT m a]
nodeChildren NodeT m [input]
xs1
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. m (NodeT m a) -> TreeT m a
TreeT forall a b. (a -> b) -> a -> b
$ do
    NodeT m [input]
xs3 <- forall (m :: * -> *) a. TreeT m a -> m (NodeT m a)
runTreeT TreeT m [input]
xs2
    let lost :: Integer
lost = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (input -> Integer
inValue forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a. NodeT m a -> a
nodeValue NodeT m [input]
xs1) forall a. Num a => a -> a -> a
- forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (input -> Integer
inValue forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a. NodeT m a -> a
nodeValue NodeT m [input]
xs3)
        ys2 :: NodeT m [output]
ys2 = forall n a.
(Num n, Ord n) =>
n -> (a -> n) -> ((n -> n) -> a -> a) -> [a] -> [a]
subFromList Integer
lost output -> Integer
outValue (Integer -> Integer) -> output -> output
modifyOutValue forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NodeT m [output]
ys1
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) input output.
Monad m =>
(input -> Integer)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> (NodeT m [input], NodeT m [output])
-> NodeT m ([input], [output])
interleaveInputOutput input -> Integer
inValue output -> Integer
outValue (Integer -> Integer) -> output -> output
modifyOutValue (NodeT m [input]
xs3, NodeT m [output]
ys2)

-- | Remove total value from a list, removing from the front
subFromList ::
  (Num n, Ord n) =>
  -- | The total value to remove from the list
  n ->
  -- | A view into the value contained in type @a@
  (a -> n) ->
  -- | A modifier for the value contained in type @a@
  ((n -> n) -> a -> a) ->
  -- | The list of @a@s to remove value from
  [a] ->
  [a]
subFromList :: forall n a.
(Num n, Ord n) =>
n -> (a -> n) -> ((n -> n) -> a -> a) -> [a] -> [a]
subFromList n
n a -> n
getVal (n -> n) -> a -> a
modifyVal = n -> [a] -> [a]
go n
n
  where
    go :: n -> [a] -> [a]
go n
0 [a]
x = [a]
x
    go n
_ [] = []
    go n
n' (a
x : [a]
xs) =
      if a -> n
getVal a
x forall a. Ord a => a -> a -> Bool
> n
n'
        then (n -> n) -> a -> a
modifyVal (forall a. Num a => a -> a -> a
subtract n
n') a
x forall a. a -> [a] -> [a]
: [a]
xs
        else n -> [a] -> [a]
go (n
n' forall a. Num a => a -> a -> a
- a -> n
getVal a
x) [a]
xs

-- | A property to check that `genSplitValue` does indeed preserve the input
propGenSplitValue :: Property
propGenSplitValue :: Property
propGenSplitValue = HasCallStack => PropertyT IO () -> Property
property forall a b. (a -> b) -> a -> b
$ do
  Integer
n <- forall (m :: * -> *) a.
(Monad m, Show a, HasCallStack) =>
Gen a -> PropertyT m a
forAll forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. (MonadGen m, Integral a) => Range a -> m a
Gen.integral (forall a. Integral a => a -> a -> Range a
Range.linear Integer
1 Integer
10000)
  [Integer]
ints <- forall (m :: * -> *) a.
(Monad m, Show a, HasCallStack) =>
Gen a -> PropertyT m a
forAll forall a b. (a -> b) -> a -> b
$ forall a.
Integer
-> Gen (Integer -> a)
-> (a -> Integer)
-> ((Integer -> Integer) -> a -> a)
-> Gen [a]
genSplitValue Integer
n (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. a -> a
id) forall a. a -> a
id forall a. a -> a
id
  forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Integer]
ints forall (m :: * -> *) a.
(MonadTest m, Eq a, Show a, HasCallStack) =>
a -> a -> m ()
=== Integer
n

-- | Given an input value and functions to generate, view, and update some type
--   'a' based on that value, split the input into a number of 'a's, preserving
--   the value throughout the shrink tree
genSplitValue ::
  -- | Total value to divide into outputs
  Integer ->
  -- | Applicative generator for an output given a value
  Gen (Integer -> a) ->
  -- | A view of the value in the output used to preserve the invariant
  (a -> Integer) ->
  -- | A modifier for the value in the output used to preserve the invariant
  ((Integer -> Integer) -> a -> a) ->
  Gen [a]
genSplitValue :: forall a.
Integer
-> Gen (Integer -> a)
-> (a -> Integer)
-> ((Integer -> Integer) -> a -> a)
-> Gen [a]
genSplitValue Integer
n Gen (Integer -> a)
genA a -> Integer
getValue (Integer -> Integer) -> a -> a
modifyValue =
  forall (m :: * -> *) a (n :: * -> *) b.
(TreeT (MaybeT m) a -> TreeT (MaybeT n) b) -> GenT m a -> GenT n b
mapGenT
    ( forall (m :: * -> *) a. m (NodeT m a) -> TreeT m a
TreeT
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall (m :: * -> *) a.
Monad m =>
(a -> Integer)
-> ((Integer -> Integer) -> a -> a)
-> [TreeT m a]
-> m (NodeT m [a])
interleaveTreeTPreserving a -> Integer
getValue (Integer -> Integer) -> a -> a
modifyValue forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. NodeT m a -> a
nodeValue forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<)
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. TreeT m a -> m (NodeT m a)
runTreeT
    )
    forall a b. (a -> b) -> a -> b
$ Integer
-> [TreeT (MaybeT Identity) a]
-> GenT Identity [TreeT (MaybeT Identity) a]
go Integer
n []
  where
    go :: Integer
-> [TreeT (MaybeT Identity) a]
-> GenT Identity [TreeT (MaybeT Identity) a]
go Integer
0 [TreeT (MaybeT Identity) a]
acc = forall (f :: * -> *) a. Applicative f => a -> f a
pure [TreeT (MaybeT Identity) a]
acc
    go Integer
left [TreeT (MaybeT Identity) a]
acc = do
      TreeT (MaybeT Identity) a
mTree <- forall (m :: * -> *) a.
MonadGen m =>
m a -> m (TreeT (MaybeT (GenBase m)) a)
toTreeMaybeT (Gen (Integer -> a)
genA forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (m :: * -> *) a. (MonadGen m, Integral a) => Range a -> m a
Gen.integral (forall a. a -> a -> Range a
Range.constant Integer
1 Integer
left))
      case forall a. Tree a -> a
treeValue (forall (m :: * -> *) a.
Monad m =>
TreeT (MaybeT m) a -> TreeT m (Maybe a)
runDiscardEffectT TreeT (MaybeT Identity) a
mTree) of
        Maybe a
Nothing -> forall (f :: * -> *) a. Alternative f => f a
empty
        Just a
a -> Integer
-> [TreeT (MaybeT Identity) a]
-> GenT Identity [TreeT (MaybeT Identity) a]
go (Integer
left forall a. Num a => a -> a -> a
- a -> Integer
getValue a
a) (TreeT (MaybeT Identity) a
mTree forall a. a -> [a] -> [a]
: [TreeT (MaybeT Identity) a]
acc)

-- | Used as part of 'genSplitValue', so see there for details of the arguments
interleaveTreeTPreserving ::
  Monad m =>
  (a -> Integer) ->
  ((Integer -> Integer) -> a -> a) ->
  [TreeT m a] ->
  m (NodeT m [a])
interleaveTreeTPreserving :: forall (m :: * -> *) a.
Monad m =>
(a -> Integer)
-> ((Integer -> Integer) -> a -> a)
-> [TreeT m a]
-> m (NodeT m [a])
interleaveTreeTPreserving a -> Integer
getValue (Integer -> Integer) -> a -> a
modifyValue =
  forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (m :: * -> *) a.
Monad m =>
(a -> Integer)
-> ((Integer -> Integer) -> a -> a) -> [NodeT m a] -> NodeT m [a]
interleavePreserving a -> Integer
getValue (Integer -> Integer) -> a -> a
modifyValue) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall (m :: * -> *) a. TreeT m a -> m (NodeT m a)
runTreeT

-- | Used as part of 'genSplitValue', so see there for details of the arguments
interleavePreserving ::
  Monad m =>
  (a -> Integer) ->
  ((Integer -> Integer) -> a -> a) ->
  [NodeT m a] ->
  NodeT m [a]
interleavePreserving :: forall (m :: * -> *) a.
Monad m =>
(a -> Integer)
-> ((Integer -> Integer) -> a -> a) -> [NodeT m a] -> NodeT m [a]
interleavePreserving a -> Integer
getValue (Integer -> Integer) -> a -> a
modifyValue [NodeT m a]
ts =
  forall (m :: * -> *) a. a -> [TreeT m a] -> NodeT m a
NodeT
    (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (m :: * -> *) a. NodeT m a -> a
nodeValue [NodeT m a]
ts)
    ( forall (m :: * -> *) a.
Monad m =>
(a -> Integer)
-> ((Integer -> Integer) -> a -> a) -> [NodeT m a] -> [TreeT m [a]]
dropOnePreserving a -> Integer
getValue (Integer -> Integer) -> a -> a
modifyValue [NodeT m a]
ts
        forall a. [a] -> [a] -> [a]
++ forall (m :: * -> *) a.
Monad m =>
(a -> Integer)
-> ((Integer -> Integer) -> a -> a) -> [NodeT m a] -> [TreeT m [a]]
shrinkOnePreserving a -> Integer
getValue (Integer -> Integer) -> a -> a
modifyValue [NodeT m a]
ts
    )

-- | Drop one of the outputs, preserving the invariant by moving its value to
--   the left
--
--   Used as part of 'genSplitValue', so see there for details of the arguments
dropOnePreserving ::
  Monad m =>
  (a -> Integer) ->
  ((Integer -> Integer) -> a -> a) ->
  [NodeT m a] ->
  [TreeT m [a]]
dropOnePreserving :: forall (m :: * -> *) a.
Monad m =>
(a -> Integer)
-> ((Integer -> Integer) -> a -> a) -> [NodeT m a] -> [TreeT m [a]]
dropOnePreserving a -> Integer
getValue (Integer -> Integer) -> a -> a
modifyValue [NodeT m a]
ts = do
  ([NodeT m a]
ws, NodeT m a
x, NodeT m a
y, [NodeT m a]
zs) <- forall a. [a] -> [([a], a, a, [a])]
viewTwo [NodeT m a]
ts
  let x' :: NodeT m a
x' = (Integer -> Integer) -> a -> a
modifyValue (forall a. Num a => a -> a -> a
+ a -> Integer
getValue (forall (m :: * -> *) a. NodeT m a -> a
nodeValue NodeT m a
y)) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NodeT m a
x
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. m (NodeT m a) -> TreeT m a
TreeT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
    forall (m :: * -> *) a.
Monad m =>
(a -> Integer)
-> ((Integer -> Integer) -> a -> a) -> [NodeT m a] -> NodeT m [a]
interleavePreserving
      a -> Integer
getValue
      (Integer -> Integer) -> a -> a
modifyValue
      ([NodeT m a]
ws forall a. [a] -> [a] -> [a]
++ [NodeT m a
x'] forall a. [a] -> [a] -> [a]
++ [NodeT m a]
zs)

-- | Shrink a value in a list, preserving the total value by moving the lost
--   value to the left
--
--   Used as part of 'genSplitValue', so see there for details of the arguments
shrinkOnePreserving ::
  Monad m =>
  (a -> Integer) ->
  ((Integer -> Integer) -> a -> a) ->
  [NodeT m a] ->
  [TreeT m [a]]
shrinkOnePreserving :: forall (m :: * -> *) a.
Monad m =>
(a -> Integer)
-> ((Integer -> Integer) -> a -> a) -> [NodeT m a] -> [TreeT m [a]]
shrinkOnePreserving a -> Integer
getValue (Integer -> Integer) -> a -> a
modifyValue [NodeT m a]
ts = do
  ([NodeT m a]
ws, NodeT m a
x, NodeT m a
y0, [NodeT m a]
zs) <- forall a. [a] -> [([a], a, a, [a])]
viewTwo [NodeT m a]
ts
  TreeT m a
y1 <- forall (m :: * -> *) a. NodeT m a -> [TreeT m a]
nodeChildren NodeT m a
y0
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. m (NodeT m a) -> TreeT m a
TreeT forall a b. (a -> b) -> a -> b
$ do
    NodeT m a
y2 <- forall (m :: * -> *) a. TreeT m a -> m (NodeT m a)
runTreeT TreeT m a
y1
    let lost :: Integer
lost = a -> Integer
getValue (forall (m :: * -> *) a. NodeT m a -> a
nodeValue NodeT m a
y0) forall a. Num a => a -> a -> a
- a -> Integer
getValue (forall (m :: * -> *) a. NodeT m a -> a
nodeValue NodeT m a
y2)
        x' :: NodeT m a
x' = (Integer -> Integer) -> a -> a
modifyValue (forall a. Num a => a -> a -> a
+ Integer
lost) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NodeT m a
x
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
Monad m =>
(a -> Integer)
-> ((Integer -> Integer) -> a -> a) -> [NodeT m a] -> NodeT m [a]
interleavePreserving a -> Integer
getValue (Integer -> Integer) -> a -> a
modifyValue ([NodeT m a]
ws forall a. [a] -> [a] -> [a]
++ [NodeT m a
x', NodeT m a
y2] forall a. [a] -> [a] -> [a]
++ [NodeT m a]
zs)

-- | All the ways of choosing two consecutive values from a list
viewTwo :: [a] -> [([a], a, a, [a])]
viewTwo :: forall a. [a] -> [([a], a, a, [a])]
viewTwo = \case
  [] -> []
  [a
_] -> []
  a
x : a
x' : [a]
xs ->
    ([], a
x, a
x', [a]
xs)
      forall a. a -> [a] -> [a]
: forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\([a]
as, a
b, a
c, [a]
ds) -> (a
x forall a. a -> [a] -> [a]
: [a]
as, a
b, a
c, [a]
ds)) (forall a. [a] -> [([a], a, a, [a])]
viewTwo (a
x' forall a. a -> [a] -> [a]
: [a]
xs))