{-# 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 =
  ([TxOut] -> Bool) -> Gen [TxOut] -> Gen [TxOut]
forall (m :: * -> *) a.
(MonadGen m, GenBase m ~ Identity) =>
(a -> Bool) -> m a -> m a
Gen.filter (Bool -> Bool
not (Bool -> Bool) -> ([TxOut] -> Bool) -> [TxOut] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TxOut] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null)
    (Gen [TxOut] -> Gen [TxOut])
-> ([Addr] -> Gen [TxOut]) -> [Addr] -> Gen [TxOut]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Addr -> Gen TxOut) -> [Addr] -> Gen [TxOut]
forall a b. (a -> Gen b) -> [a] -> Gen [b]
genTraverseSubsequence (\Addr
a -> Addr -> Lovelace -> TxOut
TxOut Addr
a (Lovelace -> TxOut) -> GenT Identity Lovelace -> Gen TxOut
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GenT Identity Lovelace
genLovelace)

genLovelace :: Gen Lovelace
genLovelace :: GenT Identity Lovelace
genLovelace =
  Integer -> Lovelace
Lovelace (Integer -> Lovelace) -> (Word32 -> Integer) -> Word32 -> Lovelace
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word32 -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32 -> Lovelace)
-> GenT Identity Word32 -> GenT Identity Lovelace
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Range Word32 -> GenT Identity Word32
forall (m :: * -> *). MonadGen m => Range Word32 -> m Word32
Gen.word32 (Word32 -> Word32 -> Word32 -> Range 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 = Double -> Word32
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Word32 -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32
mx Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
- Word32
mn) Double -> Double -> Double
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 =
  (TreeT (MaybeT Identity) [TreeT (MaybeT Identity) b]
 -> TreeT (MaybeT Identity) [b])
-> GenT Identity [TreeT (MaybeT Identity) b] -> GenT Identity [b]
forall (m :: * -> *) a (n :: * -> *) b.
(TreeT (MaybeT m) a -> TreeT (MaybeT n) b) -> GenT m a -> GenT n b
mapGenT (MaybeT Identity (NodeT (MaybeT Identity) [b])
-> TreeT (MaybeT Identity) [b]
forall (m :: * -> *) a. m (NodeT m a) -> TreeT m a
TreeT (MaybeT Identity (NodeT (MaybeT Identity) [b])
 -> TreeT (MaybeT Identity) [b])
-> (TreeT (MaybeT Identity) [TreeT (MaybeT Identity) b]
    -> MaybeT Identity (NodeT (MaybeT Identity) [b]))
-> TreeT (MaybeT Identity) [TreeT (MaybeT Identity) b]
-> TreeT (MaybeT Identity) [b]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([TreeT (MaybeT Identity) b]
-> MaybeT Identity (NodeT (MaybeT Identity) [b])
forall (m :: * -> *) a. Monad m => [TreeT m a] -> m (NodeT m [a])
interleaveTreeT ([TreeT (MaybeT Identity) b]
 -> MaybeT Identity (NodeT (MaybeT Identity) [b]))
-> (NodeT (MaybeT Identity) [TreeT (MaybeT Identity) b]
    -> [TreeT (MaybeT Identity) b])
-> NodeT (MaybeT Identity) [TreeT (MaybeT Identity) b]
-> MaybeT Identity (NodeT (MaybeT Identity) [b])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NodeT (MaybeT Identity) [TreeT (MaybeT Identity) b]
-> [TreeT (MaybeT Identity) b]
forall (m :: * -> *) a. NodeT m a -> a
nodeValue (NodeT (MaybeT Identity) [TreeT (MaybeT Identity) b]
 -> MaybeT Identity (NodeT (MaybeT Identity) [b]))
-> MaybeT
     Identity (NodeT (MaybeT Identity) [TreeT (MaybeT Identity) b])
-> MaybeT Identity (NodeT (MaybeT Identity) [b])
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<) (MaybeT
   Identity (NodeT (MaybeT Identity) [TreeT (MaybeT Identity) b])
 -> MaybeT Identity (NodeT (MaybeT Identity) [b]))
-> (TreeT (MaybeT Identity) [TreeT (MaybeT Identity) b]
    -> MaybeT
         Identity (NodeT (MaybeT Identity) [TreeT (MaybeT Identity) b]))
-> TreeT (MaybeT Identity) [TreeT (MaybeT Identity) b]
-> MaybeT Identity (NodeT (MaybeT Identity) [b])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TreeT (MaybeT Identity) [TreeT (MaybeT Identity) b]
-> MaybeT
     Identity (NodeT (MaybeT Identity) [TreeT (MaybeT Identity) b])
forall (m :: * -> *) a. TreeT m a -> m (NodeT m a)
runTreeT) (GenT Identity [TreeT (MaybeT Identity) b] -> GenT Identity [b])
-> GenT Identity [TreeT (MaybeT Identity) b] -> GenT Identity [b]
forall a b. (a -> b) -> a -> b
$ do
    [a]
sub <- [a] -> GenT Identity [a]
forall (m :: * -> *) a. MonadGen m => [a] -> m [a]
Gen.subsequence [a]
as
    (a -> GenT Identity (TreeT (MaybeT Identity) b))
