Generalize UTxO identifiers and addresses in coin selection Gen modules.

lib/core/cardano-wallet-core.cabal

@@ -263,6 +263,7 @@ library
         -- The following modules define QC generators and shrinkers that can
         -- be used by both `cardano-wallet-core` and `cardano-wallet`:
         --
+      Cardano.Wallet.CoinSelection.Gen
       Cardano.Wallet.CoinSelection.Internal.Balance.Gen
       Cardano.Wallet.Primitive.Types.Address.Gen
       Cardano.Wallet.Primitive.Types.Coin.Gen

lib/core/src/Cardano/Wallet/CoinSelection/Gen.hs

@@ -0,0 +1,49 @@
+module Cardano.Wallet.CoinSelection.Gen
+    ( coarbitraryWalletUTxO
+    , genWalletUTxO
+    , genWalletUTxOFunction
+    , genWalletUTxOLargeRange
+    , shrinkWalletUTxO
+    )
+    where
+
+import Prelude
+
+import Cardano.Wallet.CoinSelection
+    ( WalletUTxO (..) )
+import Cardano.Wallet.Primitive.Types.Address.Gen
+    ( genAddress, shrinkAddress )
+import Cardano.Wallet.Primitive.Types.Tx.Gen
+    ( genTxIn, genTxInLargeRange, shrinkTxIn )
+import Generics.SOP
+    ( NP (..) )
+import Test.QuickCheck
+    ( Gen, coarbitrary )
+import Test.QuickCheck.Extra
+    ( genFunction, genSized2, genericRoundRobinShrink, (<:>), (<@>) )
+
+--------------------------------------------------------------------------------
+-- Wallet UTxO identifiers chosen according to the size parameter
+--------------------------------------------------------------------------------
+
+coarbitraryWalletUTxO :: WalletUTxO -> Gen a -> Gen a
+coarbitraryWalletUTxO = coarbitrary . show
+
+genWalletUTxO :: Gen WalletUTxO
+genWalletUTxO = uncurry WalletUTxO <$> genSized2 genTxIn genAddress
+
+shrinkWalletUTxO :: WalletUTxO -> [WalletUTxO]
+shrinkWalletUTxO = genericRoundRobinShrink
+    <@> shrinkTxIn
+    <:> shrinkAddress
+    <:> Nil
+
+genWalletUTxOFunction :: Gen a -> Gen (WalletUTxO -> a)
+genWalletUTxOFunction = genFunction coarbitraryWalletUTxO
+
+--------------------------------------------------------------------------------
+-- Wallet UTxO identifiers chosen from a large range
+--------------------------------------------------------------------------------
+
+genWalletUTxOLargeRange :: Gen WalletUTxO
+genWalletUTxOLargeRange = WalletUTxO <$> genTxInLargeRange <*> genAddress

lib/core/src/Cardano/Wallet/CoinSelection/Internal/Balance/Gen.hs

@@ -13,16 +13,14 @@ module Cardano.Wallet.CoinSelection.Internal.Balance.Gen
 
 import Prelude
 
-import Cardano.Wallet.CoinSelection
-    ( WalletSelectionContext )
 import Cardano.Wallet.CoinSelection.Internal.Balance
     ( SelectionLimit
     , SelectionLimitOf (..)
     , SelectionSkeleton (..)
     , SelectionStrategy (..)
     )
-import Cardano.Wallet.Primitive.Types.Address.Gen
-    ( genAddress, shrinkAddress )
+import Cardano.Wallet.CoinSelection.Internal.Context
+    ( SelectionContext (..) )
 import Cardano.Wallet.Primitive.Types.Coin
     ( Coin (..) )
 import Cardano.Wallet.Primitive.Types.TokenBundle
@@ -72,8 +70,8 @@ shrinkSelectionLimit = \case
 -- Selection skeletons
 --------------------------------------------------------------------------------
 
