module Cardano.Wallet.Primitive.Types.UTxO.Gen
    ( genUTxO
    , genUTxOLarge
    , genUTxOLargeN
    , selectUTxOEntries
    , shrinkUTxO
    ) where

import Prelude

import Cardano.Wallet.Primitive.Types.Tx
    ( TxIn, TxOut )
import Cardano.Wallet.Primitive.Types.Tx.Gen
    ( genTxIn, genTxInLargeRange, genTxOut, shrinkTxIn, shrinkTxOut )
import Cardano.Wallet.Primitive.Types.UTxO
    ( UTxO (..) )
import Control.Monad
    ( replicateM )
import Test.QuickCheck
    ( Gen, choose, shrinkList, sized )
import Test.QuickCheck.Extra
    ( selectMapEntries, shrinkInterleaved )

import qualified Data.Map.Strict as Map

--------------------------------------------------------------------------------
-- UTxO sets generated according to the size parameter
--------------------------------------------------------------------------------

genUTxO :: Gen UTxO
genUTxO :: Gen UTxO
genUTxO = (Int -> Gen UTxO) -> Gen UTxO
forall a. (Int -> Gen a) -> Gen a
sized ((Int -> Gen UTxO) -> Gen UTxO) -> (Int -> Gen UTxO) -> Gen UTxO
forall a b. (a -> b) -> a -> b
$ \Int
size -> do
    Int
entryCount <- (Int, Int) -> Gen Int
forall a. Random a => (a, a) -> Gen a
choose (Int
0, Int
size)
    Map TxIn TxOut -> UTxO
UTxO (Map TxIn TxOut -> UTxO)
-> ([(TxIn, TxOut)] -> Map TxIn TxOut) -> [(TxIn, TxOut)] -> UTxO
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(TxIn, TxOut)] -> Map TxIn TxOut
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(TxIn, TxOut)] -> UTxO) -> Gen [(TxIn, TxOut)] -> Gen UTxO
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Gen (TxIn, TxOut) -> Gen [(TxIn, TxOut)]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
entryCount Gen (TxIn, TxOut)
genEntry

shrinkUTxO :: UTxO -> [UTxO]
shrinkUTxO :: UTxO -> [UTxO]
shrinkUTxO
    = Int -> [UTxO] -> [UTxO]
forall a. Int -> [a] -> [a]
take Int
16
    ([UTxO] -> [UTxO]) -> (UTxO -> [UTxO]) -> UTxO -> [UTxO]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([(TxIn, TxOut)] -> UTxO) -> [[(TxIn, TxOut)]] -> [UTxO]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Map TxIn TxOut -> UTxO
UTxO (Map TxIn TxOut -> UTxO)
-> ([(TxIn, TxOut)] -> Map TxIn TxOut) -> [(TxIn, TxOut)] -> UTxO
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(TxIn, TxOut)] -> Map TxIn TxOut
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList)
    ([[(TxIn, TxOut)]] -> [UTxO])
-> (UTxO -> [[(TxIn, TxOut)]]) -> UTxO -> [UTxO]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((TxIn, TxOut) -> [(TxIn, TxOut)])
-> [(TxIn, TxOut)] -> [[(TxIn, TxOut)]]
forall a. (a -> [a]) -> [a] -> [[a]]
shrinkList (TxIn, TxOut) -> [(TxIn, TxOut)]
shrinkEntry
    ([(TxIn, TxOut)] -> [[(TxIn, TxOut)]])