-> [a] -> GenT Identity [TreeT (MaybeT Identity) b]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (Gen b -> GenT Identity (TreeT (MaybeT Identity) b)
Gen b -> GenT Identity (TreeT (MaybeT (GenBase (GenT Identity))) b)
forall (m :: * -> *) a.
MonadGen m =>
m a -> m (TreeT (MaybeT (GenBase m)) a)
toTreeMaybeT (Gen b -> GenT Identity (TreeT (MaybeT Identity) b))
-> (a -> Gen b) -> a -> GenT Identity (TreeT (MaybeT Identity) b)
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 = (Size -> GenT Identity [a]) -> GenT Identity [a]
forall (m :: * -> *) a. MonadGen m => (Size -> m a) -> m a
Gen.sized ((Size -> GenT Identity [a]) -> GenT Identity [a])
-> (Size -> GenT Identity [a]) -> GenT Identity [a]
forall a b. (a -> b) -> a -> b
$ \Size
gSize ->
  ([a] -> Bool) -> GenT Identity [a] -> GenT Identity [a]
forall (m :: * -> *) a. MonadGen m => (a -> Bool) -> m a -> m a
ensure (Int -> [a] -> Bool
forall a. Int -> [a] -> Bool
atLeast (Int -> [a] -> Bool) -> Int -> [a] -> Bool
forall a b. (a -> b) -> a -> b
$ Size -> Range Int -> Int
forall a. Ord a => Size -> Range a -> a
Range.lowerBound Size
gSize Range Int
range) (GenT Identity [a] -> GenT Identity [a])
-> GenT Identity [a] -> GenT Identity [a]
forall a b. (a -> b) -> a -> b
$
    (() -> Gen a) -> [()] -> GenT Identity [a]
forall a b. (a -> Gen b) -> [a] -> Gen [b]
genTraverseSubsequence
      (Gen a -> () -> Gen a
forall a b. a -> b -> a
const Gen a
gen)
      (Int -> () -> [()]
forall a. Int -> a -> [a]
replicate (Size -> Range Int -> Int
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 = ([NodeT m a] -> NodeT m [a]) -> m [NodeT m a] -> m (NodeT m [a])
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [NodeT m a] -> NodeT m [a]
forall (m :: * -> *) a. Monad m => [NodeT m a] -> NodeT m [a]
Tree.interleave (m [NodeT m a] -> m (NodeT m [a]))
-> ([TreeT m a] -> m [NodeT m a]) -> [TreeT m a] -> m (NodeT m [a])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TreeT m a -> m (NodeT m a)) -> [TreeT m a] -> m [NodeT m a]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse TreeT m a -> m (NodeT m a)
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 (TxIn -> Wit) -> [TxIn] -> [Wit]
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 (Gen Tx -> Gen Tx) -> Gen Tx -> Gen Tx
forall a b. (a -> b) -> a -> b
$ Tx -> Gen Tx
forall a. a -> GenT Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Tx
tx
  where
    genTxBody :: GenT Identity TxBody
genTxBody =
      ([TxIn] -> [TxOut] -> TxBody) -> ([TxIn], [TxOut]) -> TxBody
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry [TxIn] -> [TxOut] -> TxBody
TxBody
        (([TxIn], [TxOut]) -> TxBody)
-> GenT Identity ([TxIn], [TxOut]) -> GenT Identity TxBody
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (([TxIn], [TxOut]) -> Bool)
-> GenT Identity ([TxIn], [TxOut])
-> GenT Identity ([TxIn], [TxOut])
forall (m :: * -> *) a.
(MonadGen m, GenBase m ~ Identity) =>
(a -> Bool) -> m a -> m a
Gen.filter
          (Bool -> Bool
not (Bool -> Bool)
-> (([TxIn], [TxOut]) -> Bool) -> ([TxIn], [TxOut]) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TxIn] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([TxIn] -> Bool)
-> (([TxIn], [TxOut]) -> [TxIn]) -> ([TxIn], [TxOut]) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([TxIn], [TxOut]) -> [TxIn]
forall a b. (a, b) -> a
fst)
          ( [TxIn]
-> (TxIn -> Integer)
-> Gen (Integer -> TxOut)
-> (TxOut -> Integer)
-> ((Integer -> Integer) -> TxOut -> TxOut)
-> GenT Identity ([TxIn], [TxOut])
forall input output.
[input]
-> (input -> Integer)
-> Gen (Integer -> output)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> Gen ([input], [output])
genInputOutput
              (Map TxIn TxOut -> [TxIn]
forall k a. Map k a -> [k]
M.keys Map TxIn TxOut
utxo)
              (Integer -> (TxOut -> Integer) -> Maybe TxOut -> Integer
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Integer
0 (Lovelace -> Integer
unLovelace (Lovelace -> Integer) -> (TxOut -> Lovelace) -> TxOut -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TxOut -> Lovelace
value) (Maybe TxOut -> Integer)
-> (TxIn -> Maybe TxOut) -> TxIn -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TxIn -> Map TxIn TxOut -> Maybe TxOut)
-> Map TxIn TxOut -> TxIn -> Maybe TxOut
forall a b c. (a -> b -> c) -> b -> a -> c
flip TxIn -> Map TxIn TxOut -> Maybe TxOut
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Map TxIn TxOut
utxo)
              (((Lovelace -> TxOut) -> Integer -> TxOut)
-> GenT Identity (Lovelace -> TxOut) -> Gen (Integer -> TxOut)
forall a b. (a -> b) -> GenT Identity a -> GenT Identity b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Lovelace -> TxOut) -> (Integer -> Lovelace) -> Integer -> TxOut
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Lovelace
Lovelace) (GenT Identity (Lovelace -> TxOut) -> Gen (Integer -> TxOut))
-> GenT Identity (Lovelace -> TxOut) -> Gen (Integer -> TxOut)
forall a b. (a -> b) -> a -> b
$ Addr -> Lovelace -> TxOut
TxOut (Addr -> Lovelace -> TxOut)
-> GenT Identity Addr -> GenT Identity (Lovelace -> TxOut)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Addr] -> GenT Identity Addr
forall (f :: * -> *) (m :: * -> *) a.
(Foldable f, MonadGen m) =>
f a -> m a
Gen.element [Addr]
addrs)
              (Lovelace -> Integer
unLovelace (Lovelace -> Integer) -> (TxOut -> Lovelace) -> TxOut -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TxOut -> Lovelace
value)
              (\Integer -> Integer
f TxOut
out -> TxOut
out {value = Lovelace . f . unLovelace $ value out})
          )

    witnessForTxIn :: TxBody -> TxIn -> Wit
    witnessForTxIn :: TxBody -> TxIn -> Wit
