{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeApplications #-}

module Cardano.Wallet.CoinSelection.Internal.Balance.Gen
    ( genSelectionLimit
    , genSelectionSkeleton
    , genSelectionStrategy
    , shrinkSelectionLimit
    , shrinkSelectionSkeleton
    , shrinkSelectionStrategy
    )
    where

import Prelude

import Cardano.Wallet.CoinSelection.Internal.Balance
    ( SelectionLimit
    , SelectionLimitOf (..)
    , SelectionSkeleton (..)
    , SelectionStrategy (..)
    )
import Cardano.Wallet.CoinSelection.Internal.Context
    ( SelectionContext (..) )
import Cardano.Wallet.Primitive.Types.Coin
    ( Coin (..) )
import Cardano.Wallet.Primitive.Types.TokenBundle
    ( TokenBundle )
import Cardano.Wallet.Primitive.Types.TokenBundle.Gen
    ( genTokenBundleSmallRange, shrinkTokenBundleSmallRange )
import Cardano.Wallet.Primitive.Types.TokenMap.Gen
    ( genAssetId, shrinkAssetId )
import Generics.SOP
    ( NP (..) )
import Test.QuickCheck
    ( Gen
    , NonNegative (..)
    , arbitrary
    , arbitraryBoundedEnum
    , listOf
    , oneof
    , shrink
    , shrinkList
    , shrinkMapBy
    , suchThat
    )
import Test.QuickCheck.Extra
    ( genericRoundRobinShrink, (<:>), (<@>) )

import qualified Cardano.Wallet.Primitive.Types.TokenBundle as TokenBundle
import qualified Data.Set as Set

--------------------------------------------------------------------------------
-- Selection limits
--------------------------------------------------------------------------------

genSelectionLimit :: Gen SelectionLimit
genSelectionLimit :: Gen SelectionLimit
genSelectionLimit = [Gen SelectionLimit] -> Gen SelectionLimit
forall a. [Gen a] -> Gen a
oneof
    [ Int -> SelectionLimit
forall a. a -> SelectionLimitOf a
MaximumInputLimit (Int -> SelectionLimit)
-> (NonNegative Int -> Int) -> NonNegative Int -> SelectionLimit
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonNegative Int -> Int
forall a. NonNegative a -> a
getNonNegative (NonNegative Int -> SelectionLimit)
-> Gen (NonNegative Int) -> Gen SelectionLimit
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (NonNegative Int)
forall a. Arbitrary a => Gen a
arbitrary
    , SelectionLimit -> Gen SelectionLimit
forall (f :: * -> *) a. Applicative f => a -> f a
pure SelectionLimit
forall a. SelectionLimitOf a
NoLimit
    ]

shrinkSelectionLimit :: SelectionLimit -> [SelectionLimit]
shrinkSelectionLimit :: SelectionLimit -> [SelectionLimit]
shrinkSelectionLimit = \case
    MaximumInputLimit Int
n ->
        Int -> SelectionLimit
forall a. a -> SelectionLimitOf a
MaximumInputLimit (Int -> SelectionLimit)
-> (NonNegative Int -> Int) -> NonNegative Int -> SelectionLimit
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonNegative Int -> Int
forall a. NonNegative a -> a
getNonNegative (NonNegative Int -> SelectionLimit)
-> [NonNegative Int] -> [SelectionLimit]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NonNegative Int -> [NonNegative Int]
forall a. Arbitrary a => a -> [a]
shrink (Int -> NonNegative Int
forall a. a -> NonNegative a
NonNegative Int
n)
    SelectionLimit
NoLimit ->
        []

--------------------------------------------------------------------------------
-- Selection skeletons
--------------------------------------------------------------------------------

genSelectionSkeleton :: Gen (Address ctx) -> Gen (SelectionSkeleton ctx)
genSelectionSkeleton :: Gen (Address ctx) -> Gen (SelectionSkeleton ctx)
genSelectionSkeleton Gen (Address ctx)
genAddress = Int
-> [(Address ctx, TokenBundle)]
-> [Set AssetId]
-> SelectionSkeleton ctx
forall ctx.
Int
-> [(Address ctx, TokenBundle)]
-> [Set AssetId]
-> SelectionSkeleton ctx
SelectionSkeleton
    (Int
 -> [(Address ctx, TokenBundle)]
 -> [Set AssetId]
 -> SelectionSkeleton ctx)
-> Gen Int
-> Gen
     ([(Address ctx, TokenBundle)]
      -> [Set AssetId] -> SelectionSkeleton ctx)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen Int
genSkeletonInputCount
    Gen
  ([(Address ctx, TokenBundle)]
   -> [Set AssetId] -> SelectionSkeleton ctx)
-> Gen [(Address ctx, TokenBundle)]
-> Gen ([Set AssetId] -> SelectionSkeleton ctx)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen [(Address ctx, TokenBundle)]
genSkeletonOutputs
    Gen ([Set AssetId] -> SelectionSkeleton ctx)
-> Gen [Set AssetId] -> Gen (SelectionSkeleton ctx)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen [Set AssetId]
genSkeletonChange
  where
    genSkeletonInputCount :: Gen Int
genSkeletonInputCount =
        NonNegative Int -> Int
forall a. NonNegative a -> a
getNonNegative (NonNegative Int -> Int) -> Gen (NonNegative Int) -> Gen Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Arbitrary (NonNegative Int) => Gen (NonNegative Int)
forall a. Arbitrary a => Gen a
arbitrary @(NonNegative Int)
    genSkeletonOutputs :: Gen [(Address ctx, TokenBundle)]
genSkeletonOutputs =
        Gen (Address ctx, TokenBundle) -> Gen [(Address ctx, TokenBundle)]
forall a. Gen a -> Gen [a]
listOf Gen (Address ctx, TokenBundle)
genSkeletonOutput
    genSkeletonOutput :: Gen (Address ctx, TokenBundle)
genSkeletonOutput = (,)
        (Address ctx -> TokenBundle -> (Address ctx, TokenBundle))
-> Gen (Address ctx)
-> Gen (TokenBundle -> (Address ctx, TokenBundle))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (Address ctx)
genAddress
        Gen (TokenBundle -> (Address ctx, TokenBundle))
-> Gen TokenBundle -> Gen (Address ctx, TokenBundle)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen TokenBundle
genTokenBundleSmallRange Gen TokenBundle -> (TokenBundle -> Bool) -> Gen TokenBundle
forall a. Gen a -> (a -> Bool) -> Gen a
`suchThat` TokenBundle -> Bool
tokenBundleHasNonZeroCoin
    genSkeletonChange :: Gen [Set AssetId]
genSkeletonChange =
        Gen (Set AssetId) -> Gen [Set AssetId]
forall a. Gen a -> Gen [a]
listOf ([AssetId] -> Set AssetId
forall a. Ord a => [a] -> Set a
Set.fromList ([AssetId] -> Set AssetId) -> Gen [AssetId] -> Gen (Set AssetId)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen AssetId -> Gen [AssetId]
forall a. Gen a -> Gen [a]
listOf Gen AssetId
genAssetId)

shrinkSelectionSkeleton
    :: (Address ctx -> [Address ctx])
    -> (SelectionSkeleton ctx -> [SelectionSkeleton ctx])
shrinkSelectionSkeleton :: (Address ctx -> [Address ctx])
-> SelectionSkeleton ctx -> [SelectionSkeleton ctx]
shrinkSelectionSkeleton Address ctx -> [Address ctx]
shrinkAddress = NP (I -.-> []) '[Int, [(Address ctx, TokenBundle)], [Set AssetId]]
-> SelectionSkeleton ctx -> [SelectionSkeleton ctx]
forall a (xs :: [*]).
(Generic a, GFrom a, GTo a, GCode a ~ '[xs]) =>
NP (I -.-> []) xs -> a -> [a]
genericRoundRobinShrink
    (NP (I -.-> []) '[Int, [(Address ctx, TokenBundle)], [Set AssetId]]
 -> SelectionSkeleton ctx -> [SelectionSkeleton ctx])
-> NP
     (I -.-> []) '[Int, [(Address ctx, TokenBundle)], [Set AssetId]]
-> SelectionSkeleton ctx
-> [SelectionSkeleton ctx]
forall a b. (a -> b) -> a -> b
<@> Int -> [Int]
shrinkSkeletonInputCount
    (Int -> [Int])
-> NP (I -.-> []) '[[(Address ctx, TokenBundle)], [Set AssetId]]
-> NP
     (I -.-> []) '[Int, [(Address ctx, TokenBundle)], [Set AssetId]]
forall x (xs :: [*]).
(x -> [x]) -> NP (I -.-> []) xs -> NP (I -.-> []) (x : xs)
<:> [(Address ctx, TokenBundle)] -> [[(Address ctx, TokenBundle)]]
shrinkSkeletonOutputs
    ([(Address ctx, TokenBundle)] -> [[(Address ctx, TokenBundle)]])
-> NP (I -.-> []) '[[Set AssetId]]
-> NP (I -.-> []) '[[(Address ctx, TokenBundle)], [Set AssetId]]
forall x (xs :: [*]).
(x -> [x]) -> NP (I -.-> []) xs -> NP (I -.-> []) (x : xs)
<:> [Set AssetId] -> [[Set AssetId]]
shrinkSkeletonChange
    ([Set AssetId] -> [[Set AssetId]])
-> NP (I -.-> []) '[] -> NP (I -.-> []) '[[Set AssetId]]
forall x (xs :: [*]).
(x -> [x]) -> NP (I -.-> []) xs -> NP (I -.-> []) (x : xs)
<:> NP (I -.-> []) '[]
forall k (a :: k -> *). NP a '[]
Nil
  where
    shrinkSkeletonInputCount :: Int -> [Int]
shrinkSkeletonInputCount =
        Arbitrary Int => Int -> [Int]
forall a. Arbitrary a => a -> [a]
shrink @Int
    shrinkSkeletonOutputs :: [(Address ctx, TokenBundle)] -> [[(Address ctx, TokenBundle)]]
shrinkSkeletonOutputs =
        ((Address ctx, TokenBundle) -> [(Address ctx, TokenBundle)])
-> [(Address ctx, TokenBundle)] -> [[(Address ctx, TokenBundle)]]
forall a. (a -> [a]) -> [a] -> [[a]]
shrinkList (Address ctx, TokenBundle) -> [(Address ctx, TokenBundle)]
shrinkSkeletonOutput
    shrinkSkeletonOutput :: (Address ctx, TokenBundle) -> [(Address ctx, TokenBundle)]
shrinkSkeletonOutput =
        NP (I -.-> []) '[Address ctx, TokenBundle]
-> (Address ctx, TokenBundle) -> [(Address ctx, TokenBundle)]
forall a (xs :: [*]).
(Generic a, GFrom a, GTo a, GCode a ~ '[xs]) =>
NP (I -.-> []) xs -> a -> [a]
genericRoundRobinShrink
            (NP (I -.-> []) '[Address ctx, TokenBundle]
 -> (Address ctx, TokenBundle) -> [(Address ctx, TokenBundle)])
-> NP (I -.-> []) '[Address ctx, TokenBundle]
-> (Address ctx, TokenBundle)
-> [(Address ctx, TokenBundle)]
forall a b. (a -> b) -> a -> b
<@> Address ctx -> [Address ctx]
shrinkAddress
            (Address ctx -> [Address ctx])
-> NP (I -.-> []) '[TokenBundle]
-> NP (I -.-> []) '[Address ctx, TokenBundle]
forall x (xs :: [*]).
(x -> [x]) -> NP (I -.-> []) xs -> NP (I -.-> []) (x : xs)
<:> (TokenBundle -> Bool) -> [TokenBundle] -> [TokenBundle]
forall a. (a -> Bool) -> [a] -> [a]
filter TokenBundle -> Bool
tokenBundleHasNonZeroCoin ([TokenBundle] -> [TokenBundle])
-> (TokenBundle -> [TokenBundle]) -> TokenBundle -> [TokenBundle]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TokenBundle -> [TokenBundle]
shrinkTokenBundleSmallRange
            (TokenBundle -> [TokenBundle])
-> NP (I -.-> []) '[] -> NP (I -.-> []) '[TokenBundle]
forall x (xs :: [*]).
(x -> [x]) -> NP (I -.-> []) xs -> NP (I -.-> []) (x : xs)
<:> NP (I -.-> []) '[]
forall k (a :: k -> *). NP a '[]
Nil
    shrinkSkeletonChange :: [Set AssetId] -> [[Set AssetId]]
shrinkSkeletonChange =
        (Set AssetId -> [Set AssetId]) -> [Set AssetId] -> [[Set AssetId]]
forall a. (a -> [a]) -> [a] -> [[a]]
shrinkList ((Set AssetId -> [Set AssetId])
 -> [Set AssetId] -> [[Set AssetId]])
-> (Set AssetId -> [Set AssetId])
-> [Set AssetId]
-> [[Set AssetId]]
forall a b. (a -> b) -> a -> b
$
        ([AssetId] -> Set AssetId)
-> (Set AssetId -> [AssetId])
-> ([AssetId] -> [[AssetId]])
-> Set AssetId
-> [Set AssetId]
forall a b. (a -> b) -> (b -> a) -> (a -> [a]) -> b -> [b]
shrinkMapBy [AssetId] -> Set AssetId
forall a. Ord a => [a] -> Set a
Set.fromList Set AssetId -> [AssetId]
forall a. Set a -> [a]
Set.toList ((AssetId -> [AssetId]) -> [AssetId] -> [[AssetId]]
forall a. (a -> [a]) -> [a] -> [[a]]
shrinkList AssetId -> [AssetId]
shrinkAssetId)

tokenBundleHasNonZeroCoin :: TokenBundle -> Bool
tokenBundleHasNonZeroCoin :: TokenBundle -> Bool
tokenBundleHasNonZeroCoin TokenBundle
b = TokenBundle -> Coin
TokenBundle.getCoin TokenBundle
b Coin -> Coin -> Bool
forall a. Eq a => a -> a -> Bool
/= Natural -> Coin
Coin Natural
0

--------------------------------------------------------------------------------
-- Selection strategies
--------------------------------------------------------------------------------

genSelectionStrategy :: Gen SelectionStrategy
genSelectionStrategy :: Gen SelectionStrategy
genSelectionStrategy = Gen SelectionStrategy
forall a. (Bounded a, Enum a) => Gen a
arbitraryBoundedEnum

shrinkSelectionStrategy :: SelectionStrategy -> [SelectionStrategy]
shrinkSelectionStrategy :: SelectionStrategy -> [SelectionStrategy]
shrinkSelectionStrategy = \case
    -- Shrinking from "optimal" to "minimal" should increase the likelihood of
    -- making a successful selection, as the "minimal" strategy is designed to
    -- generate smaller selections.
    SelectionStrategy
SelectionStrategyMinimal -> []
    SelectionStrategy
SelectionStrategyOptimal -> [SelectionStrategy
SelectionStrategyMinimal]