-> (UTxO -> [(TxIn, TxOut)]) -> UTxO -> [[(TxIn, TxOut)]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map TxIn TxOut -> [(TxIn, TxOut)]
forall k a. Map k a -> [(k, a)]
Map.toList
    (Map TxIn TxOut -> [(TxIn, TxOut)])
-> (UTxO -> Map TxIn TxOut) -> UTxO -> [(TxIn, TxOut)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UTxO -> Map TxIn TxOut
unUTxO

genEntry :: Gen (TxIn, TxOut)
genEntry :: Gen (TxIn, TxOut)
genEntry = (,) (TxIn -> TxOut -> (TxIn, TxOut))
-> Gen TxIn -> Gen (TxOut -> (TxIn, TxOut))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen TxIn
genTxIn Gen (TxOut -> (TxIn, TxOut)) -> Gen TxOut -> Gen (TxIn, TxOut)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen TxOut
genTxOut

shrinkEntry :: (TxIn, TxOut) -> [(TxIn, TxOut)]
shrinkEntry :: (TxIn, TxOut) -> [(TxIn, TxOut)]
shrinkEntry (TxIn
i, TxOut
o) = (TxIn -> TxOut -> (TxIn, TxOut)) -> (TxIn, TxOut) -> (TxIn, TxOut)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (,) ((TxIn, TxOut) -> (TxIn, TxOut))
-> [(TxIn, TxOut)] -> [(TxIn, TxOut)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TxIn, TxIn -> [TxIn])
-> (TxOut, TxOut -> [TxOut]) -> [(TxIn, TxOut)]
forall a b. (a, a -> [a]) -> (b, b -> [b]) -> [(a, b)]
shrinkInterleaved
    (TxIn
i, TxIn -> [TxIn]
shrinkTxIn)
    (TxOut
o, TxOut -> [TxOut]
shrinkTxOut)

--------------------------------------------------------------------------------
-- Large UTxO sets
--------------------------------------------------------------------------------

genUTxOLarge :: Gen UTxO
genUTxOLarge :: Gen UTxO
genUTxOLarge = do
    Int
entryCount <- (Int, Int) -> Gen Int
forall a. Random a => (a, a) -> Gen a
choose (Int
1024, Int
4096)
    Int -> Gen UTxO
genUTxOLargeN Int
entryCount

genUTxOLargeN :: Int -> Gen UTxO
genUTxOLargeN :: Int -> Gen UTxO
genUTxOLargeN Int
entryCount = do
    Map TxIn TxOut -> UTxO
UTxO (Map TxIn TxOut -> UTxO)
-> ([(TxIn, TxOut)] -> Map TxIn TxOut) -> [(TxIn, TxOut)] -> UTxO
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(TxIn, TxOut)] -> Map TxIn TxOut
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(TxIn, TxOut)] -> UTxO) -> Gen [(TxIn, TxOut)] -> Gen UTxO
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Gen (TxIn, TxOut) -> Gen [(TxIn, TxOut)]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
entryCount Gen (TxIn, TxOut)
genEntryLargeRange

genEntryLargeRange :: Gen (TxIn, TxOut)
genEntryLargeRange :: Gen (TxIn, TxOut)
genEntryLargeRange = (,)
    (TxIn -> TxOut -> (TxIn, TxOut))
-> Gen TxIn -> Gen (TxOut -> (TxIn, TxOut))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen TxIn
genTxInLargeRange
    -- Note that we don't need to choose outputs from a large range, as inputs
    -- are already chosen from a large range:
    Gen (TxOut -> (TxIn, TxOut)) -> Gen TxOut -> Gen (TxIn, TxOut)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen TxOut
genTxOut

--------------------------------------------------------------------------------
-- Selecting random UTxO entries
--------------------------------------------------------------------------------

-- | Selects up to a given number of entries at random from the given UTxO set.
--
-- Returns the selected entries and the remaining UTxO set with the entries
-- removed.
--
selectUTxOEntries :: UTxO -> Int -> Gen ([(TxIn, TxOut)], UTxO)
selectUTxOEntries :: UTxO -> Int -> Gen ([(TxIn, TxOut)], UTxO)
selectUTxOEntries = ((([(TxIn, TxOut)], Map TxIn TxOut) -> ([(TxIn, TxOut)], UTxO))
-> Gen ([(TxIn, TxOut)], Map TxIn TxOut)
-> Gen ([(TxIn, TxOut)], UTxO)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Map TxIn TxOut -> UTxO)
-> ([(TxIn, TxOut)], Map TxIn TxOut) -> ([(TxIn, TxOut)], UTxO)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Map TxIn TxOut -> UTxO
UTxO) (Gen ([(TxIn, TxOut)], Map TxIn TxOut)
 -> Gen ([(TxIn, TxOut)], UTxO))
-> (Int -> Gen ([(TxIn, TxOut)], Map TxIn TxOut))
-> Int
-> Gen ([(TxIn, TxOut)], UTxO)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) ((Int -> Gen ([(TxIn, TxOut)], Map TxIn TxOut))
 -> Int -> Gen ([(TxIn, TxOut)], UTxO))
-> (UTxO -> Int -> Gen ([(TxIn, TxOut)], Map TxIn TxOut))
-> UTxO
-> Int
-> Gen ([(TxIn, TxOut)], UTxO)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map TxIn TxOut -> Int -> Gen ([(TxIn, TxOut)], Map TxIn TxOut)
forall k v. Ord k => Map k v -> Int -> Gen ([(k, v)], Map k v)
selectMapEntries (Map TxIn TxOut -> Int -> Gen ([(TxIn, TxOut)], Map TxIn TxOut))
-> (UTxO -> Map TxIn TxOut)
-> UTxO
-> Int
-> Gen ([(TxIn, TxOut)], Map TxIn TxOut)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UTxO -> Map TxIn TxOut
unUTxO