witnessForTxIn TxBody
tx TxIn
txin =
      case TxIn -> Map TxIn TxOut -> Maybe TxOut
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 (Owner -> KeyPair) -> Owner -> KeyPair
forall a b. (a -> b) -> a -> b
$ VKey -> Owner
forall a. HasOwner a => a -> Owner
owner VKey
pay) TxBody
tx
        Maybe TxOut
Nothing ->
          [Char] -> Wit
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 (SKey -> TxBody -> Sig TxBody
forall a. SKey -> a -> Sig a
sign SKey
sk TxBody
tx)
    subtractFees :: Gen Tx -> Gen Tx
    subtractFees :: Gen Tx -> Gen Tx
subtractFees =
      (Tx -> Tx) -> Gen Tx -> Gen Tx
forall a b. (a -> b) -> GenT Identity a -> GenT Identity b
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.
        (Gen Tx -> Gen Tx) -> (Gen Tx -> Gen Tx) -> Gen Tx -> Gen Tx
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Tx -> Bool) -> Gen Tx -> Gen Tx
forall (m :: * -> *) a.
(MonadGen m, GenBase m ~ Identity) =>
(a -> Bool) -> m a -> m a
Gen.filter (\tx :: Tx
tx@(Tx TxBody
txb [Wit]
_) -> [Lovelace] -> Lovelace
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (TxOut -> Lovelace
value (TxOut -> Lovelace) -> [TxOut] -> [Lovelace]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TxBody -> [TxOut]
outputs TxBody
txb) Lovelace -> Lovelace -> Bool
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 = subFromList (txfee tx) value updateValue (outputs txb)}
           in TxBody -> [Wit] -> Tx
