module Cardano.Wallet.Primitive.Types.TokenBundle.Gen
    ( genTokenBundleSmallRange
    , genTokenBundleSmallRangePositive
    , genTokenBundle
    , shrinkTokenBundle
    , shrinkTokenBundleSmallRange
    , shrinkTokenBundleSmallRangePositive
    , genTokenBundlePartition
    , genTokenBundlePartitionNonNull
    ) where

import Prelude

import Cardano.Wallet.Primitive.Types.Coin.Gen
    ( genCoin
    , genCoinPartition
    , genCoinPositive
    , shrinkCoin
    , shrinkCoinPositive
    )
import Cardano.Wallet.Primitive.Types.TokenBundle
    ( TokenBundle (..) )
import Cardano.Wallet.Primitive.Types.TokenMap.Gen
    ( genTokenMap
    , genTokenMapPartition
    , genTokenMapSmallRange
    , shrinkTokenMap
    )
import Data.List.NonEmpty
    ( NonEmpty )
import Test.QuickCheck
    ( Gen )
import Test.QuickCheck.Extra
    ( shrinkInterleaved )

import qualified Data.Foldable as F
import qualified Data.List.NonEmpty as NE

--------------------------------------------------------------------------------
-- Token bundles with variable numbers of assets, the upper bound being
-- QuickCheck's size parameter.
--
-- Policy identifiers, asset names, token quantities are all allowed to vary.
--------------------------------------------------------------------------------

genTokenBundle :: Gen TokenBundle
genTokenBundle :: Gen TokenBundle
genTokenBundle = Coin -> TokenMap -> TokenBundle
TokenBundle
    (Coin -> TokenMap -> TokenBundle)
-> Gen Coin -> Gen (TokenMap -> TokenBundle)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Coin
genCoin
    Gen (TokenMap -> TokenBundle) -> Gen TokenMap -> Gen TokenBundle
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen TokenMap
genTokenMap

shrinkTokenBundle :: TokenBundle -> [TokenBundle]
shrinkTokenBundle :: TokenBundle -> [TokenBundle]
shrinkTokenBundle (TokenBundle Coin
c TokenMap
m)=
    (Coin -> TokenMap -> TokenBundle)
-> (Coin, TokenMap) -> TokenBundle
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Coin -> TokenMap -> TokenBundle
TokenBundle ((Coin, TokenMap) -> TokenBundle)
-> [(Coin, TokenMap)] -> [TokenBundle]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Coin, Coin -> [Coin])
-> (TokenMap, TokenMap -> [TokenMap]) -> [(Coin, TokenMap)]
forall a b. (a, a -> [a]) -> (b, b -> [b]) -> [(a, b)]
shrinkInterleaved
        (Coin
c, Coin -> [Coin]
shrinkCoin)
        (TokenMap
m, TokenMap -> [TokenMap]
shrinkTokenMap)

--------------------------------------------------------------------------------
-- Token bundles with coins, assets, and quantities chosen from small ranges
--------------------------------------------------------------------------------

genTokenBundleSmallRange :: Gen TokenBundle
genTokenBundleSmallRange :: Gen TokenBundle
genTokenBundleSmallRange = Coin -> TokenMap -> TokenBundle
TokenBundle
    (Coin -> TokenMap -> TokenBundle)
-> Gen Coin -> Gen (TokenMap -> TokenBundle)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Coin
genCoin
    Gen (TokenMap -> TokenBundle) -> Gen TokenMap -> Gen TokenBundle
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen TokenMap
genTokenMapSmallRange

shrinkTokenBundleSmallRange :: TokenBundle -> [TokenBundle]
shrinkTokenBundleSmallRange :: TokenBundle -> [TokenBundle]
shrinkTokenBundleSmallRange = TokenBundle -> [TokenBundle]
shrinkTokenBundle

genTokenBundleSmallRangePositive :: Gen TokenBundle
genTokenBundleSmallRangePositive :: Gen TokenBundle
genTokenBundleSmallRangePositive = Coin -> TokenMap -> TokenBundle
TokenBundle
    (Coin -> TokenMap -> TokenBundle)
