From 652323a625ec8ca4ecf97f86168680f1fc51a78d Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Thu, 14 Jan 2021 14:31:47 +0000
Subject: [PATCH] Fork `genTokenBundleSmallRange` into positive and
 non-positive variants.

This change forks `genTokenBundleSmallRange` into two variants:

  - `genTokenBundleSmallRange`
     Generates token bundles where the ada quantity may be zero.

  - `genTokenBundleSmallRangePositive`
     Generates token bundles where the ada quantity is always non-zero.

This is necessary, as some QC properties require token bundles with ada
quantities of zero.

But coin selection QC properties typically require token bundles (within
transaction outputs) to have non-zero ada quantities.

This change also forks the associated shrinker function
`shrinkTokenBundleSmallRange` in a similar fashion.
---
 .../Wallet/Primitive/Types/TokenBundle/Gen.hs | 21 +++++++++++++++++--
 .../CoinSelection/MA/RoundRobinSpec.hs        | 10 ++++-----
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/lib/core/src/Cardano/Wallet/Primitive/Types/TokenBundle/Gen.hs b/lib/core/src/Cardano/Wallet/Primitive/Types/TokenBundle/Gen.hs
index 989327e95e1..9133ddb10ca 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/Types/TokenBundle/Gen.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/Types/TokenBundle/Gen.hs
@@ -1,12 +1,18 @@
 module Cardano.Wallet.Primitive.Types.TokenBundle.Gen
     ( genTokenBundleSmallRange
+    , genTokenBundleSmallRangePositive
     , shrinkTokenBundleSmallRange
+    , shrinkTokenBundleSmallRangePositive
     ) where
 
 import Prelude
 
 import Cardano.Wallet.Primitive.Types.Coin.Gen
-    ( genCoinSmallPositive, shrinkCoinSmallPositive )
+    ( genCoinSmall
+    , genCoinSmallPositive
+    , shrinkCoinSmall
+    , shrinkCoinSmallPositive
+    )
 import Cardano.Wallet.Primitive.Types.TokenBundle
     ( TokenBundle (..) )
 import Cardano.Wallet.Primitive.Types.TokenMap.Gen
@@ -22,11 +28,22 @@ import Test.QuickCheck.Extra
 
 genTokenBundleSmallRange :: Gen TokenBundle
 genTokenBundleSmallRange = TokenBundle
-    <$> genCoinSmallPositive
+    <$> genCoinSmall
     <*> genTokenMapSmallRange
 
 shrinkTokenBundleSmallRange :: TokenBundle -> [TokenBundle]
 shrinkTokenBundleSmallRange (TokenBundle c m) =
+    uncurry TokenBundle <$> shrinkInterleaved
+        (c, shrinkCoinSmall)
+        (m, shrinkTokenMapSmallRange)
+
+genTokenBundleSmallRangePositive :: Gen TokenBundle
+genTokenBundleSmallRangePositive = TokenBundle
+    <$> genCoinSmallPositive
+    <*> genTokenMapSmallRange
+
+shrinkTokenBundleSmallRangePositive :: TokenBundle -> [TokenBundle]
+shrinkTokenBundleSmallRangePositive (TokenBundle c m) =
     uncurry TokenBundle <$> shrinkInterleaved
         (c, shrinkCoinSmallPositive)
         (m, shrinkTokenMapSmallRange)
diff --git a/lib/core/test/unit/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobinSpec.hs b/lib/core/test/unit/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobinSpec.hs
index 3d3cfee9bba..6dfa4aab386 100644
--- a/lib/core/test/unit/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobinSpec.hs
+++ b/lib/core/test/unit/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobinSpec.hs
@@ -38,7 +38,7 @@ import Cardano.Wallet.Primitive.Types.Coin
 import Cardano.Wallet.Primitive.Types.TokenBundle
     ( TokenBundle )
 import Cardano.Wallet.Primitive.Types.TokenBundle.Gen
-    ( genTokenBundleSmallRange, shrinkTokenBundleSmallRange )
+    ( genTokenBundleSmallRangePositive, shrinkTokenBundleSmallRangePositive )
 import Cardano.Wallet.Primitive.Types.TokenMap
     ( AssetId (..) )
 import Cardano.Wallet.Primitive.Types.TokenMap.Gen
@@ -518,8 +518,8 @@ genMakeChangeData = flip suchThat isValidMakeChangeData $ do
   where
     genTokenBundles :: Int -> Gen (NonEmpty TokenBundle)
     genTokenBundles count = (:|)
-        <$> genTokenBundleSmallRange
-        <*> replicateM count genTokenBundleSmallRange
+        <$> genTokenBundleSmallRangePositive
+        <*> replicateM count genTokenBundleSmallRangePositive
 
 prop_makeChange_identity
     :: NonEmpty TokenBundle -> Property
@@ -762,8 +762,8 @@ instance Arbitrary (MockRoundRobinState TokenName Word8) where
     shrink = shrinkMockRoundRobinState shrink
 
 instance Arbitrary TokenBundle where
-    arbitrary = genTokenBundleSmallRange
-    shrink = shrinkTokenBundleSmallRange
+    arbitrary = genTokenBundleSmallRangePositive
+    shrink = shrinkTokenBundleSmallRangePositive
 
 instance Arbitrary TokenQuantity where
     arbitrary = genTokenQuantitySmallPositive