Tx TxBody
newBody (TxBody -> TxIn -> Wit
witnessForTxIn TxBody
newBody (TxIn -> Wit) -> [TxIn] -> [Wit]
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 = f (value 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
PropertyT IO () -> Property
property (PropertyT IO () -> Property) -> PropertyT IO () -> Property
forall a b. (a -> b) -> a -> b
$ do
  [Integer]
ins <- Gen [Integer] -> PropertyT IO [Integer]
forall (m :: * -> *) a.
(Monad m, Show a, HasCallStack) =>
Gen a -> PropertyT m a
forAll (Gen [Integer] -> PropertyT IO [Integer])
-> Gen [Integer] -> PropertyT IO [Integer]
forall a b. (a -> b) -> a -> b
$ Range Int -> GenT Identity Integer -> Gen [Integer]
forall (m :: * -> *) a. MonadGen m => Range Int -> m a -> m [a]
Gen.list (Int -> Int -> Range Int
forall a. Integral a => a -> a -> Range a
Range.linear Int
0 Int
5) (Range Integer -> GenT Identity Integer
forall (m :: * -> *) a. (MonadGen m, Integral a) => Range a -> m a
Gen.integral (Integer -> Integer -> Range Integer
forall a. Integral a => a -> a -> Range a
Range.linear Integer
0 Integer
5))
  Tree ([Integer], [Integer])
insOutsTree <- Gen (Tree ([Integer], [Integer]))
-> PropertyT IO (Tree ([Integer], [Integer]))
forall (m :: * -> *) a.
(Monad m, Show a, HasCallStack) =>
Gen a -> PropertyT m a
forAll (Gen (Tree ([Integer], [Integer]))
 -> PropertyT IO (Tree ([Integer], [Integer])))
-> Gen (Tree ([Integer], [Integer]))
-> PropertyT IO (Tree ([Integer], [Integer]))
forall a b. (a -> b) -> a -> b
$ GenT Identity ([Integer], [Integer])
-> Gen (Tree ([Integer], [Integer]))
forall (m :: * -> *) a.
(MonadGen m, GenBase m ~ Identity) =>
m a -> m (Tree a)
toTree ([Integer]
-> (Integer -> Integer)
-> Gen (Integer -> Integer)
-> (Integer -> Integer)
-> ((Integer -> Integer) -> Integer -> Integer)
-> GenT Identity ([Integer], [Integer])
forall input output.
[input]
-> (input -> Integer)
-> Gen (Integer -> output)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> Gen ([input], [output])
genInputOutput [Integer]
ins Integer -> Integer
forall a. a -> a
id ((Integer -> Integer) -> Gen (Integer -> Integer)
forall a. a -> GenT Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Integer -> Integer
forall a. a -> a
id) Integer -> Integer
forall a. a -> a
id (Integer -> Integer) -> Integer -> Integer
forall a. a -> a
id)
  Bool -> PropertyT IO ()
forall (m :: * -> *). (MonadTest m, HasCallStack) => Bool -> m ()
assert (Bool -> PropertyT IO ()) -> Bool -> PropertyT IO ()
forall a b. (a -> b) -> a -> b
$ (([Integer], [Integer]) -> Bool)
-> Tree ([Integer], [Integer]) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\([Integer]
ins', [Integer]
outs) -> [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Integer]
ins' Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== [Integer] -> Integer
forall a. Num a => [a] -> a
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 =
  (TreeT
   (MaybeT Identity)
   (TreeT (MaybeT Identity) [input], TreeT (MaybeT Identity) [output])
 -> TreeT (MaybeT Identity) ([input], [output]))
-> GenT
     Identity
     (TreeT (MaybeT Identity) [input], TreeT (MaybeT Identity) [output])
-> GenT Identity ([input], [output])
forall (m :: * -> *) a (n :: * -> *) b.
(TreeT (MaybeT m) a -> TreeT (MaybeT n) b) -> GenT m a -> GenT n b
mapGenT
    ( MaybeT Identity (NodeT (MaybeT Identity) ([input], [output]))
-> TreeT (MaybeT Identity) ([input], [output])
forall (m :: * -> *) a. m (NodeT m a) -> TreeT m a
TreeT
        (MaybeT Identity (NodeT (MaybeT Identity) ([input], [output]))
 -> TreeT (MaybeT Identity) ([input], [output]))
-> (TreeT
      (MaybeT Identity)
      (TreeT (MaybeT Identity) [input], TreeT (MaybeT Identity) [output])
    -> MaybeT Identity (NodeT (MaybeT Identity) ([input], [output])))
-> TreeT
     (MaybeT Identity)
     (TreeT (MaybeT Identity) [input], TreeT (MaybeT Identity) [output])
-> TreeT (MaybeT Identity) ([input], [output])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ( (input -> Integer)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> (TreeT (MaybeT Identity) [input],
    TreeT (MaybeT Identity) [output])
-> MaybeT Identity (NodeT (MaybeT Identity) ([input], [output]))
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
              ((TreeT (MaybeT Identity) [input],
  TreeT (MaybeT Identity) [output])
 -> MaybeT Identity (NodeT (MaybeT Identity) ([input], [output])))
-> (NodeT
      (MaybeT Identity)
      (TreeT (MaybeT Identity) [input], TreeT (MaybeT Identity) [output])
    -> (TreeT (MaybeT Identity) [input],
        TreeT (MaybeT Identity) [output]))
-> NodeT
     (MaybeT Identity)
     (TreeT (MaybeT Identity) [input], TreeT (MaybeT Identity) [output])
-> MaybeT Identity (NodeT (MaybeT Identity) ([input], [output]))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NodeT
  (MaybeT Identity)
  (TreeT (MaybeT Identity) [input], TreeT (MaybeT Identity) [output])
-> (TreeT (MaybeT Identity) [input],
    TreeT (MaybeT Identity) [output])
forall (m :: * -> *) a. NodeT m a -> a
nodeValue
              (NodeT
   (MaybeT Identity)
   (TreeT (MaybeT Identity) [input], TreeT (MaybeT Identity) [output])
 -> MaybeT Identity (NodeT (MaybeT Identity) ([input], [output])))
-> MaybeT
     Identity
     (NodeT
        (MaybeT Identity)
        (TreeT (MaybeT Identity) [input],
         TreeT (MaybeT Identity) [output]))
-> MaybeT Identity (NodeT (MaybeT Identity) ([input], [output]))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<
          )
        (MaybeT
   Identity
   (NodeT
      (MaybeT Identity)
      (TreeT (MaybeT Identity) [input],
       TreeT (MaybeT Identity) [output]))
 -> MaybeT Identity (NodeT (MaybeT Identity) ([input], [output])))
-> (TreeT
      (MaybeT Identity)
      (TreeT (MaybeT Identity) [input], TreeT (MaybeT Identity) [output])
    -> MaybeT
         Identity
         (NodeT
            (MaybeT Identity)
            (TreeT (MaybeT Identity) [input],
             TreeT (MaybeT Identity) [output])))
-> TreeT
     (MaybeT Identity)
     (TreeT (MaybeT Identity) [input], TreeT (MaybeT Identity) [output])
-> MaybeT Identity (NodeT (MaybeT Identity) ([input], [output]))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TreeT
  (MaybeT Identity)
  (TreeT (MaybeT Identity) [input], TreeT (MaybeT Identity) [output])
-> MaybeT
     Identity
     (NodeT
        (MaybeT Identity)
        (TreeT (MaybeT Identity) [input],
         TreeT (MaybeT Identity) [output]))
forall (m :: * -> *) a. TreeT m a -> m (NodeT m a)
runTreeT
    )
    (GenT
   Identity
   (TreeT (MaybeT Identity) [input], TreeT (MaybeT Identity) [output])
 -> GenT Identity ([input], [output]))
-> GenT
     Identity
     (TreeT (MaybeT Identity) [input], TreeT (MaybeT Identity) [output])
-> GenT Identity ([input], [output])
forall a b. (a -> b) -> a -> b
$ do
      TreeT (MaybeT Identity) [input]
insTree <- GenT Identity [input]
-> GenT Identity (TreeT (MaybeT (GenBase (GenT Identity))) [input])
forall (m :: * -> *) a.
MonadGen m =>
m a -> m (TreeT (MaybeT (GenBase m)) a)
toTreeMaybeT ([input] -> GenT Identity [input]
forall (m :: * -> *) a. MonadGen m => [a] -> m [a]
Gen.subsequence [input]
ins)
      case Tree (Maybe [input]) -> Maybe [input]
forall a. Tree a -> a
treeValue (TreeT (MaybeT Identity) [input] -> Tree (Maybe [input])
forall (m :: * -> *) a.
Monad m =>
TreeT (MaybeT m) a -> TreeT m (Maybe a)
runDiscardEffectT TreeT (MaybeT Identity) [input]
insTree) of
        Maybe [input]
Nothing -> GenT
  Identity
  (TreeT (MaybeT Identity) [input], TreeT (MaybeT Identity) [output])
forall a. GenT Identity a
forall (f :: * -> *) a. Alternative f => f a
empty
        Just [input]
is ->
          (,)
            (TreeT (MaybeT Identity) [input]
 -> TreeT (MaybeT Identity) [output]
 -> (TreeT (MaybeT Identity) [input],
     TreeT (MaybeT Identity) [output]))
-> GenT Identity (TreeT (MaybeT Identity) [input])
-> GenT
     Identity
     (TreeT (MaybeT Identity) [output]
      -> (TreeT (MaybeT Identity) [input],
          TreeT (MaybeT Identity) [output]))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TreeT (MaybeT Identity) [input]
-> GenT Identity (TreeT (MaybeT Identity) [input])
forall a. a -> GenT Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TreeT (MaybeT Identity) [input]
insTree
            GenT
  Identity
  (TreeT (MaybeT Identity) [output]
   -> (TreeT (MaybeT Identity) [input],
       TreeT (MaybeT Identity) [output]))
-> GenT Identity (TreeT (MaybeT Identity) [output])
-> GenT
     Identity
     (TreeT (MaybeT Identity) [input], TreeT (MaybeT Identity) [output])
forall a b.
GenT Identity (a -> b) -> GenT Identity a -> GenT Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> GenT Identity [output]
-> GenT
     Identity (TreeT (MaybeT (GenBase (GenT Identity))) [output])
forall (m :: * -> *) a.
MonadGen m =>
m a -> m (TreeT (MaybeT (GenBase m)) a)
toTreeMaybeT
              ( Integer
-> Gen (Integer -> output)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> GenT Identity [output]
forall a.
Integer
-> Gen (Integer -> a)
-> (a -> Integer)
-> ((Integer -> Integer) -> a -> a)
-> Gen [a]
genSplitValue
                  ([Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Integer] -> Integer) -> [Integer] -> Integer
forall a b. (a -> b) -> a -> b
$ input -> Integer
inValue (input -> Integer) -> [input] -> [Integer]
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 =
  ((NodeT m [input], NodeT m [output])
 -> NodeT m ([input], [output]))
-> m (NodeT m [input], NodeT m [output])
-> m (NodeT m ([input], [output]))
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((input -> Integer)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> (NodeT m [input], NodeT m [output])
-> NodeT m ([input], [output])
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)
    (m (NodeT m [input], NodeT m [output])
 -> m (NodeT m ([input], [output])))
-> ((TreeT m [input], TreeT m [output])
    -> m (NodeT m [input], NodeT m [output]))
-> (TreeT m [input], TreeT m [output])
-> m (NodeT m ([input], [output]))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TreeT m [input] -> m (NodeT m [input]))
-> (TreeT m [output] -> m (NodeT m [output]))
-> (TreeT m [input], TreeT m [output])
-> m (NodeT m [input], NodeT m [output])
forall (f :: * -> *) a c b d.
Applicative f =>
(a -> f c) -> (b -> f d) -> (a, b) -> f (c, d)
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 TreeT m [input] -> m (NodeT m [input])
forall (m :: * -> *) a. TreeT m a -> m (NodeT m a)
runTreeT TreeT m [output] -> m (NodeT m [output])
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 NodeT m [input] -> [input]
forall (m :: * -> *) a. NodeT m a -> a
nodeValue NodeT m [input]
as of
    [] -> ([input], [output])
-> [TreeT m ([input], [output])] -> NodeT m ([input], [output])
forall (m :: * -> *) a. a -> [TreeT m a] -> NodeT m a
NodeT ([], []) []
    [input]
_ ->
      ([input], [output])
-> [TreeT m ([input], [output])] -> NodeT m ([input], [output])
forall (m :: * -> *) a. a -> [TreeT m a] -> NodeT m a
NodeT
        (NodeT m [input] -> [input]
forall (m :: * -> *) a. NodeT m a -> a
nodeValue NodeT m [input]
as, NodeT m [output] -> [output]
forall (m :: * -> *) a. NodeT m a -> a
nodeValue NodeT m [output]
bs)
        ( (input -> Integer)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> (NodeT m [input], NodeT m [output])
-> [TreeT m ([input], [output])]
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)
            [TreeT m ([input], [output])]
-> [TreeT m ([input], [output])] -> [TreeT m ([input], [output])]
forall a. [a] -> [a] -> [a]
++ (input -> Integer)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> (NodeT m [input], NodeT m [output])
-> [TreeT m ([input], [output])]
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 <- NodeT m [output] -> [TreeT m [output]]
forall (m :: * -> *) a. NodeT m a -> [TreeT m a]
nodeChildren NodeT m [output]
ys1
  TreeT m ([input], [output]) -> [TreeT m ([input], [output])]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TreeT m ([input], [output]) -> [TreeT m ([input], [output])])
-> (m (NodeT m ([input], [output])) -> TreeT m ([input], [output]))
-> m (NodeT m ([input], [output]))
-> [TreeT m ([input], [output])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m (NodeT m ([input], [output])) -> TreeT m ([input], [output])
forall (m :: * -> *) a. m (NodeT m a) -> TreeT m a
TreeT (m (NodeT m ([input], [output])) -> [TreeT m ([input], [output])])
-> m (NodeT m ([input], [output])) -> [TreeT m ([input], [output])]
forall a b. (a -> b) -> a -> b
$ do
    NodeT m [output]
ys3 <- TreeT m [output] -> m (NodeT m [output])
forall (m :: * -> *) a. TreeT m a -> m (NodeT m a)
runTreeT TreeT m [output]
ys2
    NodeT m ([input], [output]) -> m (NodeT m ([input], [output]))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NodeT m ([input], [output]) -> m (NodeT m ([input], [output])))
-> NodeT m ([input], [output]) -> m (NodeT m ([input], [output]))
forall a b. (a -> b) -> a -> b
$ (input -> Integer)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> (NodeT m [input], NodeT m [output])
-> NodeT m ([input], [output])
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 <- NodeT m [input] -> [TreeT m [input]]
forall (m :: * -> *) a. NodeT m a -> [TreeT m a]
nodeChildren NodeT m [input]
xs1
  TreeT m ([input], [output]) -> [TreeT m ([input], [output])]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TreeT m ([input], [output]) -> [TreeT m ([input], [output])])