-> Gen Coin -> Gen (TokenMap -> TokenBundle)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Coin
genCoinPositive
    Gen (TokenMap -> TokenBundle) -> Gen TokenMap -> Gen TokenBundle
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen TokenMap
genTokenMapSmallRange

shrinkTokenBundleSmallRangePositive :: TokenBundle -> [TokenBundle]
shrinkTokenBundleSmallRangePositive :: TokenBundle -> [TokenBundle]
shrinkTokenBundleSmallRangePositive (TokenBundle Coin
c TokenMap
m) =
    (Coin -> TokenMap -> TokenBundle)
-> (Coin, TokenMap) -> TokenBundle
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Coin -> TokenMap -> TokenBundle
TokenBundle ((Coin, TokenMap) -> TokenBundle)
-> [(Coin, TokenMap)] -> [TokenBundle]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Coin, Coin -> [Coin])
-> (TokenMap, TokenMap -> [TokenMap]) -> [(Coin, TokenMap)]
forall a b. (a, a -> [a]) -> (b, b -> [b]) -> [(a, b)]
shrinkInterleaved
        (Coin
c, Coin -> [Coin]
shrinkCoinPositive)
        (TokenMap
m, TokenMap -> [TokenMap]
shrinkTokenMap)

--------------------------------------------------------------------------------
-- Partitioning token bundles
--------------------------------------------------------------------------------

-- | Partitions a token bundle randomly into a given number of parts.
--
-- Satisfies the following properties:
--
-- prop> forAll (genTokenBundlePartition b i) $ (==       b) . fold
-- prop> forAll (genTokenBundlePartition b i) $ (== max 1 i) . length
--
genTokenBundlePartition :: TokenBundle -> Int -> Gen (NonEmpty TokenBundle)
genTokenBundlePartition :: TokenBundle -> Int -> Gen (NonEmpty TokenBundle)
genTokenBundlePartition (TokenBundle Coin
c TokenMap
m) Int
i =
    (Coin -> TokenMap -> TokenBundle)
-> NonEmpty Coin -> NonEmpty TokenMap -> NonEmpty TokenBundle
forall a b c.
(a -> b -> c) -> NonEmpty a -> NonEmpty b -> NonEmpty c
NE.zipWith Coin -> TokenMap -> TokenBundle
TokenBundle
        (NonEmpty Coin -> NonEmpty TokenMap -> NonEmpty TokenBundle)
-> Gen (NonEmpty Coin)
-> Gen (NonEmpty TokenMap -> NonEmpty TokenBundle)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Coin -> Int -> Gen (NonEmpty Coin)
genCoinPartition     Coin
c Int
i
        Gen (NonEmpty TokenMap -> NonEmpty TokenBundle)
-> Gen (NonEmpty TokenMap) -> Gen (NonEmpty TokenBundle)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TokenMap -> Int -> Gen (NonEmpty TokenMap)
genTokenMapPartition TokenMap
m Int
i

-- | Like 'genTokenBundlePartition', but with empty values removed from the
--   result.
--
genTokenBundlePartitionNonNull :: TokenBundle -> Int -> Gen [TokenBundle]
genTokenBundlePartitionNonNull :: TokenBundle -> Int -> Gen [TokenBundle]
genTokenBundlePartitionNonNull TokenBundle
m Int
i =
    (TokenBundle -> Bool) -> [TokenBundle] -> [TokenBundle]
forall a. (a -> Bool) -> [a] -> [a]
filter (TokenBundle -> TokenBundle -> Bool
forall a. Eq a => a -> a -> Bool
/= TokenBundle
forall a. Monoid a => a
mempty) ([TokenBundle] -> [TokenBundle])
-> (NonEmpty TokenBundle -> [TokenBundle])
-> NonEmpty TokenBundle
-> [TokenBundle]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty TokenBundle -> [TokenBundle]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList (NonEmpty TokenBundle -> [TokenBundle])
-> Gen (NonEmpty TokenBundle) -> Gen [TokenBundle]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TokenBundle -> Int -> Gen (NonEmpty TokenBundle)
genTokenBundlePartition TokenBundle
m Int
i