-genSelectionSkeleton :: Gen (SelectionSkeleton WalletSelectionContext)
-genSelectionSkeleton = SelectionSkeleton
+genSelectionSkeleton :: Gen (Address ctx) -> Gen (SelectionSkeleton ctx)
+genSelectionSkeleton genAddress = SelectionSkeleton
     <$> genSkeletonInputCount
     <*> genSkeletonOutputs
     <*> genSkeletonChange
@@ -89,9 +87,9 @@ genSelectionSkeleton = SelectionSkeleton
         listOf (Set.fromList <$> listOf genAssetId)
 
 shrinkSelectionSkeleton
-    :: SelectionSkeleton WalletSelectionContext
-    -> [SelectionSkeleton WalletSelectionContext]
-shrinkSelectionSkeleton = genericRoundRobinShrink
+    :: (Address ctx -> [Address ctx])
+    -> (SelectionSkeleton ctx -> [SelectionSkeleton ctx])
+shrinkSelectionSkeleton shrinkAddress = genericRoundRobinShrink
     <@> shrinkSkeletonInputCount
     <:> shrinkSkeletonOutputs
     <:> shrinkSkeletonChange

lib/core/src/Cardano/Wallet/Primitive/Types/UTxOIndex/Gen.hs

@@ -1,3 +1,5 @@
+{-# LANGUAGE ScopedTypeVariables #-}
+
 module Cardano.Wallet.Primitive.Types.UTxOIndex.Gen
     ( genUTxOIndex
     , genUTxOIndexLarge
@@ -7,16 +9,10 @@ module Cardano.Wallet.Primitive.Types.UTxOIndex.Gen
 
 import Prelude
 
-import Cardano.Wallet.CoinSelection
-    ( WalletUTxO (..) )
-import Cardano.Wallet.Primitive.Types.Address.Gen
-    ( genAddress, shrinkAddress )
 import Cardano.Wallet.Primitive.Types.TokenBundle
     ( TokenBundle )
 import Cardano.Wallet.Primitive.Types.TokenBundle.Gen
     ( genTokenBundleSmallRangePositive, shrinkTokenBundleSmallRangePositive )
-import Cardano.Wallet.Primitive.Types.Tx.Gen
-    ( genTxIn, genTxInLargeRange, shrinkTxIn )
 import Cardano.Wallet.Primitive.Types.UTxOIndex
     ( UTxOIndex )
 import Control.Monad
@@ -26,51 +22,40 @@ import Generics.SOP
 import Test.QuickCheck
     ( Gen, choose, listOf, shrinkList, shrinkMapBy )
 import Test.QuickCheck.Extra
-    ( genSized2, genericRoundRobinShrink, (<:>), (<@>) )
+    ( genericRoundRobinShrink, (<:>), (<@>) )
 
 import qualified Cardano.Wallet.Primitive.Types.UTxOIndex as UTxOIndex
 
 --------------------------------------------------------------------------------
 -- Indices generated according to the size parameter
 --------------------------------------------------------------------------------
 
-genUTxOIndex :: Gen (UTxOIndex WalletUTxO)
-genUTxOIndex = UTxOIndex.fromSequence <$> listOf genEntry
+genUTxOIndex :: forall u. Ord u => Gen u -> Gen (UTxOIndex u)
+genUTxOIndex genUTxO = UTxOIndex.fromSequence <$> listOf genEntry
   where
-    genEntry :: Gen (WalletUTxO, TokenBundle)
-    genEntry = (,) <$> genWalletUTxO <*> genTokenBundleSmallRangePositive
-
-    genWalletUTxO :: Gen WalletUTxO
-    genWalletUTxO = uncurry WalletUTxO <$> genSized2 genTxIn genAddress
+    genEntry :: Gen (u, TokenBundle)
+    genEntry = (,) <$> genUTxO <*> genTokenBundleSmallRangePositive
 
-shrinkUTxOIndex :: UTxOIndex WalletUTxO -> [UTxOIndex WalletUTxO]
-shrinkUTxOIndex =
+shrinkUTxOIndex :: forall u. Ord u => (u -> [u]) -> UTxOIndex u -> [UTxOIndex u]
+shrinkUTxOIndex shrinkUTxO =
     shrinkMapBy UTxOIndex.fromSequence UTxOIndex.toList (shrinkList shrinkEntry)
   where
-    shrinkEntry :: (WalletUTxO, TokenBundle) -> [(WalletUTxO, TokenBundle)]
+    shrinkEntry :: (u, TokenBundle) -> [(u, TokenBundle)]
     shrinkEntry = genericRoundRobinShrink
-        <@> shrinkWalletUTxO
+        <@> shrinkUTxO
         <:> shrinkTokenBundleSmallRangePositive
         <:> Nil
 
-    shrinkWalletUTxO :: WalletUTxO -> [WalletUTxO]
-    shrinkWalletUTxO = genericRoundRobinShrink
-        <@> shrinkTxIn
-        <:> shrinkAddress
-        <:> Nil
-
 --------------------------------------------------------------------------------
 -- Large indices
 --------------------------------------------------------------------------------
 
-genUTxOIndexLarge :: Gen (UTxOIndex WalletUTxO)
-genUTxOIndexLarge = genUTxOIndexLargeN =<< choose (1024, 4096)
+genUTxOIndexLarge :: Ord u => Gen u -> Gen (UTxOIndex u)
+genUTxOIndexLarge genUTxO =
+    genUTxOIndexLargeN genUTxO =<< choose (1024, 4096)
 
-genUTxOIndexLargeN :: Int -> Gen (UTxOIndex WalletUTxO)
-genUTxOIndexLargeN n = UTxOIndex.fromSequence <$> replicateM n genEntry
+genUTxOIndexLargeN :: forall u. Ord u => Gen u -> Int -> Gen (UTxOIndex u)
+genUTxOIndexLargeN genUTxO n = UTxOIndex.fromSequence <$> replicateM n genEntry
   where
-    genEntry :: Gen (WalletUTxO, TokenBundle)
-    genEntry = (,) <$> genWalletUTxO <*> genTokenBundleSmallRangePositive
-
-    genWalletUTxO :: Gen WalletUTxO
-    genWalletUTxO = WalletUTxO <$> genTxInLargeRange <*> genAddress
+    genEntry :: Gen (u, TokenBundle)
+    genEntry = (,) <$> genUTxO <*> genTokenBundleSmallRangePositive

lib/core/src/Cardano/Wallet/Primitive/Types/UTxOSelection/Gen.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeApplications #-}
 
 module Cardano.Wallet.Primitive.Types.UTxOSelection.Gen
@@ -10,8 +11,6 @@ module Cardano.Wallet.Primitive.Types.UTxOSelection.Gen
 
 import Prelude
 
-import Cardano.Wallet.CoinSelection
-    ( WalletUTxO )
 import Cardano.Wallet.Primitive.Types.UTxOIndex.Gen
     ( genUTxOIndex, shrinkUTxOIndex )
 import Cardano.Wallet.Primitive.Types.UTxOSelection
@@ -29,39 +28,40 @@ import qualified Cardano.Wallet.Primitive.Types.UTxOSelection as UTxOSelection
 -- Selections that may be empty
 --------------------------------------------------------------------------------
 
-coarbitraryWalletUTxO :: WalletUTxO -> Gen a -> Gen a
-coarbitraryWalletUTxO = coarbitrary . show
+coarbitraryUTxO :: Show u => u -> Gen a -> Gen a
+coarbitraryUTxO = coarbitrary . show
 
-genWalletUTxOFunction :: Gen a -> Gen (WalletUTxO -> a)
-genWalletUTxOFunction = genFunction coarbitraryWalletUTxO
+genUTxOFunction :: Show u => Gen a -> Gen (u -> a)
+genUTxOFunction = genFunction coarbitraryUTxO
 
-genUTxOSelection :: Gen (UTxOSelection WalletUTxO)
-genUTxOSelection = UTxOSelection.fromIndexFiltered
-    <$> genFilter
-    <*> genUTxOIndex
+genUTxOSelection :: forall u. (Ord u, Show u) => Gen u -> Gen (UTxOSelection u)
+genUTxOSelection genUTxO = UTxOSelection.fromIndexFiltered
+    <$> genUTxOFilter
+    <*> genUTxOIndex genUTxO
   where
-    genFilter :: Gen (WalletUTxO -> Bool)
-    genFilter = genWalletUTxOFunction (arbitrary @Bool)
+    genUTxOFilter :: Gen (u -> Bool)
+    genUTxOFilter = genUTxOFunction (arbitrary @Bool)
 
-shrinkUTxOSelection :: UTxOSelection WalletUTxO -> [UTxOSelection WalletUTxO]
-shrinkUTxOSelection =
+shrinkUTxOSelection
+    :: Ord u => (u -> [u]) -> (UTxOSelection u -> [UTxOSelection u])
+shrinkUTxOSelection shrinkUTxO =
     shrinkMapBy UTxOSelection.fromIndexPair UTxOSelection.toIndexPair $
         liftShrink2
-            shrinkUTxOIndex
-            shrinkUTxOIndex
+            (shrinkUTxOIndex shrinkUTxO)
+            (shrinkUTxOIndex shrinkUTxO)
 
 --------------------------------------------------------------------------------
 -- Selections that are non-empty
 --------------------------------------------------------------------------------
 
-genUTxOSelectionNonEmpty :: Gen (UTxOSelectionNonEmpty WalletUTxO)
-genUTxOSelectionNonEmpty =
-    genUTxOSelection `suchThatMap` UTxOSelection.toNonEmpty
+genUTxOSelectionNonEmpty
+    :: (Ord u, Show u) => Gen u -> Gen (UTxOSelectionNonEmpty u)
+genUTxOSelectionNonEmpty genUTxO =
+    genUTxOSelection genUTxO `suchThatMap` UTxOSelection.toNonEmpty
 
 shrinkUTxOSelectionNonEmpty
-    :: UTxOSelectionNonEmpty WalletUTxO
-    -> [UTxOSelectionNonEmpty WalletUTxO]
-shrinkUTxOSelectionNonEmpty
+    :: Ord u => (u -> [u]) -> (UTxOSelectionNonEmpty u -> [UTxOSelectionNonEmpty u])
+shrinkUTxOSelectionNonEmpty shrinkUTxO
     = mapMaybe UTxOSelection.toNonEmpty
-    . shrinkUTxOSelection
+    . shrinkUTxOSelection shrinkUTxO
     . UTxOSelection.fromNonEmpty

lib/core/test/unit/Cardano/Wallet/CoinSelection/Internal/BalanceSpec.hs

@@ -42,6 +42,12 @@ import Cardano.Numeric.Util
     ( inAscendingPartialOrder )
 import Cardano.Wallet.CoinSelection
     ( WalletSelectionContext, WalletUTxO (..) )
+import Cardano.Wallet.CoinSelection.Gen
+    ( genWalletUTxO
+    , genWalletUTxOFunction
+    , genWalletUTxOLargeRange
+    , shrinkWalletUTxO
+    )
 import Cardano.Wallet.CoinSelection.Internal.Balance
     ( AssetCount (..)
     , BalanceInsufficientError (..)
@@ -214,7 +220,6 @@ import Test.QuickCheck
     , arbitraryBoundedEnum
     , checkCoverage
     , choose
-    , coarbitrary
     , conjoin
     , counterexample
     , cover
@@ -240,7 +245,7 @@ import Test.QuickCheck
 import Test.QuickCheck.Classes
     ( eqLaws, ordLaws )
 import Test.QuickCheck.Extra
-    ( genFunction, genericRoundRobinShrink, report, verify, (<:>), (<@>) )
+    ( genericRoundRobinShrink, report, verify, (<:>), (<@>) )
 import Test.QuickCheck.Monadic
     ( PropertyM (..), assert, monadicIO, monitor, run )
 import Test.Utils.Laws
@@ -654,7 +659,7 @@ shrinkSelectionParams
     -> [SelectionParams WalletSelectionContext]
 shrinkSelectionParams = genericRoundRobinShrink
     <@> shrinkList shrinkOutput
-    <:> shrinkUTxOSelection
+    <:> shrinkUTxOSelection shrinkWalletUTxO
     <:> shrinkCoin
     <:> shrinkCoin
     <:> shrinkTokenMap
@@ -868,7 +873,7 @@ prop_performSelection_huge = ioProperty $
     -- the cost of re-generating it on every pass. This will still generate
     -- interesting cases, since selection within that large index is random.
     property . prop_performSelection_huge_inner
-        <$> generate (genUTxOIndexLargeN 50000)
+        <$> generate (genUTxOIndexLargeN genWalletUTxOLargeRange 50000)
 
 prop_performSelection_huge_inner
     :: UTxOIndex WalletUTxO
@@ -4371,16 +4376,6 @@ unitTests lbl cases =
     forM_ (zip [1..] cases) $ \(i, test) ->
         it (lbl <> " example #" <> show @Int i) test
 
---------------------------------------------------------------------------------
--- Wallet UTxO identifiers
---------------------------------------------------------------------------------
-
-coarbitraryWalletUTxO :: WalletUTxO -> Gen a -> Gen a
-coarbitraryWalletUTxO = coarbitrary . show
-
-genWalletUTxOFunction :: Gen a -> Gen (WalletUTxO -> a)
-genWalletUTxOFunction = genFunction coarbitraryWalletUTxO
-
 --------------------------------------------------------------------------------
 -- Arbitrary instances
 --------------------------------------------------------------------------------
@@ -4444,8 +4439,8 @@ instance Arbitrary TxOut where
     shrink = shrinkTxOut
 
 instance Arbitrary (UTxOSelection WalletUTxO) where
-    arbitrary = genUTxOSelection
-    shrink = shrinkUTxOSelection
+    arbitrary = genUTxOSelection genWalletUTxO
+    shrink = shrinkUTxOSelection shrinkWalletUTxO
 
 newtype Large a = Large
     { getLarge :: a }
@@ -4458,22 +4453,22 @@ newtype Small a = Small
 instance Arbitrary (Large (SelectionParams WalletSelectionContext)) where
     arbitrary = Large <$> genSelectionParams
         (genWalletUTxOFunction (arbitrary @Bool))
-        (genUTxOIndexLarge)
+        (genUTxOIndexLarge genWalletUTxOLargeRange)
     shrink = shrinkMapBy Large getLarge shrinkSelectionParams
 
 instance Arbitrary (Small (SelectionParams WalletSelectionContext)) where
     arbitrary = Small <$> genSelectionParams
         (genWalletUTxOFunction (arbitrary @Bool))
-        (genUTxOIndex)
+        (genUTxOIndex genWalletUTxO)
     shrink = shrinkMapBy Small getSmall shrinkSelectionParams
 
 instance Arbitrary (Large (UTxOIndex WalletUTxO)) where
-    arbitrary = Large <$> genUTxOIndexLarge
-    shrink = shrinkMapBy Large getLarge shrinkUTxOIndex
+    arbitrary = Large <$> genUTxOIndexLarge genWalletUTxOLargeRange
+    shrink = shrinkMapBy Large getLarge (shrinkUTxOIndex shrinkWalletUTxO)
 
 instance Arbitrary (Small (UTxOIndex WalletUTxO)) where
-    arbitrary = Small <$> genUTxOIndex
-    shrink = shrinkMapBy Small getSmall shrinkUTxOIndex
+    arbitrary = Small <$> genUTxOIndex genWalletUTxO
+    shrink = shrinkMapBy Small getSmall (shrinkUTxOIndex shrinkWalletUTxO)
 
 instance Arbitrary Coin where
     arbitrary = genCoinPositive