-> (m (NodeT m ([input], [output])) -> TreeT m ([input], [output]))
-> m (NodeT m ([input], [output]))
-> [TreeT m ([input], [output])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m (NodeT m ([input], [output])) -> TreeT m ([input], [output])
forall (m :: * -> *) a. m (NodeT m a) -> TreeT m a
TreeT (m (NodeT m ([input], [output])) -> [TreeT m ([input], [output])])
-> m (NodeT m ([input], [output])) -> [TreeT m ([input], [output])]
forall a b. (a -> b) -> a -> b
$ do
    NodeT m [input]
xs3 <- TreeT m [input] -> m (NodeT m [input])
forall (m :: * -> *) a. TreeT m a -> m (NodeT m a)
runTreeT TreeT m [input]
xs2
    let lost :: Integer
lost = [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (input -> Integer
inValue (input -> Integer) -> [input] -> [Integer]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NodeT m [input] -> [input]
forall (m :: * -> *) a. NodeT m a -> a
nodeValue NodeT m [input]
xs1) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (input -> Integer
inValue (input -> Integer) -> [input] -> [Integer]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NodeT m [input] -> [input]
forall (m :: * -> *) a. NodeT m a -> a
nodeValue NodeT m [input]
xs3)
        ys2 :: NodeT m [output]
ys2 = Integer
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> [output]
-> [output]
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 ([output] -> [output]) -> NodeT m [output] -> NodeT m [output]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NodeT m [output]
ys1
    NodeT m ([input], [output]) -> m (NodeT m ([input], [output]))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NodeT m ([input], [output]) -> m (NodeT m ([input], [output])))
-> NodeT m ([input], [output]) -> m (NodeT m ([input], [output]))
forall a b. (a -> b) -> a -> b
$ (input -> Integer)
-> (output -> Integer)
-> ((Integer -> Integer) -> output -> output)
-> (NodeT m [input], NodeT m [output])
-> NodeT m ([input], [output])
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 n -> n -> Bool
forall a. Ord a => a -> a -> Bool
> n
n'
        then (n -> n) -> a -> a
modifyVal (n -> n -> n
forall a. Num a => a -> a -> a
subtract n
n') a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
xs
        else n -> [a] -> [a]
go (n
n' n -> 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
PropertyT IO () -> Property
property (PropertyT IO () -> Property) -> PropertyT IO () -> Property
forall a b. (a -> b) -> a -> b
$ do
  Integer
n <- GenT Identity Integer -> PropertyT IO Integer
forall (m :: * -> *) a.
(Monad m, Show a, HasCallStack) =>
Gen a -> PropertyT m a
forAll (GenT Identity Integer -> PropertyT IO Integer)
-> GenT Identity Integer -> PropertyT IO Integer
forall a b. (a -> b) -> a -> b
$ Range Integer -> GenT Identity Integer
forall (m :: * -> *) a. (MonadGen m, Integral a) => Range a -> m a
Gen.integral (Integer -> Integer -> Range Integer
forall a. Integral a => a -> a -> Range a
Range.linear Integer
1 Integer
10000)
  [Integer]
ints <- Gen [Integer] -> PropertyT IO [Integer]
forall (m :: * -> *) a.
(Monad m, Show a, HasCallStack) =>
Gen a -> PropertyT m a
forAll (Gen [Integer] -> PropertyT IO [Integer])
-> Gen [Integer] -> PropertyT IO [Integer]
forall a b. (a -> b) -> a -> b
$ Integer
-> Gen (Integer -> Integer)
-> (Integer -> Integer)
-> ((Integer -> Integer) -> Integer -> Integer)
-> Gen [Integer]
forall a.
Integer
-> Gen (Integer -> a)
-> (a -> Integer)
-> ((Integer -> Integer) -> a -> a)
-> Gen [a]
genSplitValue Integer
n ((Integer -> Integer) -> Gen (Integer -> Integer)
forall a. a -> GenT Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Integer -> Integer
forall a. a -> a
id) Integer -> Integer
forall a. a -> a
id (Integer -> Integer) -> Integer -> Integer
forall a. a -> a
id
  [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Integer]
ints Integer -> Integer -> PropertyT IO ()
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 =
  (TreeT (MaybeT Identity) [TreeT (MaybeT Identity) a]
 -> TreeT (MaybeT Identity) [a])
-> GenT Identity [TreeT (MaybeT Identity) a] -> GenT Identity [a]
forall (m :: * -> *) a (n :: * -> *) b.
(TreeT (MaybeT m) a -> TreeT (MaybeT n) b) -> GenT m a -> GenT n b
mapGenT
    ( MaybeT Identity (NodeT (MaybeT Identity) [a])
-> TreeT (MaybeT Identity) [a]
forall (m :: * -> *) a. m (NodeT m a) -> TreeT m a
TreeT
        (MaybeT Identity (NodeT (MaybeT Identity) [a])
 -> TreeT (MaybeT Identity) [a])
-> (TreeT (MaybeT Identity) [TreeT (MaybeT Identity) a]
    -> MaybeT Identity (NodeT (MaybeT Identity) [a]))
-> TreeT (MaybeT Identity) [TreeT (MaybeT Identity) a]
-> TreeT (MaybeT Identity) [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a -> Integer)
-> ((Integer -> Integer) -> a -> a)
-> [TreeT (MaybeT Identity) a]
-> MaybeT Identity (NodeT (MaybeT Identity) [a])
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 ([TreeT (MaybeT Identity) a]
 -> MaybeT Identity (NodeT (MaybeT Identity) [a]))
-> (NodeT (MaybeT Identity) [TreeT (MaybeT Identity) a]
    -> [TreeT (MaybeT Identity) a])
-> NodeT (MaybeT Identity) [TreeT (MaybeT Identity) a]
-> MaybeT Identity (NodeT (MaybeT Identity) [a])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NodeT (MaybeT Identity) [TreeT (MaybeT Identity) a]
-> [TreeT (MaybeT Identity) a]
forall (m :: * -> *) a. NodeT m a -> a
nodeValue (NodeT (MaybeT Identity) [TreeT (MaybeT Identity) a]
 -> MaybeT Identity (NodeT (MaybeT Identity) [a]))
-> MaybeT
     Identity (NodeT (MaybeT Identity) [TreeT (MaybeT Identity) a])
-> MaybeT Identity (NodeT (MaybeT Identity) [a])
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<)
        (MaybeT
   Identity (NodeT (MaybeT Identity) [TreeT (MaybeT Identity) a])
 -> MaybeT Identity (NodeT (MaybeT Identity) [a]))
-> (TreeT (MaybeT Identity) [TreeT (MaybeT Identity) a]
    -> MaybeT
         Identity (NodeT (MaybeT Identity) [TreeT (MaybeT Identity) a]))
-> TreeT (MaybeT Identity) [TreeT (MaybeT Identity) a]
-> MaybeT Identity (NodeT (MaybeT Identity) [a])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TreeT (MaybeT Identity) [TreeT (MaybeT Identity) a]
-> MaybeT
     Identity (NodeT (MaybeT Identity) [TreeT (MaybeT Identity) a])
forall (m :: * -> *) a. TreeT m a -> m (NodeT m a)
runTreeT
    )
    (GenT Identity [TreeT (MaybeT Identity) a] -> GenT Identity [a])
-> GenT Identity [TreeT (MaybeT Identity) a] -> GenT Identity [a]
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 = [TreeT (MaybeT Identity) a]
-> GenT Identity [TreeT (MaybeT Identity) a]
forall a. a -> GenT Identity a
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 <- GenT Identity a
-> GenT Identity (TreeT (MaybeT (GenBase (GenT Identity))) a)
forall (m :: * -> *) a.
MonadGen m =>
m a -> m (TreeT (MaybeT (GenBase m)) a)
toTreeMaybeT (Gen (Integer -> a)
genA Gen (Integer -> a) -> GenT Identity Integer -> GenT Identity a
forall a b.
GenT Identity (a -> b) -> GenT Identity a -> GenT Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Range Integer -> GenT Identity Integer
forall (m :: * -> *) a. (MonadGen m, Integral a) => Range a -> m a
Gen.integral (Integer -> Integer -> Range Integer
forall a. a -> a -> Range a
Range.constant Integer
1 Integer
left))
      case Tree (Maybe a) -> Maybe a
forall a. Tree a -> a
treeValue (TreeT (MaybeT Identity) a -> Tree (Maybe a)
forall (m :: * -> *) a.
Monad m =>
TreeT (MaybeT m) a -> TreeT m (Maybe a)
runDiscardEffectT TreeT (MaybeT Identity) a
mTree) of
        Maybe a
Nothing -> GenT Identity [TreeT (MaybeT Identity) a]
forall a. GenT Identity a
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 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- a -> Integer
getValue a
a) (TreeT (MaybeT Identity) a
mTree TreeT (MaybeT Identity) a
-> [TreeT (MaybeT Identity) a] -> [TreeT (MaybeT Identity) a]
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 =
  ([NodeT m a] -> NodeT m [a]) -> m [NodeT m a] -> m (NodeT m [a])
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> Integer)
-> ((Integer -> Integer) -> a -> a) -> [NodeT m a] -> NodeT m [a]
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) (m [NodeT m a] -> m (NodeT m [a]))
-> ([TreeT m a] -> m [NodeT m a]) -> [TreeT m a] -> m (NodeT m [a])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TreeT m a -> m (NodeT m a)) -> [TreeT m a] -> m [NodeT m a]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse TreeT m a -> m (NodeT m a)
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 =
  [a] -> [TreeT m [a]] -> NodeT m [a]
