{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE NamedFieldPuns #-}

module Cardano.Wallet.Primitive.Types.TokenMap.Gen
    ( genAssetId
    , genAssetIdLargeRange
    , genTokenMap
    , genTokenMapSmallRange
    , shrinkAssetId
    , shrinkTokenMap
    , AssetIdF (..)
    , genTokenMapPartition
    , genTokenMapPartitionNonNull
    ) where

import Prelude

import Cardano.Wallet.Primitive.Types.TokenMap
    ( AssetId (..), TokenMap )
import Cardano.Wallet.Primitive.Types.TokenPolicy.Gen
    ( genTokenName
    , genTokenNameLargeRange
    , genTokenPolicyId
    , genTokenPolicyIdLargeRange
    , shrinkTokenName
    , shrinkTokenPolicyId
    , testTokenNames
    , testTokenPolicyIds
    )
import Cardano.Wallet.Primitive.Types.TokenQuantity
    ( TokenQuantity (..) )
import Cardano.Wallet.Primitive.Types.TokenQuantity.Gen
    ( genTokenQuantity, genTokenQuantityPartition, shrinkTokenQuantity )
import Control.Monad
    ( replicateM )
import Data.List
    ( elemIndex, transpose )
import Data.List.NonEmpty
    ( NonEmpty )
import Data.Maybe
    ( fromMaybe )
import GHC.Generics
    ( Generic )
import Safe
    ( fromJustNote )
import Test.QuickCheck
    ( CoArbitrary (..)
    , Function (..)
    , Gen
    , choose
    , functionMap
    , oneof
    , shrinkList
    , sized
    , variant
    )
import Test.QuickCheck.Extra
    ( genSized2With, shrinkInterleaved )

import qualified Cardano.Wallet.Primitive.Types.TokenMap as TokenMap
import qualified Data.Foldable as F
import qualified Data.List.NonEmpty as NE

--------------------------------------------------------------------------------
-- Asset identifiers chosen from a range that depends on the size parameter
--------------------------------------------------------------------------------

genAssetId :: Gen AssetId
genAssetId :: Gen AssetId
genAssetId = (TokenPolicyId -> TokenName -> AssetId)
-> Gen TokenPolicyId -> Gen TokenName -> Gen AssetId
forall a b c. (a -> b -> c) -> Gen a -> Gen b -> Gen c
genSized2With TokenPolicyId -> TokenName -> AssetId
AssetId Gen TokenPolicyId
genTokenPolicyId Gen TokenName
genTokenName

shrinkAssetId :: AssetId -> [AssetId]
shrinkAssetId :: AssetId -> [AssetId]
shrinkAssetId (AssetId TokenPolicyId
p TokenName
t) = (TokenPolicyId -> TokenName -> AssetId)
-> (TokenPolicyId, TokenName) -> AssetId
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry TokenPolicyId -> TokenName -> AssetId
AssetId ((TokenPolicyId, TokenName) -> AssetId)
-> [(TokenPolicyId, TokenName)] -> [AssetId]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TokenPolicyId, TokenPolicyId -> [TokenPolicyId])
-> (TokenName, TokenName -> [TokenName])
-> [(TokenPolicyId, TokenName)]
forall a b. (a, a -> [a]) -> (b, b -> [b]) -> [(a, b)]
shrinkInterleaved
    (TokenPolicyId
p, TokenPolicyId -> [TokenPolicyId]
shrinkTokenPolicyId)
    (TokenName
t, TokenName -> [TokenName]
shrinkTokenName)

--------------------------------------------------------------------------------
-- Asset identifiers chosen from a large range (to minimize collisions)
--------------------------------------------------------------------------------

genAssetIdLargeRange :: Gen AssetId
genAssetIdLargeRange :: Gen AssetId
genAssetIdLargeRange = TokenPolicyId -> TokenName -> AssetId
AssetId
    (TokenPolicyId -> TokenName -> AssetId)
-> Gen TokenPolicyId -> Gen (TokenName -> AssetId)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen TokenPolicyId
genTokenPolicyIdLargeRange
    Gen (TokenName -> AssetId) -> Gen TokenName -> Gen AssetId
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen TokenName
genTokenNameLargeRange

--------------------------------------------------------------------------------
-- Token maps with assets and quantities chosen from ranges that depend on the
-- size parameter
--------------------------------------------------------------------------------

genTokenMap :: Gen TokenMap
genTokenMap :: Gen TokenMap
genTokenMap = (Int -> Gen TokenMap) -> Gen TokenMap
forall a. (Int -> Gen a) -> Gen a
sized ((Int -> Gen TokenMap) -> Gen TokenMap)
-> (Int -> Gen TokenMap) -> Gen TokenMap
forall a b. (a -> b) -> a -> b
$ \Int
size -> do
    Int
