{-# LANGUAGE ScopedTypeVariables #-}

module Cardano.Wallet.Primitive.Types.UTxOIndex.Gen
    ( genUTxOIndex
    , genUTxOIndexLarge
    , genUTxOIndexLargeN
    , shrinkUTxOIndex
    ) where

import Prelude

import Cardano.Wallet.Primitive.Types.TokenBundle
    ( TokenBundle )
import Cardano.Wallet.Primitive.Types.TokenBundle.Gen
    ( genTokenBundleSmallRangePositive, shrinkTokenBundleSmallRangePositive )
import Cardano.Wallet.Primitive.Types.UTxOIndex
    ( UTxOIndex )
import Control.Monad
    ( replicateM )
import Generics.SOP
    ( NP (..) )
import Test.QuickCheck
    ( Gen, choose, listOf, shrinkList, shrinkMapBy )
import Test.QuickCheck.Extra
    ( genericRoundRobinShrink, (<:>), (<@>) )

import qualified Cardano.Wallet.Primitive.Types.UTxOIndex as UTxOIndex

--------------------------------------------------------------------------------
-- Indices generated according to the size parameter
--------------------------------------------------------------------------------

genUTxOIndex :: forall u. Ord u => Gen u -> Gen (UTxOIndex u)
genUTxOIndex :: Gen u -> Gen (UTxOIndex u)
genUTxOIndex Gen u
genUTxO = [(u, TokenBundle)] -> UTxOIndex u
forall (f :: * -> *) u.
(Foldable f, Ord u) =>
f (u, TokenBundle) -> UTxOIndex u
UTxOIndex.fromSequence ([(u, TokenBundle)] -> UTxOIndex u)
-> Gen [(u, TokenBundle)] -> Gen (UTxOIndex u)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (u, TokenBundle) -> Gen [(u, TokenBundle)]
forall a. Gen a -> Gen [a]
listOf Gen (u, TokenBundle)
genEntry
  where
    genEntry :: Gen (u, TokenBundle)
    genEntry :: Gen (u, TokenBundle)
genEntry = (,) (u -> TokenBundle -> (u, TokenBundle))
-> Gen u -> Gen (TokenBundle -> (u, TokenBundle))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen u
genUTxO Gen (TokenBundle -> (u, TokenBundle))
-> Gen TokenBundle -> Gen (u, TokenBundle)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen TokenBundle
genTokenBundleSmallRangePositive

shrinkUTxOIndex :: forall u. Ord u => (u -> [u]) -> UTxOIndex u -> [UTxOIndex u]
shrinkUTxOIndex :: (u -> [u]) -> UTxOIndex u -> [UTxOIndex u]
shrinkUTxOIndex u -> [u]
shrinkUTxO =
    ([(u, TokenBundle)] -> UTxOIndex u)
-> (UTxOIndex u -> [(u, TokenBundle)])
-> ([(u, TokenBundle)] -> [[(u, TokenBundle)]])
-> UTxOIndex u
-> [UTxOIndex u]
forall a b. (a -> b) -> (b -> a) -> (a -> [a]) -> b -> [b]
shrinkMapBy [(u, TokenBundle)] -> UTxOIndex u
forall (f :: * -> *) u.
(Foldable f, Ord u) =>
f (u, TokenBundle) -> UTxOIndex u
UTxOIndex.fromSequence UTxOIndex u -> [(u, TokenBundle)]
forall u. UTxOIndex u -> [(u, TokenBundle)]
UTxOIndex.toList (((u, TokenBundle) -> [(u, TokenBundle)])
-> [(u, TokenBundle)] -> [[(u, TokenBundle)]]
forall a. (a -> [a]) -> [a] -> [[a]]
shrinkList (u, TokenBundle) -> [(u, TokenBundle)]
shrinkEntry)
  where
    shrinkEntry :: (u, TokenBundle) -> [(u, TokenBundle)]
    shrinkEntry :: (u, TokenBundle) -> [(u, TokenBundle)]
shrinkEntry = NP (I -.-> []) '[u, TokenBundle]
-> (u, TokenBundle) -> [(u, TokenBundle)]
forall a (xs :: [*]).
(Generic a, GFrom a, GTo a, GCode a ~ '[xs]) =>
NP (I -.-> []) xs -> a -> [a]
genericRoundRobinShrink
        (NP (I -.-> []) '[u, TokenBundle]
 -> (u, TokenBundle) -> [(u, TokenBundle)])
-> NP (I -.-> []) '[u, TokenBundle]
-> (u, TokenBundle)
-> [(u, TokenBundle)]
forall a b. (a -> b) -> a -> b
<@> u -> [u]
shrinkUTxO
        (u -> [u])
-> NP (I -.-> []) '[TokenBundle]
-> NP (I -.-> []) '[u, TokenBundle]
forall x (xs :: [*]).
(x -> [x]) -> NP (I -.-> []) xs -> NP (I -.-> []) (x : xs)
<:> TokenBundle -> [TokenBundle]
shrinkTokenBundleSmallRangePositive
        (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

--------------------------------------------------------------------------------
-- Large indices
--------------------------------------------------------------------------------

genUTxOIndexLarge :: Ord u => Gen u -> Gen (UTxOIndex u)
genUTxOIndexLarge :: Gen u -> Gen (UTxOIndex u)
genUTxOIndexLarge Gen u
genUTxO =
    Gen u -> Int -> Gen (UTxOIndex u)
forall u. Ord u => Gen u -> Int -> Gen (UTxOIndex u)
genUTxOIndexLargeN Gen u
genUTxO (Int -> Gen (UTxOIndex u)) -> Gen Int -> Gen (UTxOIndex u)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Int, Int) -> Gen Int
forall a. Random a => (a, a) -> Gen a
choose (Int
1024, Int
4096)

genUTxOIndexLargeN :: forall u. Ord u => Gen u -> Int -> Gen (UTxOIndex u)
genUTxOIndexLargeN :: Gen u -> Int -> Gen (UTxOIndex u)
genUTxOIndexLargeN Gen u
genUTxO Int
n = [(u, TokenBundle)] -> UTxOIndex u
forall (f :: * -> *) u.
(Foldable f, Ord u) =>
f (u, TokenBundle) -> UTxOIndex u
UTxOIndex.fromSequence ([(u, TokenBundle)] -> UTxOIndex u)
-> Gen [(u, TokenBundle)] -> Gen (UTxOIndex u)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Gen (u, TokenBundle) -> Gen [(u, TokenBundle)]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n Gen (u, TokenBundle)
genEntry
  where
    genEntry :: Gen (u, TokenBundle)
    genEntry :: Gen (u, TokenBundle)
genEntry = (,) (u -> TokenBundle -> (u, TokenBundle))
-> Gen u -> Gen (TokenBundle -> (u, TokenBundle))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen u
genUTxO Gen (TokenBundle -> (u, TokenBundle))
-> Gen TokenBundle -> Gen (u, TokenBundle)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen TokenBundle
genTokenBundleSmallRangePositive