forall (m :: * -> *) a. a -> [TreeT m a] -> NodeT m a
NodeT
    ((NodeT m a -> a) -> [NodeT m a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap NodeT m a -> a
forall (m :: * -> *) a. NodeT m a -> a
nodeValue [NodeT m a]
ts)
    ( (a -> Integer)
-> ((Integer -> Integer) -> a -> a) -> [NodeT m a] -> [TreeT m [a]]
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
        [TreeT m [a]] -> [TreeT m [a]] -> [TreeT m [a]]
forall a. [a] -> [a] -> [a]
++ (a -> Integer)
-> ((Integer -> Integer) -> a -> a) -> [NodeT m a] -> [TreeT m [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) <- [NodeT m a] -> [([NodeT m a], NodeT m a, NodeT m a, [NodeT m a])]
forall a. [a] -> [([a], a, a, [a])]
viewTwo [NodeT m a]
ts
  let x' :: NodeT m a
x' = (Integer -> Integer) -> a -> a
modifyValue (Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ a -> Integer
getValue (NodeT m a -> a
forall (m :: * -> *) a. NodeT m a -> a
nodeValue NodeT m a
y)) (a -> a) -> NodeT m a -> NodeT m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NodeT m a
x
  TreeT m [a] -> [TreeT m [a]]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TreeT m [a] -> [TreeT m [a]])
-> (NodeT m [a] -> TreeT m [a]) -> NodeT m [a] -> [TreeT m [a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m (NodeT m [a]) -> TreeT m [a]
forall (m :: * -> *) a. m (NodeT m a) -> TreeT m a
TreeT (m (NodeT m [a]) -> TreeT m [a])
-> (NodeT m [a] -> m (NodeT m [a])) -> NodeT m [a] -> TreeT m [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NodeT m [a] -> m (NodeT m [a])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NodeT m [a] -> [TreeT m [a]]) -> NodeT m [a] -> [TreeT m [a]]
forall a b. (a -> b) -> a -> b
$
    (a -> Integer)
-> ((Integer -> Integer) -> a -> a) -> [NodeT m a] -> NodeT m [a]
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 [NodeT m a] -> [NodeT m a] -> [NodeT m a]
forall a. [a] -> [a] -> [a]
++ [NodeT m a
x'] [NodeT m a] -> [NodeT m a] -> [NodeT m a]
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) <- [NodeT m a] -> [([NodeT m a], NodeT m a, NodeT m a, [NodeT m a])]
forall a. [a] -> [([a], a, a, [a])]
viewTwo [NodeT m a]
ts
  TreeT m a
y1 <- NodeT m a -> [TreeT m a]
forall (m :: * -> *) a. NodeT m a -> [TreeT m a]
nodeChildren NodeT m a
y0
  TreeT m [a] -> [TreeT m [a]]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TreeT m [a] -> [TreeT m [a]])
-> (m (NodeT m [a]) -> TreeT m [a])
-> m (NodeT m [a])
-> [TreeT m [a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m (NodeT m [a]) -> TreeT m [a]
forall (m :: * -> *) a. m (NodeT m a) -> TreeT m a
TreeT (m (NodeT m [a]) -> [TreeT m [a]])
-> m (NodeT m [a]) -> [TreeT m [a]]
forall a b. (a -> b) -> a -> b
$ do
    NodeT m a
y2 <- TreeT m a -> m (NodeT m a)
forall (m :: * -> *) a. TreeT m a -> m (NodeT m a)
runTreeT TreeT m a
y1
    let lost :: Integer
lost = a -> Integer
getValue (NodeT m a -> a
forall (m :: * -> *) a. NodeT m a -> a
nodeValue NodeT m a
y0) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- a -> Integer
getValue (NodeT m a -> a
forall (m :: * -> *) a. NodeT m a -> a
nodeValue NodeT m a
y2)
        x' :: NodeT m a
x' = (Integer -> Integer) -> a -> a
modifyValue (Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
lost) (a -> a) -> NodeT m a -> NodeT m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NodeT m a
x
    NodeT m [a] -> m (NodeT m [a])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NodeT m [a] -> m (NodeT m [a])) -> NodeT m [a] -> m (NodeT m [a])
forall a b. (a -> b) -> a -> b
$ (a -> Integer)
-> ((Integer -> Integer) -> a -> a) -> [NodeT m a] -> NodeT m [a]
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 [NodeT m a] -> [NodeT m a] -> [NodeT m a]
forall a. [a] -> [a] -> [a]
++ [NodeT m a
x', NodeT m a
y2] [NodeT m a] -> [NodeT m a] -> [NodeT m a]
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)
      ([a], a, a, [a]) -> [([a], a, a, [a])] -> [([a], a, a, [a])]
forall a. a -> [a] -> [a]
: (([a], a, a, [a]) -> ([a], a, a, [a]))
-> [([a], a, a, [a])] -> [([a], a, a, [a])]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\([a]
as, a
b, a
c, [a]
ds) -> (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
as, a
b, a
c, [a]
ds)) ([a] -> [([a], a, a, [a])]
forall a. [a] -> [([a], a, a, [a])]
viewTwo (a
x' a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
xs))