assetCount <- (Int, Int) -> Gen Int
forall a. Random a => (a, a) -> Gen a
choose (Int
0, Int
size)
    [(AssetId, TokenQuantity)] -> TokenMap
TokenMap.fromFlatList ([(AssetId, TokenQuantity)] -> TokenMap)
-> Gen [(AssetId, TokenQuantity)] -> Gen TokenMap
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int
-> Gen (AssetId, TokenQuantity) -> Gen [(AssetId, TokenQuantity)]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
assetCount Gen (AssetId, TokenQuantity)
genAssetQuantity
  where
    genAssetQuantity :: Gen (AssetId, TokenQuantity)
genAssetQuantity = (,)
        (AssetId -> TokenQuantity -> (AssetId, TokenQuantity))
-> Gen AssetId -> Gen (TokenQuantity -> (AssetId, TokenQuantity))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen AssetId
genAssetId
        Gen (TokenQuantity -> (AssetId, TokenQuantity))
-> Gen TokenQuantity -> Gen (AssetId, TokenQuantity)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen TokenQuantity
genTokenQuantity

--------------------------------------------------------------------------------
-- Token maps with assets and quantities chosen from small ranges
--------------------------------------------------------------------------------

genTokenMapSmallRange :: Gen TokenMap
genTokenMapSmallRange :: Gen TokenMap
genTokenMapSmallRange = do
    Int
assetCount <- [Gen Int] -> Gen Int
forall a. [Gen a] -> Gen a
oneof
        [ Int -> Gen Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
0
        , Int -> Gen Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
1
        , (Int, Int) -> Gen Int
forall a. Random a => (a, a) -> Gen a
choose (Int
2, Int
16)
        ]
    [(AssetId, TokenQuantity)] -> TokenMap
TokenMap.fromFlatList ([(AssetId, TokenQuantity)] -> TokenMap)
-> Gen [(AssetId, TokenQuantity)] -> Gen TokenMap
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int
-> Gen (AssetId, TokenQuantity) -> Gen [(AssetId, TokenQuantity)]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
assetCount Gen (AssetId, TokenQuantity)
genAssetQuantity
  where
    genAssetQuantity :: Gen (AssetId, TokenQuantity)
genAssetQuantity = (,)
        (AssetId -> TokenQuantity -> (AssetId, TokenQuantity))
-> Gen AssetId -> Gen (TokenQuantity -> (AssetId, TokenQuantity))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen AssetId
genAssetId
        Gen (TokenQuantity -> (AssetId, TokenQuantity))
-> Gen TokenQuantity -> Gen (AssetId, TokenQuantity)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen TokenQuantity
genTokenQuantity

shrinkTokenMap :: TokenMap -> [TokenMap]
shrinkTokenMap :: TokenMap -> [TokenMap]
shrinkTokenMap
    = ([(AssetId, TokenQuantity)] -> TokenMap)
-> [[(AssetId, TokenQuantity)]] -> [TokenMap]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [(AssetId, TokenQuantity)] -> TokenMap
TokenMap.fromFlatList
    ([[(AssetId, TokenQuantity)]] -> [TokenMap])
-> (TokenMap -> [[(AssetId, TokenQuantity)]])
-> TokenMap
-> [TokenMap]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((AssetId, TokenQuantity) -> [(AssetId, TokenQuantity)])
-> [(AssetId, TokenQuantity)] -> [[(AssetId, TokenQuantity)]]
forall a. (a -> [a]) -> [a] -> [[a]]
shrinkList (AssetId, TokenQuantity) -> [(AssetId, TokenQuantity)]
shrinkAssetQuantity
    ([(AssetId, TokenQuantity)] -> [[(AssetId, TokenQuantity)]])
-> (TokenMap -> [(AssetId, TokenQuantity)])
-> TokenMap
-> [[(AssetId, TokenQuantity)]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TokenMap -> [(AssetId, TokenQuantity)]
TokenMap.toFlatList
  where
    shrinkAssetQuantity :: (AssetId, TokenQuantity) -> [(AssetId, TokenQuantity)]
shrinkAssetQuantity (AssetId
a, TokenQuantity
q) = (AssetId, AssetId -> [AssetId])
-> (TokenQuantity, TokenQuantity -> [TokenQuantity])
-> [(AssetId, TokenQuantity)]
forall a b. (a, a -> [a]) -> (b, b -> [b]) -> [(a, b)]
shrinkInterleaved
        (AssetId
a, AssetId -> [AssetId]
shrinkAssetId)
        (TokenQuantity
q, TokenQuantity -> [TokenQuantity]
shrinkTokenQuantity)

--------------------------------------------------------------------------------
-- Filtering functions
--------------------------------------------------------------------------------

newtype AssetIdF = AssetIdF AssetId
    deriving ((forall x. AssetIdF -> Rep AssetIdF x)
-> (forall x. Rep AssetIdF x -> AssetIdF) -> Generic AssetIdF
forall x. Rep AssetIdF x -> AssetIdF
forall x. AssetIdF -> Rep AssetIdF x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep AssetIdF x -> AssetIdF
$cfrom :: forall x. AssetIdF -> Rep AssetIdF x
Generic, AssetIdF -> AssetIdF -> Bool
(AssetIdF -> AssetIdF -> Bool)
-> (AssetIdF -> AssetIdF -> Bool) -> Eq AssetIdF
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: AssetIdF -> AssetIdF -> Bool
$c/= :: AssetIdF -> AssetIdF -> Bool
== :: AssetIdF -> AssetIdF -> Bool
$c== :: AssetIdF -> AssetIdF -> Bool
Eq, Int -> AssetIdF -> ShowS
[AssetIdF] -> ShowS
AssetIdF -> String
(Int -> AssetIdF -> ShowS)
-> (AssetIdF -> String) -> ([AssetIdF] -> ShowS) -> Show AssetIdF
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [AssetIdF] -> ShowS
$cshowList :: [AssetIdF] -> ShowS
show :: AssetIdF -> String
$cshow :: AssetIdF -> String
showsPrec :: Int -> AssetIdF -> ShowS
$cshowsPrec :: Int -> AssetIdF -> ShowS
Show, ReadPrec [AssetIdF]
ReadPrec AssetIdF
Int -> ReadS AssetIdF
ReadS [AssetIdF]
(Int -> ReadS AssetIdF)
-> ReadS [AssetIdF]
-> ReadPrec AssetIdF
-> ReadPrec [AssetIdF]
-> Read AssetIdF
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
readListPrec :: ReadPrec [AssetIdF]
$creadListPrec :: ReadPrec [AssetIdF]
readPrec :: ReadPrec AssetIdF
$creadPrec :: ReadPrec AssetIdF
readList :: ReadS [AssetIdF]
$creadList :: ReadS [AssetIdF]
readsPrec :: Int -> ReadS AssetIdF
$creadsPrec :: Int -> ReadS AssetIdF
Read)

instance Function AssetIdF where
    function :: (AssetIdF -> b) -> AssetIdF :-> b
function = (AssetIdF -> String)
-> (String -> AssetIdF) -> (AssetIdF -> b) -> AssetIdF :-> b
forall b a c.
Function b =>
(a -> b) -> (b -> a) -> (a -> c) -> a :-> c
functionMap AssetIdF -> String
forall a. Show a => a -> String
show String -> AssetIdF
forall a. Read a => String -> a
read

instance CoArbitrary AssetIdF where
    coarbitrary :: AssetIdF -> Gen b -> Gen b
coarbitrary (AssetIdF AssetId{TokenName
$sel:tokenName:AssetId :: AssetId -> TokenName
tokenName :: TokenName
tokenName, TokenPolicyId
$sel:tokenPolicyId:AssetId :: AssetId -> TokenPolicyId
tokenPolicyId :: TokenPolicyId
tokenPolicyId}) Gen b
genB = do
        let n :: Int
n = Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe Int
0 (TokenName -> [TokenName] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
elemIndex TokenName
tokenName [TokenName]
testTokenNames)
        let m :: Int
m = Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe Int
0 (TokenPolicyId -> [TokenPolicyId] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
elemIndex TokenPolicyId
tokenPolicyId [TokenPolicyId]
testTokenPolicyIds)
        Int -> Gen b -> Gen b
forall n a. Integral n => n -> Gen a -> Gen a
variant (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
m) Gen b
genB

--------------------------------------------------------------------------------
-- Partitioning token maps
--------------------------------------------------------------------------------

-- | Partitions a token map randomly into a given number of parts.
--
-- Satisfies the following properties:
--
-- prop> forAll (genTokenMapPartition m i) $ (== m      ) . fold
-- prop> forAll (genTokenMapPartition m i) $ (== max 1 i) . length
--
genTokenMapPartition :: TokenMap -> Int -> Gen (NonEmpty TokenMap)
genTokenMapPartition :: TokenMap -> Int -> Gen (NonEmpty TokenMap)
genTokenMapPartition TokenMap
m Int
i
    | TokenMap -> Bool
TokenMap.isEmpty TokenMap
m =
        NonEmpty TokenMap -> Gen (NonEmpty TokenMap)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NonEmpty TokenMap -> Gen (NonEmpty TokenMap))
-> NonEmpty TokenMap -> Gen (NonEmpty TokenMap)
forall a b. (a -> b) -> a -> b
$ [TokenMap] -> NonEmpty TokenMap
forall a. [a] -> NonEmpty a
NE.fromList ([TokenMap] -> NonEmpty TokenMap)
-> [TokenMap] -> NonEmpty TokenMap
forall a b. (a -> b) -> a -> b
$ Int -> TokenMap -> [TokenMap]
forall a. Int -> a -> [a]
replicate (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 Int
i) TokenMap
forall a. Monoid a => a
mempty
    | Bool
otherwise =
        ([(AssetId, TokenQuantity)] -> TokenMap)
-> NonEmpty [(AssetId, TokenQuantity)] -> NonEmpty TokenMap
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [(AssetId, TokenQuantity)] -> TokenMap
TokenMap.fromFlatList (NonEmpty [(AssetId, TokenQuantity)] -> NonEmpty TokenMap)
-> ([NonEmpty (AssetId, TokenQuantity)]
    -> NonEmpty [(AssetId, TokenQuantity)])
-> [NonEmpty (AssetId, TokenQuantity)]
-> NonEmpty TokenMap
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [NonEmpty (AssetId, TokenQuantity)]
-> NonEmpty [(AssetId, TokenQuantity)]
forall a. [NonEmpty a] -> NonEmpty [a]
transposeNE ([NonEmpty (AssetId, TokenQuantity)] -> NonEmpty TokenMap)
-> Gen [NonEmpty (AssetId, TokenQuantity)]
-> Gen (NonEmpty TokenMap)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
        ((AssetId, TokenQuantity)
 -> Gen (NonEmpty (AssetId, TokenQuantity)))
-> [(AssetId, TokenQuantity)]
-> Gen [NonEmpty (AssetId, TokenQuantity)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (AssetId, TokenQuantity) -> Gen (NonEmpty (AssetId, TokenQuantity))
forall a. (a, TokenQuantity) -> Gen (NonEmpty (a, TokenQuantity))
partitionAQ (TokenMap -> [(AssetId, TokenQuantity)]
TokenMap.toFlatList TokenMap
m)
  where
    partitionAQ :: (a, TokenQuantity) -> Gen (NonEmpty (a, TokenQuantity))
    partitionAQ :: (a, TokenQuantity) -> Gen (NonEmpty (a, TokenQuantity))
partitionAQ = ((a, NonEmpty TokenQuantity) -> NonEmpty (a, TokenQuantity))
-> Gen (a, NonEmpty TokenQuantity)
-> Gen (NonEmpty (a, TokenQuantity))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, NonEmpty TokenQuantity) -> NonEmpty (a, TokenQuantity)
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA (Gen (a, NonEmpty TokenQuantity)
 -> Gen (NonEmpty (a, TokenQuantity)))
-> ((a, TokenQuantity) -> Gen (a, NonEmpty TokenQuantity))
-> (a, TokenQuantity)
-> Gen (NonEmpty (a, TokenQuantity))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TokenQuantity -> Gen (NonEmpty TokenQuantity))
-> (a, TokenQuantity) -> Gen (a, NonEmpty TokenQuantity)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (TokenQuantity -> Int -> Gen (NonEmpty TokenQuantity)
`genTokenQuantityPartition` Int
i)

    transposeNE :: [NonEmpty a] -> NonEmpty [a]
    transposeNE :: [NonEmpty a] -> NonEmpty [a]
transposeNE = String -> Maybe (NonEmpty [a]) -> NonEmpty [a]
forall a. Partial => String -> Maybe a -> a
fromJustNote String
note (Maybe (NonEmpty [a]) -> NonEmpty [a])
-> ([NonEmpty a] -> Maybe (NonEmpty [a]))
-> [NonEmpty a]
-> NonEmpty [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[a]] -> Maybe (NonEmpty [a])
forall a. [a] -> Maybe (NonEmpty a)
NE.nonEmpty ([[a]] -> Maybe (NonEmpty [a]))
-> ([NonEmpty a] -> [[a]]) -> [NonEmpty a] -> Maybe (NonEmpty [a])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[a]] -> [[a]]
forall a. [[a]] -> [[a]]
transpose ([[a]] -> [[a]])
-> ([NonEmpty a] -> [[a]]) -> [NonEmpty a] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (NonEmpty a -> [a]) -> [NonEmpty a] -> [[a]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap NonEmpty a -> [a]
forall a. NonEmpty a -> [a]
NE.toList
      where
        note :: String
note = String
"genTokenMapPartition.transposeNE: unexpected empty list"

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