From 316b5f688dadc9170fd386889375e2f18ca809e9 Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Tue, 2 Mar 2021 03:07:12 +0000
Subject: [PATCH 01/24] Move `unsafePartitionNatural` to `Numeric.Util`.

---
 .../Primitive/CoinSelection/MA/RoundRobin.hs  | 27 +----------------
 lib/numeric/src/Cardano/Numeric/Util.hs       | 30 +++++++++++++++++++
 2 files changed, 31 insertions(+), 26 deletions(-)

diff --git a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
index 769a8076468..bd2b69276d7 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
@@ -89,7 +89,7 @@ import Prelude
 import Algebra.PartialOrd
     ( PartialOrd (..) )
 import Cardano.Numeric.Util
-    ( padCoalesce, partitionNatural )
+    ( padCoalesce, partitionNatural, unsafePartitionNatural )
 import Cardano.Wallet.Primitive.Types.Coin
     ( Coin (..), addCoin, subtractCoin, sumCoins )
 import Cardano.Wallet.Primitive.Types.TokenBundle
@@ -1346,31 +1346,6 @@ equipartitionTokenMapWithMaxQuantity m (TokenQuantity maxQuantity)
         , "the maximum allowable token quantity cannot be zero."
         ]
 
---------------------------------------------------------------------------------
--- Unsafe partitioning
---------------------------------------------------------------------------------
-
--- | Partitions a natural number into a number of parts, where the size of each
---   part is proportional to the size of its corresponding element in the given
---   list of weights, and the number of parts is equal to the number of weights.
---
--- Throws a run-time error if the sum of weights is equal to zero.
---
-unsafePartitionNatural
-    :: HasCallStack
-    => Natural
-    -- ^ Natural number to partition
-    -> NonEmpty Natural
-    -- ^ List of weights
-    -> NonEmpty Natural
-unsafePartitionNatural target =
-    fromMaybe zeroWeightSumError . partitionNatural target
-  where
-    zeroWeightSumError = error $ unwords
-        [ "unsafePartitionNatural:"
-        , "specified weights must have a non-zero sum."
-        ]
-
 --------------------------------------------------------------------------------
 -- Grouping and ungrouping
 --------------------------------------------------------------------------------
diff --git a/lib/numeric/src/Cardano/Numeric/Util.hs b/lib/numeric/src/Cardano/Numeric/Util.hs
index 1589ba52b3a..e9e2684dd7a 100644
--- a/lib/numeric/src/Cardano/Numeric/Util.hs
+++ b/lib/numeric/src/Cardano/Numeric/Util.hs
@@ -5,6 +5,7 @@
 module Cardano.Numeric.Util
     ( padCoalesce
     , partitionNatural
+    , unsafePartitionNatural
     ) where
 
 import Prelude hiding
@@ -16,10 +17,14 @@ import Data.Function
     ( (&) )
 import Data.List.NonEmpty
     ( NonEmpty (..) )
+import Data.Maybe
+    ( fromMaybe )
 import Data.Ord
     ( Down (..), comparing )
 import Data.Ratio
     ( (%) )
+import GHC.Stack
+    ( HasCallStack )
 import Numeric.Natural
     ( Natural )
 
@@ -181,6 +186,31 @@ partitionNatural target weights
     totalWeight :: Natural
     totalWeight = F.sum weights
 
+--------------------------------------------------------------------------------
+-- Unsafe partitioning
+--------------------------------------------------------------------------------
+
+-- | Partitions a natural number into a number of parts, where the size of each
+--   part is proportional to the size of its corresponding element in the given
+--   list of weights, and the number of parts is equal to the number of weights.
+--
+-- Throws a run-time error if the sum of weights is equal to zero.
+--
+unsafePartitionNatural
+    :: HasCallStack
+    => Natural
+    -- ^ Natural number to partition
+    -> NonEmpty Natural
+    -- ^ List of weights
+    -> NonEmpty Natural
+unsafePartitionNatural target =
+    fromMaybe zeroWeightSumError . partitionNatural target
+  where
+    zeroWeightSumError = error $ unwords
+        [ "unsafePartitionNatural:"
+        , "specified weights must have a non-zero sum."
+        ]
+
 --------------------------------------------------------------------------------
 -- Internal types and functions
 --------------------------------------------------------------------------------

From e0a195a16fdfc7633300a39d5f395a5a3711ab3b Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Tue, 2 Mar 2021 03:29:28 +0000
Subject: [PATCH 02/24] Move `equipartitionNatural` to `Numeric.Util`.

---
 .../Primitive/CoinSelection/MA/RoundRobin.hs  | 19 +-------
 .../CoinSelection/MA/RoundRobinSpec.hs        | 47 ------------------
 lib/numeric/src/Cardano/Numeric/Util.hs       | 29 ++++++++++-
 .../test/unit/Cardano/Numeric/UtilSpec.hs     | 48 ++++++++++++++++++-
 4 files changed, 76 insertions(+), 67 deletions(-)

diff --git a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
index bd2b69276d7..f0e25d2ad6c 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
@@ -54,7 +54,6 @@ module Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     , assignCoinsToChangeMaps
 
     -- * Partitioning
-    , equipartitionNatural
     , equipartitionTokenBundleWithMaxQuantity
     , equipartitionTokenBundlesWithMaxQuantity
     , equipartitionTokenMap
@@ -89,7 +88,7 @@ import Prelude
 import Algebra.PartialOrd
     ( PartialOrd (..) )
 import Cardano.Numeric.Util
-    ( padCoalesce, partitionNatural, unsafePartitionNatural )
+    ( equipartitionNatural, padCoalesce, partitionNatural )
 import Cardano.Wallet.Primitive.Types.Coin
     ( Coin (..), addCoin, subtractCoin, sumCoins )
 import Cardano.Wallet.Primitive.Types.TokenBundle
@@ -1222,22 +1221,6 @@ equipartitionCoin c =
     -- value.
     fmap unsafeNaturalToCoin . equipartitionNatural (coinToNatural c)
 
--- | Computes the equipartition of a natural number into 'n' smaller numbers.
---
-equipartitionNatural
-    :: HasCallStack
-    => Natural
-    -- ^ The natural number to be partitioned.
-    -> NonEmpty a
-    -- ^ Represents the number of portions in which to partition the number.
-    -> NonEmpty Natural
-    -- ^ The partitioned numbers.
-equipartitionNatural n count =
-    -- Note: due to the behaviour of the underlying partition algorithm, a
-    -- simple list reversal is enough to ensure that the resultant list is
-    -- sorted in ascending order.
-    NE.reverse $ unsafePartitionNatural n (1 <$ count)
-
 -- | Computes the equipartition of a token map into 'n' smaller maps.
 --
 -- Each asset is partitioned independently.
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 f0e46cc36c7..54257975d1c 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
@@ -36,7 +36,6 @@ import Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     , assignCoinsToChangeMaps
     , balanceMissing
     , coinSelectionLens
-    , equipartitionNatural
     , equipartitionTokenBundleWithMaxQuantity
     , equipartitionTokenBundlesWithMaxQuantity
     , equipartitionTokenMap
@@ -163,7 +162,6 @@ import Test.QuickCheck
     , suchThat
     , withMaxSuccess
     , (.&&.)
-    , (.||.)
     , (===)
     , (==>)
     )
@@ -322,17 +320,6 @@ spec = describe "Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobinSpec" $
         unitTests "makeChangeForUserSpecifiedAsset"
             unit_makeChangeForUserSpecifiedAsset
 
-    parallel $ describe "Equipartitioning natural numbers" $ do
-
-        it "prop_equipartitionNatural_fair" $
-            property prop_equipartitionNatural_fair
-        it "prop_equipartitionNatural_length" $
-            property prop_equipartitionNatural_length
-        it "prop_equipartitionNatural_order" $
-            property prop_equipartitionNatural_order
-        it "prop_equipartitionNatural_sum" $
-            property prop_equipartitionNatural_sum
-
     parallel $ describe "Equipartitioning token maps" $ do
 
         it "prop_equipartitionTokenMap_fair" $
@@ -1875,40 +1862,6 @@ unit_makeChangeForUserSpecifiedAsset =
     assetC :: AssetId
     assetC = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "2")
 
---------------------------------------------------------------------------------
--- Equipartitioning natural numbers
---------------------------------------------------------------------------------
-
--- Test that natural numbers are equipartitioned fairly:
---
--- Each portion must be within unity of the ideal portion.
---
-prop_equipartitionNatural_fair
-    :: Natural -> NonEmpty () -> Property
-prop_equipartitionNatural_fair n count = (.||.)
-    (difference === 0)
-    (difference === 1)
-  where
-    difference :: Natural
-    difference = F.maximum results - F.minimum results
-
-    results :: NonEmpty Natural
-    results = equipartitionNatural n count
-
-prop_equipartitionNatural_length :: Natural -> NonEmpty () -> Property
-prop_equipartitionNatural_length n count =
-    NE.length (equipartitionNatural n count) === NE.length count
-
-prop_equipartitionNatural_order :: Natural -> NonEmpty () -> Property
-prop_equipartitionNatural_order n count =
-    NE.sort results === results
-  where
-    results = equipartitionNatural n count
-
-prop_equipartitionNatural_sum :: Natural -> NonEmpty () -> Property
-prop_equipartitionNatural_sum n count =
-    F.sum (equipartitionNatural n count) === n
-
 --------------------------------------------------------------------------------
 -- Equipartitioning token maps
 --------------------------------------------------------------------------------
diff --git a/lib/numeric/src/Cardano/Numeric/Util.hs b/lib/numeric/src/Cardano/Numeric/Util.hs
index e9e2684dd7a..45da3721ec7 100644
--- a/lib/numeric/src/Cardano/Numeric/Util.hs
+++ b/lib/numeric/src/Cardano/Numeric/Util.hs
@@ -3,9 +3,15 @@
 {-# LANGUAGE TypeApplications #-}
 
 module Cardano.Numeric.Util
-    ( padCoalesce
+    (
+      -- * Coalescing values
+      padCoalesce
+
+      -- * Partitioning natural numbers
+    , equipartitionNatural
     , partitionNatural
     , unsafePartitionNatural
+
     ) where
 
 import Prelude hiding
@@ -101,6 +107,27 @@ padCoalesce sourceUnsorted target
 -- Partitioning natural numbers
 --------------------------------------------------------------------------------
 
+-- | Computes the equipartition of a natural number into 'n' smaller numbers.
+--
+-- An /equipartition/ of a natural number 'n' is a /partition/ of that number
+-- into 'n' smaller numbers whose values differ by no more than 1.
+--
+-- The resultant list is sorted in ascending order.
+--
+equipartitionNatural
+    :: HasCallStack
+    => Natural
+    -- ^ The natural number to be partitioned.
+    -> NonEmpty a
+    -- ^ Represents the number of portions in which to partition the number.
+    -> NonEmpty Natural
+    -- ^ The partitioned numbers.
+equipartitionNatural n count =
+    -- Note: due to the behaviour of the underlying partition algorithm, a
+    -- simple list reversal is enough to ensure that the resultant list is
+    -- sorted in ascending order.
+    NE.reverse $ unsafePartitionNatural n (1 <$ count)
+
 -- | Partitions a natural number into a number of parts, where the size of each
 --   part is proportional to the size of its corresponding element in the given
 --   list of weights, and the number of parts is equal to the number of weights.
diff --git a/lib/numeric/test/unit/Cardano/Numeric/UtilSpec.hs b/lib/numeric/test/unit/Cardano/Numeric/UtilSpec.hs
index 9dc6df5748e..65647095eb1 100644
--- a/lib/numeric/test/unit/Cardano/Numeric/UtilSpec.hs
+++ b/lib/numeric/test/unit/Cardano/Numeric/UtilSpec.hs
@@ -8,7 +8,7 @@ module Cardano.Numeric.UtilSpec
 import Prelude
 
 import Cardano.Numeric.Util
-    ( padCoalesce, partitionNatural )
+    ( equipartitionNatural, padCoalesce, partitionNatural )
 import Data.List.NonEmpty
     ( NonEmpty (..) )
 import Data.Maybe
@@ -31,6 +31,7 @@ import Test.QuickCheck
     , shrinkIntegral
     , withMaxSuccess
     , (.&&.)
+    , (.||.)
     , (===)
     )
 
@@ -49,6 +50,17 @@ spec = do
         it "prop_padCoalesce_sum" $
             property $ prop_padCoalesce_sum @(Sum Int)
 
+    describe "equipartitionNatural" $ do
+
+        it "prop_equipartitionNatural_fair" $
+            property prop_equipartitionNatural_fair
+        it "prop_equipartitionNatural_length" $
+            property prop_equipartitionNatural_length
+        it "prop_equipartitionNatural_order" $
+            property prop_equipartitionNatural_order
+        it "prop_equipartitionNatural_sum" $
+            property prop_equipartitionNatural_sum
+
     describe "partitionNatural" $ do
 
         it "prop_partitionNatural_length" $
@@ -79,6 +91,40 @@ prop_padCoalesce_sum
 prop_padCoalesce_sum source target =
     F.fold source === F.fold (padCoalesce source target)
 
+--------------------------------------------------------------------------------
+-- Equipartitioning natural numbers
+--------------------------------------------------------------------------------
+
+-- Test that natural numbers are equipartitioned fairly:
+--
+-- Each portion must be within unity of the ideal portion.
+--
+prop_equipartitionNatural_fair
+    :: Natural -> NonEmpty () -> Property
+prop_equipartitionNatural_fair n count = (.||.)
+    (difference === 0)
+    (difference === 1)
+  where
+    difference :: Natural
+    difference = F.maximum results - F.minimum results
+
+    results :: NonEmpty Natural
+    results = equipartitionNatural n count
+
+prop_equipartitionNatural_length :: Natural -> NonEmpty () -> Property
+prop_equipartitionNatural_length n count =
+    NE.length (equipartitionNatural n count) === NE.length count
+
+prop_equipartitionNatural_order :: Natural -> NonEmpty () -> Property
+prop_equipartitionNatural_order n count =
+    NE.sort results === results
+  where
+    results = equipartitionNatural n count
+
+prop_equipartitionNatural_sum :: Natural -> NonEmpty () -> Property
+prop_equipartitionNatural_sum n count =
+    F.sum (equipartitionNatural n count) === n
+
 --------------------------------------------------------------------------------
 -- Partitioning natural numbers
 --------------------------------------------------------------------------------

From 03a193c3f5fe007dc9f96c976b8dd1efac8ea15b Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Tue, 2 Mar 2021 03:41:58 +0000
Subject: [PATCH 03/24] Move `equipartitionCoin` to `Coin` module.

---
 .../Primitive/CoinSelection/MA/RoundRobin.hs  | 18 +-----------
 .../Cardano/Wallet/Primitive/Types/Coin.hs    | 29 ++++++++++++++++++-
 2 files changed, 29 insertions(+), 18 deletions(-)

diff --git a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
index f0e25d2ad6c..cb41d3f0f87 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
@@ -1205,22 +1205,6 @@ makeChangeForCoin targets excess =
 --
 --------------------------------------------------------------------------------
 
--- | Computes the equipartition of a coin into 'n' smaller coins.
---
-equipartitionCoin
-    :: HasCallStack
-    => Coin
-    -- ^ The coin to be partitioned.
-    -> NonEmpty a
-    -- ^ Represents the number of portions in which to partition the coin.
-    -> NonEmpty Coin
-    -- ^ The partitioned coins.
-equipartitionCoin c =
-    -- Note: the natural-to-coin conversion is safe, as equipartitioning always
-    -- guarantees to produce values that are less than or equal to the original
-    -- value.
-    fmap unsafeNaturalToCoin . equipartitionNatural (coinToNatural c)
-
 -- | Computes the equipartition of a token map into 'n' smaller maps.
 --
 -- Each asset is partitioned independently.
@@ -1277,7 +1261,7 @@ equipartitionTokenBundleWithMaxQuantity
 equipartitionTokenBundleWithMaxQuantity b maxQuantity =
     NE.zipWith TokenBundle cs ms
   where
-    cs = equipartitionCoin (view #coin b) ms
+    cs = Coin.equipartition (view #coin b) ms
     ms = equipartitionTokenMapWithMaxQuantity (view #tokens b) maxQuantity
 
 -- | Applies 'equipartitionTokenBundleWithMaxQuantity' to a list of bundles.
diff --git a/lib/core/src/Cardano/Wallet/Primitive/Types/Coin.hs b/lib/core/src/Cardano/Wallet/Primitive/Types/Coin.hs
index 68184b0215a..0776ac5f44b 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/Types/Coin.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/Types/Coin.hs
@@ -25,11 +25,14 @@ module Cardano.Wallet.Primitive.Types.Coin
     , subtractCoin
     , sumCoins
     , distance
+    , equipartition
 
     ) where
 
 import Prelude
 
+import Cardano.Numeric.Util
+    ( equipartitionNatural )
 import Control.DeepSeq
     ( NFData (..) )
 import Control.Monad
@@ -38,6 +41,8 @@ import Data.Foldable
     ( foldl' )
 import Data.Hashable
     ( Hashable )
+import Data.List.NonEmpty
+    ( NonEmpty (..) )
 import Data.Quantity
     ( Quantity (..) )
 import Data.Text.Class
@@ -111,6 +116,9 @@ coinToInteger = fromIntegral . unCoin
 coinToNatural :: Coin -> Natural
 coinToNatural = fromIntegral . unCoin
 
+unsafeNaturalToCoin :: Natural -> Coin
+unsafeNaturalToCoin = Coin . fromIntegral
+
 {-------------------------------------------------------------------------------
                                      Checks
 -------------------------------------------------------------------------------}
@@ -146,7 +154,26 @@ addCoin (Coin a) (Coin b) = Coin (a + b)
 sumCoins :: Foldable t => t Coin -> Coin
 sumCoins = foldl' addCoin (Coin 0)
 
-
 -- | Absolute difference between two coin amounts. The result is never negative.
 distance :: Coin -> Coin -> Coin
 distance (Coin a) (Coin b) = if a < b then Coin (b - a) else Coin (a - b)
+
+-- | Computes the equipartition of a coin into 'n' smaller coins.
+--
+-- An /equipartition/ of a coin is a /partition/ of that coin into 'n' smaller
+-- coins whose values differ by no more than 1.
+--
+-- The resultant list is sorted in ascending order.
+--
+equipartition
+    :: Coin
+    -- ^ The coin to be partitioned.
+    -> NonEmpty a
+    -- ^ Represents the number of portions in which to partition the coin.
+    -> NonEmpty Coin
+    -- ^ The partitioned coins.
+equipartition c =
+    -- Note: the natural-to-coin conversion is safe, as equipartitioning always
+    -- guarantees to produce values that are less than or equal to the original
+    -- value.
+    fmap unsafeNaturalToCoin . equipartitionNatural (coinToNatural c)

From f7c46c2fc973c81be1da9b94dca87a7a2fb9147c Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Tue, 2 Mar 2021 03:57:59 +0000
Subject: [PATCH 04/24] Move `equipartitionTokenQuantity` to `TokenQuantity`
 module.

---
 .../Primitive/CoinSelection/MA/RoundRobin.hs  | 17 ++------------
 .../Wallet/Primitive/Types/TokenQuantity.hs   | 22 +++++++++++++++++++
 .../CoinSelection/MA/RoundRobinSpec.hs        |  5 ++---
 3 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
index cb41d3f0f87..654d6b15b51 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
@@ -58,7 +58,6 @@ module Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     , equipartitionTokenBundlesWithMaxQuantity
     , equipartitionTokenMap
     , equipartitionTokenMapWithMaxQuantity
-    , equipartitionTokenQuantity
 
     -- * Grouping and ungrouping
     , groupByKey
@@ -144,6 +143,7 @@ import Numeric.Natural
 import qualified Cardano.Wallet.Primitive.Types.Coin as Coin
 import qualified Cardano.Wallet.Primitive.Types.TokenBundle as TokenBundle
 import qualified Cardano.Wallet.Primitive.Types.TokenMap as TokenMap
+import qualified Cardano.Wallet.Primitive.Types.TokenQuantity as TokenQuantity
 import qualified Cardano.Wallet.Primitive.Types.Tx as Tx
 import qualified Cardano.Wallet.Primitive.Types.UTxOIndex as UTxOIndex
 import qualified Data.Foldable as F
@@ -1226,20 +1226,7 @@ equipartitionTokenMap m count =
         -> NonEmpty TokenMap
     accumulate maps (asset, quantity) = NE.zipWith (<>) maps $
         TokenMap.singleton asset <$>
-            equipartitionTokenQuantity quantity count
-
--- | Computes the equipartition of a token quantity into 'n' smaller quantities.
---
-equipartitionTokenQuantity
-    :: HasCallStack
-    => TokenQuantity
-    -- ^ The token quantity to be partitioned.
-    -> NonEmpty a
-    -- ^ Represents the number of portions in which to partition the quantity.
-    -> NonEmpty TokenQuantity
-    -- ^ The partitioned quantities.
-equipartitionTokenQuantity q =
-    fmap TokenQuantity . equipartitionNatural (unTokenQuantity q)
+            TokenQuantity.equipartition quantity count
 
 --------------------------------------------------------------------------------
 -- Equipartitioning according to a maximum token quantity
diff --git a/lib/core/src/Cardano/Wallet/Primitive/Types/TokenQuantity.hs b/lib/core/src/Cardano/Wallet/Primitive/Types/TokenQuantity.hs
index cc5040bd89a..23e154940b4 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/Types/TokenQuantity.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/Types/TokenQuantity.hs
@@ -16,6 +16,7 @@ module Cardano.Wallet.Primitive.Types.TokenQuantity
     , subtract
     , pred
     , succ
+    , equipartition
 
       -- * Tests
     , isNonZero
@@ -29,6 +30,8 @@ module Cardano.Wallet.Primitive.Types.TokenQuantity
 import Prelude hiding
     ( pred, subtract, succ )
 
+import Cardano.Numeric.Util
+    ( equipartitionNatural )
 import Control.DeepSeq
     ( NFData (..) )
 import Control.Monad
@@ -39,6 +42,8 @@ import Data.Functor
     ( ($>) )
 import Data.Hashable
     ( Hashable )
+import Data.List.NonEmpty
+    ( NonEmpty (..) )
 import Data.Text.Class
     ( FromText (..), ToText (..) )
 import Fmt
@@ -122,6 +127,23 @@ pred (TokenQuantity q) = TokenQuantity $ Prelude.pred q
 succ :: TokenQuantity -> TokenQuantity
 succ (TokenQuantity q) = TokenQuantity $ Prelude.succ q
 
+-- | Computes the equipartition of a token quantity into 'n' smaller quantities.
+--
+-- An /equipartition/ of a token quantity is a /partition/ of that quantity
+-- into 'n' smaller quantities whose values differ by no more than 1.
+--
+-- The resultant list is sorted in ascending order.
+--
+equipartition
+    :: TokenQuantity
+    -- ^ The token quantity to be partitioned.
+    -> NonEmpty a
+    -- ^ Represents the number of portions in which to partition the quantity.
+    -> NonEmpty TokenQuantity
+    -- ^ The partitioned quantities.
+equipartition q =
+    fmap TokenQuantity . equipartitionNatural (unTokenQuantity q)
+
 --------------------------------------------------------------------------------
 -- Tests
 --------------------------------------------------------------------------------
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 54257975d1c..72b82cb4a20 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
@@ -40,7 +40,6 @@ import Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     , equipartitionTokenBundlesWithMaxQuantity
     , equipartitionTokenMap
     , equipartitionTokenMapWithMaxQuantity
-    , equipartitionTokenQuantity
     , fullBalance
     , groupByKey
     , makeChange
@@ -1218,7 +1217,7 @@ boundaryTest2 = BoundaryTestData
     }
   where
     assetA = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "1")
-    q1 :| [q2] = equipartitionTokenQuantity maxTxOutTokenQuantity (() :| [()])
+    q1 :| [q2] = TokenQuantity.equipartition maxTxOutTokenQuantity (() :| [()])
     boundaryTestOutputs =
       [ (Coin 1_500_000, []) ]
     boundaryTestUTxO =
@@ -1247,7 +1246,7 @@ boundaryTest3 = BoundaryTestData
     }
   where
     assetA = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "1")
-    q1 :| [q2] = equipartitionTokenQuantity
+    q1 :| [q2] = TokenQuantity.equipartition
         (TokenQuantity.succ maxTxOutTokenQuantity) (() :| [()])
     boundaryTestOutputs =
       [ (Coin 1_500_000, []) ]

From 407445e62b5d0afe2965e5e5fc16d5ece6d1425a Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Tue, 2 Mar 2021 04:28:30 +0000
Subject: [PATCH 05/24] Move `inAscendingPartialOrder` to `Numeric.Util`.

---
 lib/core/cardano-wallet-core.cabal            |  1 +
 .../CoinSelection/MA/RoundRobinSpec.hs        |  5 ++---
 lib/numeric/cardano-numeric.cabal             |  2 ++
 lib/numeric/src/Cardano/Numeric/Util.hs       | 19 +++++++++++++++++++
 nix/.stack.nix/cardano-numeric.nix            |  6 +++++-
 nix/.stack.nix/cardano-wallet-core.nix        |  1 +
 6 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/lib/core/cardano-wallet-core.cabal b/lib/core/cardano-wallet-core.cabal
index 29c55e87f61..be53338c2e4 100644
--- a/lib/core/cardano-wallet-core.cabal
+++ b/lib/core/cardano-wallet-core.cabal
@@ -243,6 +243,7 @@ test-suite unit
     , cardano-addresses
     , cardano-api
     , cardano-crypto
+    , cardano-numeric
     , cardano-wallet-core
     , cardano-wallet-launcher
     , cardano-wallet-test-utils
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 72b82cb4a20..61303780b92 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
@@ -19,6 +19,8 @@ import Prelude
 
 import Algebra.PartialOrd
     ( PartialOrd (..) )
+import Cardano.Numeric.Util
+    ( inAscendingPartialOrder )
 import Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     ( AssetCount (..)
     , BalanceInsufficientError (..)
@@ -2237,9 +2239,6 @@ consecutivePairs xs = case tailMay xs of
     Nothing -> []
     Just ys -> xs `zip` ys
 
-inAscendingPartialOrder :: (Foldable f, PartialOrd a) => f a -> Bool
-inAscendingPartialOrder = all (uncurry leq) . consecutivePairs . F.toList
-
 addExtraSource :: Maybe Coin -> TokenBundle -> TokenBundle
 addExtraSource extraSource =
     TokenBundle.add
diff --git a/lib/numeric/cardano-numeric.cabal b/lib/numeric/cardano-numeric.cabal
index aceee1794a7..3a02e9b546b 100644
--- a/lib/numeric/cardano-numeric.cabal
+++ b/lib/numeric/cardano-numeric.cabal
@@ -29,6 +29,8 @@ library
     ghc-options: -O2 -Werror
   build-depends:
       base
+    , lattices
+    , safe
   hs-source-dirs:
       src
   exposed-modules:
diff --git a/lib/numeric/src/Cardano/Numeric/Util.hs b/lib/numeric/src/Cardano/Numeric/Util.hs
index 45da3721ec7..da97ab8c6b6 100644
--- a/lib/numeric/src/Cardano/Numeric/Util.hs
+++ b/lib/numeric/src/Cardano/Numeric/Util.hs
@@ -12,11 +12,16 @@ module Cardano.Numeric.Util
     , partitionNatural
     , unsafePartitionNatural
 
+      -- * Partial orders
+    , inAscendingPartialOrder
+
     ) where
 
 import Prelude hiding
     ( round )
 
+import Algebra.PartialOrd
+    ( PartialOrd (..) )
 import Control.Arrow
     ( (&&&) )
 import Data.Function
@@ -33,6 +38,8 @@ import GHC.Stack
     ( HasCallStack )
 import Numeric.Natural
     ( Natural )
+import Safe
+    ( tailMay )
 
 import qualified Data.Foldable as F
 import qualified Data.List.NonEmpty as NE
@@ -238,6 +245,13 @@ unsafePartitionNatural target =
         , "specified weights must have a non-zero sum."
         ]
 
+--------------------------------------------------------------------------------
+-- Partial orders
+--------------------------------------------------------------------------------
+
+inAscendingPartialOrder :: (Foldable f, PartialOrd a) => f a -> Bool
+inAscendingPartialOrder = all (uncurry leq) . consecutivePairs . F.toList
+
 --------------------------------------------------------------------------------
 -- Internal types and functions
 --------------------------------------------------------------------------------
@@ -247,6 +261,11 @@ unsafePartitionNatural target =
 applyN :: Int -> (a -> a) -> a -> a
 applyN n f = F.foldr (.) id (replicate n f)
 
+consecutivePairs :: [a] -> [(a, a)]
+consecutivePairs xs = case tailMay xs of
+    Nothing -> []
+    Just ys -> xs `zip` ys
+
 -- Extract the fractional part of a rational number.
 --
 -- Examples:
diff --git a/nix/.stack.nix/cardano-numeric.nix b/nix/.stack.nix/cardano-numeric.nix
index f1dbbc62a36..a0257761a0a 100644
--- a/nix/.stack.nix/cardano-numeric.nix
+++ b/nix/.stack.nix/cardano-numeric.nix
@@ -25,7 +25,11 @@
       };
     components = {
       "library" = {
-        depends = [ (hsPkgs."base" or (errorHandler.buildDepError "base")) ];
+        depends = [
+          (hsPkgs."base" or (errorHandler.buildDepError "base"))
+          (hsPkgs."lattices" or (errorHandler.buildDepError "lattices"))
+          (hsPkgs."safe" or (errorHandler.buildDepError "safe"))
+          ];
         buildable = true;
         };
       tests = {
diff --git a/nix/.stack.nix/cardano-wallet-core.nix b/nix/.stack.nix/cardano-wallet-core.nix
index c112293c6d6..79faf211c21 100644
--- a/nix/.stack.nix/cardano-wallet-core.nix
+++ b/nix/.stack.nix/cardano-wallet-core.nix
@@ -128,6 +128,7 @@
             (hsPkgs."cardano-addresses" or (errorHandler.buildDepError "cardano-addresses"))
             (hsPkgs."cardano-api" or (errorHandler.buildDepError "cardano-api"))
             (hsPkgs."cardano-crypto" or (errorHandler.buildDepError "cardano-crypto"))
+            (hsPkgs."cardano-numeric" or (errorHandler.buildDepError "cardano-numeric"))
             (hsPkgs."cardano-wallet-core" or (errorHandler.buildDepError "cardano-wallet-core"))
             (hsPkgs."cardano-wallet-launcher" or (errorHandler.buildDepError "cardano-wallet-launcher"))
             (hsPkgs."cardano-wallet-test-utils" or (errorHandler.buildDepError "cardano-wallet-test-utils"))

From e8c07f3e8a0ec8ddce48d9e735c9a4756e7cca2d Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Tue, 2 Mar 2021 04:40:05 +0000
Subject: [PATCH 06/24] Move `equipartitionTokenMap` into `TokenMap`.

---
 .../Primitive/CoinSelection/MA/RoundRobin.hs  | 29 +--------
 .../Wallet/Primitive/Types/TokenMap.hs        | 34 ++++++++++
 .../CoinSelection/MA/RoundRobinSpec.hs        | 61 ------------------
 .../Wallet/Primitive/Types/TokenMapSpec.hs    | 63 ++++++++++++++++++-
 4 files changed, 98 insertions(+), 89 deletions(-)

diff --git a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
index 654d6b15b51..f56c61ae2df 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
@@ -56,7 +56,6 @@ module Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     -- * Partitioning
     , equipartitionTokenBundleWithMaxQuantity
     , equipartitionTokenBundlesWithMaxQuantity
-    , equipartitionTokenMap
     , equipartitionTokenMapWithMaxQuantity
 
     -- * Grouping and ungrouping
@@ -87,7 +86,7 @@ import Prelude
 import Algebra.PartialOrd
     ( PartialOrd (..) )
 import Cardano.Numeric.Util
-    ( equipartitionNatural, padCoalesce, partitionNatural )
+    ( padCoalesce, partitionNatural )
 import Cardano.Wallet.Primitive.Types.Coin
     ( Coin (..), addCoin, subtractCoin, sumCoins )
 import Cardano.Wallet.Primitive.Types.TokenBundle
@@ -143,7 +142,6 @@ import Numeric.Natural
 import qualified Cardano.Wallet.Primitive.Types.Coin as Coin
 import qualified Cardano.Wallet.Primitive.Types.TokenBundle as TokenBundle
 import qualified Cardano.Wallet.Primitive.Types.TokenMap as TokenMap
-import qualified Cardano.Wallet.Primitive.Types.TokenQuantity as TokenQuantity
 import qualified Cardano.Wallet.Primitive.Types.Tx as Tx
 import qualified Cardano.Wallet.Primitive.Types.UTxOIndex as UTxOIndex
 import qualified Data.Foldable as F
@@ -1205,29 +1203,6 @@ makeChangeForCoin targets excess =
 --
 --------------------------------------------------------------------------------
 
--- | Computes the equipartition of a token map into 'n' smaller maps.
---
--- Each asset is partitioned independently.
---
-equipartitionTokenMap
-    :: HasCallStack
-    => TokenMap
-    -- ^ The map to be partitioned.
-    -> NonEmpty a
-    -- ^ Represents the number of portions in which to partition the map.
-    -> NonEmpty TokenMap
-    -- ^ The partitioned maps.
-equipartitionTokenMap m count =
-    F.foldl' accumulate (TokenMap.empty <$ count) (TokenMap.toFlatList m)
-  where
-    accumulate
-        :: NonEmpty TokenMap
-        -> (AssetId, TokenQuantity)
-        -> NonEmpty TokenMap
-    accumulate maps (asset, quantity) = NE.zipWith (<>) maps $
-        TokenMap.singleton asset <$>
-            TokenQuantity.equipartition quantity count
-
 --------------------------------------------------------------------------------
 -- Equipartitioning according to a maximum token quantity
 --------------------------------------------------------------------------------
@@ -1288,7 +1263,7 @@ equipartitionTokenMapWithMaxQuantity m (TokenQuantity maxQuantity)
     | currentMaxQuantity <= maxQuantity =
         m :| []
     | otherwise =
-        equipartitionTokenMap m (() :| replicate extraPartCount ())
+        TokenMap.equipartitionQuantities m (() :| replicate extraPartCount ())
   where
     TokenQuantity currentMaxQuantity = TokenMap.maximumQuantity m
 
diff --git a/lib/core/src/Cardano/Wallet/Primitive/Types/TokenMap.hs b/lib/core/src/Cardano/Wallet/Primitive/Types/TokenMap.hs
index 8411ce0ff35..41844c6fe33 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/Types/TokenMap.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/Types/TokenMap.hs
@@ -71,6 +71,9 @@ module Cardano.Wallet.Primitive.Types.TokenMap
     , removeQuantity
     , maximumQuantity
 
+    -- * Partitioning
+    , equipartitionQuantities
+
     -- * Policies
     , hasPolicy
 
@@ -640,6 +643,37 @@ maximumQuantity =
         | otherwise =
             champion
 
+--------------------------------------------------------------------------------
+-- Partitioning
+--------------------------------------------------------------------------------
+
+-- | Partitions a token map into 'n' smaller maps, where the quantity of each
+--   token is equipartitioned across the resultant maps.
+--
+-- In the resultant maps, the smallest quantity and largest quantity of a given
+-- token will differ by no more than 1.
+--
+-- The resultant list is sorted into ascending order when maps are compared
+-- with the 'leq' function.
+--
+equipartitionQuantities
+    :: TokenMap
+    -- ^ The map to be partitioned.
+    -> NonEmpty a
+    -- ^ Represents the number of portions in which to partition the map.
+    -> NonEmpty TokenMap
+    -- ^ The partitioned maps.
+equipartitionQuantities m count =
+    F.foldl' accumulate (empty <$ count) (toFlatList m)
+  where
+    accumulate
+        :: NonEmpty TokenMap
+        -> (AssetId, TokenQuantity)
+        -> NonEmpty TokenMap
+    accumulate maps (asset, quantity) = NE.zipWith (<>) maps $
+        singleton asset <$>
+            TokenQuantity.equipartition quantity count
+
 --------------------------------------------------------------------------------
 -- Policies
 --------------------------------------------------------------------------------
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 61303780b92..6c4bed3b3b3 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
@@ -40,7 +40,6 @@ import Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     , coinSelectionLens
     , equipartitionTokenBundleWithMaxQuantity
     , equipartitionTokenBundlesWithMaxQuantity
-    , equipartitionTokenMap
     , equipartitionTokenMapWithMaxQuantity
     , fullBalance
     , groupByKey
@@ -321,17 +320,6 @@ spec = describe "Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobinSpec" $
         unitTests "makeChangeForUserSpecifiedAsset"
             unit_makeChangeForUserSpecifiedAsset
 
-    parallel $ describe "Equipartitioning token maps" $ do
-
-        it "prop_equipartitionTokenMap_fair" $
-            property prop_equipartitionTokenMap_fair
-        it "prop_equipartitionTokenMap_length" $
-            property prop_equipartitionTokenMap_length
-        it "prop_equipartitionTokenMap_order" $
-            property prop_equipartitionTokenMap_order
-        it "prop_equipartitionTokenMap_sum" $
-            property prop_equipartitionTokenMap_sum
-
     parallel $ describe "Equipartitioning token bundles by max quantity" $ do
 
         describe "Individual token bundles" $ do
@@ -1863,55 +1851,6 @@ unit_makeChangeForUserSpecifiedAsset =
     assetC :: AssetId
     assetC = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "2")
 
---------------------------------------------------------------------------------
--- Equipartitioning token maps
---------------------------------------------------------------------------------
-
--- Test that token maps are equipartitioned fairly:
---
--- Each token quantity portion must be within unity of the ideal portion.
---
-prop_equipartitionTokenMap_fair :: TokenMap -> NonEmpty () -> Property
-prop_equipartitionTokenMap_fair m count = property $
-    isZeroOrOne maximumDifference
-  where
-    -- Here we take advantage of the fact that the resultant maps are sorted
-    -- into ascending order when compared with the 'leq' function.
-    --
-    -- Consequently:
-    --
-    --  - the head map will be the smallest;
-    --  - the last map will be the greatest.
-    --
-    -- Therefore, subtracting the head map from the last map will produce a map
-    -- where each token quantity is equal to the difference between:
-    --
-    --  - the smallest quantity of that token in the resulting maps;
-    --  - the greatest quantity of that token in the resulting maps.
-    --
-    differences :: TokenMap
-    differences = NE.last results `TokenMap.unsafeSubtract` NE.head results
-
-    isZeroOrOne :: TokenQuantity -> Bool
-    isZeroOrOne (TokenQuantity q) = q == 0 || q == 1
-
-    maximumDifference :: TokenQuantity
-    maximumDifference = TokenMap.maximumQuantity differences
-
-    results = equipartitionTokenMap m count
-
-prop_equipartitionTokenMap_length :: TokenMap -> NonEmpty () -> Property
-prop_equipartitionTokenMap_length m count =
-    NE.length (equipartitionTokenMap m count) === NE.length count
-
-prop_equipartitionTokenMap_order :: TokenMap -> NonEmpty () -> Property
-prop_equipartitionTokenMap_order m count = property $
-    inAscendingPartialOrder (equipartitionTokenMap m count)
-
-prop_equipartitionTokenMap_sum :: TokenMap -> NonEmpty () -> Property
-prop_equipartitionTokenMap_sum m count =
-    F.fold (equipartitionTokenMap m count) === m
-
 --------------------------------------------------------------------------------
 -- Equipartitioning token bundles according to a maximum quantity
 --------------------------------------------------------------------------------
diff --git a/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenMapSpec.hs b/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenMapSpec.hs
index 8f6bf00259c..3514b46376b 100644
--- a/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenMapSpec.hs
+++ b/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenMapSpec.hs
@@ -13,6 +13,8 @@ import Prelude
 
 import Algebra.PartialOrd
     ( PartialOrd (..) )
+import Cardano.Numeric.Util
+    ( inAscendingPartialOrder )
 import Cardano.Wallet.Primitive.Types.TokenMap
     ( AssetId (..), Flat (..), Nested (..), TokenMap, difference )
 import Cardano.Wallet.Primitive.Types.TokenMap.Gen
@@ -103,7 +105,6 @@ import qualified Data.Set as Set
 import qualified Data.Text as T
 import qualified Test.Utils.Roundtrip as Roundtrip
 
-
 spec :: Spec
 spec =
     describe "Token map properties" $
@@ -209,6 +210,17 @@ spec =
         it "prop_maximumQuantity_all" $
             property prop_maximumQuantity_all
 
+    parallel $ describe "Partitioning" $ do
+
+        it "prop_equipartitionQuantities_fair" $
+            property prop_equipartitionQuantities_fair
+        it "prop_equipartitionQuantities_length" $
+            property prop_equipartitionQuantities_length
+        it "prop_equipartitionQuantities_order" $
+            property prop_equipartitionQuantities_order
+        it "prop_equipartitionQuantities_sum" $
+            property prop_equipartitionQuantities_sum
+
     parallel $ describe "JSON serialization" $ do
 
         describe "Roundtrip tests" $ do
@@ -500,6 +512,55 @@ prop_maximumQuantity_all b =
   where
     maxQ = TokenMap.maximumQuantity b
 
+--------------------------------------------------------------------------------
+-- Partitioning
+--------------------------------------------------------------------------------
+
+-- Test that token map quantities are equipartitioned fairly:
+--
+-- Each token quantity portion must be within unity of the ideal portion.
+--
+prop_equipartitionQuantities_fair :: TokenMap -> NonEmpty () -> Property
+prop_equipartitionQuantities_fair m count = property $
+    isZeroOrOne maximumDifference
+  where
+    -- Here we take advantage of the fact that the resultant maps are sorted
+    -- into ascending order when compared with the 'leq' function.
+    --
+    -- Consequently:
+    --
+    --  - the head map will be the smallest;
+    --  - the last map will be the greatest.
+    --
+    -- Therefore, subtracting the head map from the last map will produce a map
+    -- where each token quantity is equal to the difference between:
+    --
+    --  - the smallest quantity of that token in the resulting maps;
+    --  - the greatest quantity of that token in the resulting maps.
+    --
+    differences :: TokenMap
+    differences = NE.last results `TokenMap.unsafeSubtract` NE.head results
+
+    isZeroOrOne :: TokenQuantity -> Bool
+    isZeroOrOne (TokenQuantity q) = q == 0 || q == 1
+
+    maximumDifference :: TokenQuantity
+    maximumDifference = TokenMap.maximumQuantity differences
+
+    results = TokenMap.equipartitionQuantities m count
+
+prop_equipartitionQuantities_length :: TokenMap -> NonEmpty () -> Property
+prop_equipartitionQuantities_length m count =
+    NE.length (TokenMap.equipartitionQuantities m count) === NE.length count
+
+prop_equipartitionQuantities_order :: TokenMap -> NonEmpty () -> Property
+prop_equipartitionQuantities_order m count = property $
+    inAscendingPartialOrder (TokenMap.equipartitionQuantities m count)
+
+prop_equipartitionQuantities_sum :: TokenMap -> NonEmpty () -> Property
+prop_equipartitionQuantities_sum m count =
+    F.fold (TokenMap.equipartitionQuantities m count) === m
+
 --------------------------------------------------------------------------------
 -- JSON serialization tests
 --------------------------------------------------------------------------------

From 96d844bfeca1c8d3382ee0213d9483f0023f8b4c Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Tue, 2 Mar 2021 05:01:44 +0000
Subject: [PATCH 07/24] Move `equipartitionTokenMapWithMaxQuantity` to
 `TokenMap`.

---
 .../Primitive/CoinSelection/MA/RoundRobin.hs  | 38 +-------
 .../Wallet/Primitive/Types/TokenMap.hs        | 35 +++++++
 .../CoinSelection/MA/RoundRobinSpec.hs        | 92 +------------------
 .../Wallet/Primitive/Types/TokenMapSpec.hs    | 90 +++++++++++++++++-
 4 files changed, 131 insertions(+), 124 deletions(-)

diff --git a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
index f56c61ae2df..ebd30affd32 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
@@ -56,7 +56,6 @@ module Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     -- * Partitioning
     , equipartitionTokenBundleWithMaxQuantity
     , equipartitionTokenBundlesWithMaxQuantity
-    , equipartitionTokenMapWithMaxQuantity
 
     -- * Grouping and ungrouping
     , groupByKey
@@ -1220,11 +1219,11 @@ equipartitionTokenBundleWithMaxQuantity
     -- ^ Maximum allowable token quantity.
     -> NonEmpty TokenBundle
     -- ^ The partitioned bundles.
-equipartitionTokenBundleWithMaxQuantity b maxQuantity =
+equipartitionTokenBundleWithMaxQuantity (TokenBundle c m) maxQuantity =
     NE.zipWith TokenBundle cs ms
   where
-    cs = Coin.equipartition (view #coin b) ms
-    ms = equipartitionTokenMapWithMaxQuantity (view #tokens b) maxQuantity
+    cs = Coin.equipartition c ms
+    ms = TokenMap.equipartitionQuantitiesWithUpperBound m maxQuantity
 
 -- | Applies 'equipartitionTokenBundleWithMaxQuantity' to a list of bundles.
 --
@@ -1244,37 +1243,6 @@ equipartitionTokenBundlesWithMaxQuantity
 equipartitionTokenBundlesWithMaxQuantity bs maxQuantity =
     (`equipartitionTokenBundleWithMaxQuantity` maxQuantity) =<< bs
 
--- | Computes the equipartition of a token map into 'n' smaller maps, according
---   to the given maximum token quantity.
---
--- The value 'n' is computed automatically, and is the minimum value required
--- to achieve the goal that no token quantity in any of the resulting maps
--- exceeds the maximum allowable token quantity.
---
-equipartitionTokenMapWithMaxQuantity
-    :: TokenMap
-    -> TokenQuantity
-    -- ^ Maximum allowable token quantity.
-    -> NonEmpty TokenMap
-    -- ^ The partitioned maps.
-equipartitionTokenMapWithMaxQuantity m (TokenQuantity maxQuantity)
-    | maxQuantity == 0 =
-        maxQuantityZeroError
-    | currentMaxQuantity <= maxQuantity =
-        m :| []
-    | otherwise =
-        TokenMap.equipartitionQuantities m (() :| replicate extraPartCount ())
-  where
-    TokenQuantity currentMaxQuantity = TokenMap.maximumQuantity m
-
-    extraPartCount :: Int
-    extraPartCount = floor $ pred currentMaxQuantity % maxQuantity
-
-    maxQuantityZeroError = error $ unwords
-        [ "equipartitionTokenMapWithMaxQuantity:"
-        , "the maximum allowable token quantity cannot be zero."
-        ]
-
 --------------------------------------------------------------------------------
 -- Grouping and ungrouping
 --------------------------------------------------------------------------------
diff --git a/lib/core/src/Cardano/Wallet/Primitive/Types/TokenMap.hs b/lib/core/src/Cardano/Wallet/Primitive/Types/TokenMap.hs
index 41844c6fe33..b20c6786cb3 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/Types/TokenMap.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/Types/TokenMap.hs
@@ -73,6 +73,7 @@ module Cardano.Wallet.Primitive.Types.TokenMap
 
     -- * Partitioning
     , equipartitionQuantities
+    , equipartitionQuantitiesWithUpperBound
 
     -- * Policies
     , hasPolicy
@@ -120,6 +121,8 @@ import Data.Map.Strict.NonEmptyMap
     ( NonEmptyMap )
 import Data.Maybe
     ( fromMaybe, isJust )
+import Data.Ratio
+    ( (%) )
 import Data.Set
     ( Set )
 import Data.Text.Class
@@ -674,6 +677,38 @@ equipartitionQuantities m count =
         singleton asset <$>
             TokenQuantity.equipartition quantity count
 
+-- | Partitions a token map into 'n' smaller maps, where the quantity of each
+--   token is equipartitioned across the resultant maps, with the goal that no
+--   token quantity in any of the resultant maps exceeds the given upper bound.
+--
+-- The value 'n' is computed automatically, and is the minimum value required
+-- to achieve the goal that no token quantity in any of the resulting maps
+-- exceeds the maximum allowable token quantity.
+--
+equipartitionQuantitiesWithUpperBound
+    :: TokenMap
+    -> TokenQuantity
+    -- ^ Maximum allowable token quantity.
+    -> NonEmpty TokenMap
+    -- ^ The partitioned maps.
+equipartitionQuantitiesWithUpperBound m (TokenQuantity maxQuantity)
+    | maxQuantity == 0 =
+        maxQuantityZeroError
+    | currentMaxQuantity <= maxQuantity =
+        m :| []
+    | otherwise =
+        equipartitionQuantities m (() :| replicate extraPartCount ())
+  where
+    TokenQuantity currentMaxQuantity = maximumQuantity m
+
+    extraPartCount :: Int
+    extraPartCount = floor $ pred currentMaxQuantity % maxQuantity
+
+    maxQuantityZeroError = error $ unwords
+        [ "equipartitionQuantitiesWithUpperBound:"
+        , "the maximum allowable token quantity cannot be zero."
+        ]
+
 --------------------------------------------------------------------------------
 -- Policies
 --------------------------------------------------------------------------------
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 6c4bed3b3b3..aca94d14dd6 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
@@ -40,7 +40,6 @@ import Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     , coinSelectionLens
     , equipartitionTokenBundleWithMaxQuantity
     , equipartitionTokenBundlesWithMaxQuantity
-    , equipartitionTokenMapWithMaxQuantity
     , fullBalance
     , groupByKey
     , makeChange
@@ -338,19 +337,6 @@ spec = describe "Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobinSpec" $
             it "prop_equipartitionTokenBundlesWithMaxQuantity_sum" $
                 property prop_equipartitionTokenBundlesWithMaxQuantity_sum
 
-    parallel $ describe "Equipartitioning token maps by max quantity" $ do
-
-        it "prop_equipartitionTokenMapWithMaxQuantity_coverage" $
-            property prop_equipartitionTokenMapWithMaxQuantity_coverage
-        it "prop_equipartitionTokenMapWithMaxQuantity_length" $
-            property prop_equipartitionTokenMapWithMaxQuantity_length
-        it "prop_equipartitionTokenMapWithMaxQuantity_max" $
-            property prop_equipartitionTokenMapWithMaxQuantity_max
-        it "prop_equipartitionTokenMapWithMaxQuantity_order" $
-            property prop_equipartitionTokenMapWithMaxQuantity_order
-        it "prop_equipartitionTokenMapWithMaxQuantity_sum" $
-            property prop_equipartitionTokenMapWithMaxQuantity_sum
-
     parallel $ describe "Grouping and ungrouping" $ do
 
         it "prop_groupByKey_ungroupByKey" $
@@ -1860,9 +1846,11 @@ unit_makeChangeForUserSpecifiedAsset =
 --
 equipartitionTokenBundleWithMaxQuantity_expectedLength
     :: TokenBundle -> TokenQuantity -> Int
-equipartitionTokenBundleWithMaxQuantity_expectedLength m =
-    equipartitionTokenMapWithMaxQuantity_expectedLength
-        (view #tokens m)
+equipartitionTokenBundleWithMaxQuantity_expectedLength
+    (TokenBundle _ m) (TokenQuantity maxQuantity) =
+        max 1 $ ceiling $ currentMaxQuantity % maxQuantity
+  where
+    TokenQuantity currentMaxQuantity = TokenMap.maximumQuantity m
 
 prop_equipartitionTokenBundleWithMaxQuantity_length
     :: TokenBundle -> TokenQuantity -> Property
@@ -1925,76 +1913,6 @@ prop_equipartitionTokenBundlesWithMaxQuantity_sum ms maxQuantity =
         F.fold (equipartitionTokenBundlesWithMaxQuantity ms maxQuantity)
             === F.fold ms
 
---------------------------------------------------------------------------------
--- Equipartitioning token maps according to a maximum quantity
---------------------------------------------------------------------------------
-
--- | Computes the number of parts that 'equipartitionTokenMapWithMaxQuantity'
---   should return.
---
-equipartitionTokenMapWithMaxQuantity_expectedLength
-    :: TokenMap -> TokenQuantity -> Int
-equipartitionTokenMapWithMaxQuantity_expectedLength
-    m (TokenQuantity maxQuantity) =
-        max 1 $ ceiling $ currentMaxQuantity % maxQuantity
-  where
-    TokenQuantity currentMaxQuantity = TokenMap.maximumQuantity m
-
-prop_equipartitionTokenMapWithMaxQuantity_coverage
-    :: TokenMap -> TokenQuantity -> Property
-prop_equipartitionTokenMapWithMaxQuantity_coverage m maxQuantity =
-    maxQuantity > TokenQuantity.zero ==>
-        checkCoverage $
-        cover 8 (maxQuantity == TokenQuantity 1)
-            "Maximum allowable quantity == 1" $
-        cover 8 (maxQuantity == TokenQuantity 2)
-            "Maximum allowable quantity == 2" $
-        cover 8 (maxQuantity >= TokenQuantity 3)
-            "Maximum allowable quantity >= 3" $
-        cover 8 (expectedLength == 1)
-            "Expected number of parts == 1" $
-        cover 8 (expectedLength == 2)
-            "Expected number of parts == 2" $
-        cover 8 (expectedLength >= 3)
-            "Expected number of parts >= 3" $
-        property $ expectedLength > 0
-  where
-    expectedLength = equipartitionTokenMapWithMaxQuantity_expectedLength
-        m maxQuantity
-
-prop_equipartitionTokenMapWithMaxQuantity_length
-    :: TokenMap -> TokenQuantity -> Property
-prop_equipartitionTokenMapWithMaxQuantity_length m maxQuantity =
-    maxQuantity > TokenQuantity.zero ==>
-        length (equipartitionTokenMapWithMaxQuantity m maxQuantity)
-            === equipartitionTokenMapWithMaxQuantity_expectedLength
-                m maxQuantity
-
-prop_equipartitionTokenMapWithMaxQuantity_max
-    :: TokenMap -> TokenQuantity -> Property
-prop_equipartitionTokenMapWithMaxQuantity_max m maxQuantity =
-    maxQuantity > TokenQuantity.zero ==>
-        checkCoverage $
-        cover 10 (maxResultQuantity == maxQuantity)
-            "At least one resultant token map has a maximal quantity" $
-        property $ maxResultQuantity <= maxQuantity
-  where
-    results = equipartitionTokenMapWithMaxQuantity m maxQuantity
-    maxResultQuantity = F.maximum (TokenMap.maximumQuantity <$> results)
-
-prop_equipartitionTokenMapWithMaxQuantity_order
-    :: TokenMap -> TokenQuantity -> Property
-prop_equipartitionTokenMapWithMaxQuantity_order m maxQuantity =
-    maxQuantity > TokenQuantity.zero ==>
-        inAscendingPartialOrder
-            (equipartitionTokenMapWithMaxQuantity m maxQuantity)
-
-prop_equipartitionTokenMapWithMaxQuantity_sum
-    :: TokenMap -> TokenQuantity -> Property
-prop_equipartitionTokenMapWithMaxQuantity_sum m maxQuantity =
-    maxQuantity > TokenQuantity.zero ==>
-        F.fold (equipartitionTokenMapWithMaxQuantity m maxQuantity) === m
-
 --------------------------------------------------------------------------------
 -- Grouping and ungrouping
 --------------------------------------------------------------------------------
diff --git a/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenMapSpec.hs b/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenMapSpec.hs
index 3514b46376b..91ba2596094 100644
--- a/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenMapSpec.hs
+++ b/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenMapSpec.hs
@@ -54,6 +54,8 @@ import Data.Maybe
     ( mapMaybe )
 import Data.Proxy
     ( Proxy (..) )
+import Data.Ratio
+    ( (%) )
 import Data.String.QQ
     ( s )
 import Data.Text
@@ -210,7 +212,7 @@ spec =
         it "prop_maximumQuantity_all" $
             property prop_maximumQuantity_all
 
-    parallel $ describe "Partitioning" $ do
+    parallel $ describe "Partitioning quantities" $ do
 
         it "prop_equipartitionQuantities_fair" $
             property prop_equipartitionQuantities_fair
@@ -221,6 +223,19 @@ spec =
         it "prop_equipartitionQuantities_sum" $
             property prop_equipartitionQuantities_sum
 
+    parallel $ describe "Partitioning quantities with an upper bound" $ do
+
+        it "prop_equipartitionQuantitiesWithUpperBound_coverage" $
+            property prop_equipartitionQuantitiesWithUpperBound_coverage
+        it "prop_equipartitionQuantitiesWithUpperBound_length" $
+            property prop_equipartitionQuantitiesWithUpperBound_length
+        it "prop_equipartitionQuantitiesWithUpperBound_max" $
+            property prop_equipartitionQuantitiesWithUpperBound_max
+        it "prop_equipartitionQuantitiesWithUpperBound_order" $
+            property prop_equipartitionQuantitiesWithUpperBound_order
+        it "prop_equipartitionQuantitiesWithUpperBound_sum" $
+            property prop_equipartitionQuantitiesWithUpperBound_sum
+
     parallel $ describe "JSON serialization" $ do
 
         describe "Roundtrip tests" $ do
@@ -513,7 +528,7 @@ prop_maximumQuantity_all b =
     maxQ = TokenMap.maximumQuantity b
 
 --------------------------------------------------------------------------------
--- Partitioning
+-- Partitioning quantities
 --------------------------------------------------------------------------------
 
 -- Test that token map quantities are equipartitioned fairly:
@@ -561,6 +576,77 @@ prop_equipartitionQuantities_sum :: TokenMap -> NonEmpty () -> Property
 prop_equipartitionQuantities_sum m count =
     F.fold (TokenMap.equipartitionQuantities m count) === m
 
+--------------------------------------------------------------------------------
+-- Partitioning quantities according to an upper bound
+--------------------------------------------------------------------------------
+
+-- | Computes the number of parts that 'equipartitionQuantitiesWithUpperBound'
+--   should return.
+--
+equipartitionQuantitiesWithUpperBound_expectedLength
+    :: TokenMap -> TokenQuantity -> Int
+equipartitionQuantitiesWithUpperBound_expectedLength
+    m (TokenQuantity maxQuantity) =
+        max 1 $ ceiling $ currentMaxQuantity % maxQuantity
+  where
+    TokenQuantity currentMaxQuantity = TokenMap.maximumQuantity m
+
+prop_equipartitionQuantitiesWithUpperBound_coverage
+    :: TokenMap -> TokenQuantity -> Property
+prop_equipartitionQuantitiesWithUpperBound_coverage m maxQuantity =
+    maxQuantity > TokenQuantity.zero ==>
+        checkCoverage $
+        cover 8 (maxQuantity == TokenQuantity 1)
+            "Maximum allowable quantity == 1" $
+        cover 8 (maxQuantity == TokenQuantity 2)
+            "Maximum allowable quantity == 2" $
+        cover 8 (maxQuantity >= TokenQuantity 3)
+            "Maximum allowable quantity >= 3" $
+        cover 8 (expectedLength == 1)
+            "Expected number of parts == 1" $
+        cover 8 (expectedLength == 2)
+            "Expected number of parts == 2" $
+        cover 8 (expectedLength >= 3)
+            "Expected number of parts >= 3" $
+        property $ expectedLength > 0
+  where
+    expectedLength = equipartitionQuantitiesWithUpperBound_expectedLength
+        m maxQuantity
+
+prop_equipartitionQuantitiesWithUpperBound_length
+    :: TokenMap -> TokenQuantity -> Property
+prop_equipartitionQuantitiesWithUpperBound_length m maxQuantity =
+    maxQuantity > TokenQuantity.zero ==>
+        length (TokenMap.equipartitionQuantitiesWithUpperBound m maxQuantity)
+            === equipartitionQuantitiesWithUpperBound_expectedLength
+                m maxQuantity
+
+prop_equipartitionQuantitiesWithUpperBound_max
+    :: TokenMap -> TokenQuantity -> Property
+prop_equipartitionQuantitiesWithUpperBound_max m maxQuantity =
+    maxQuantity > TokenQuantity.zero ==>
+        checkCoverage $
+        cover 10 (maxResultQuantity == maxQuantity)
+            "At least one resultant token map has a maximal quantity" $
+        property $ maxResultQuantity <= maxQuantity
+  where
+    results = TokenMap.equipartitionQuantitiesWithUpperBound m maxQuantity
+    maxResultQuantity = F.maximum (TokenMap.maximumQuantity <$> results)
+
+prop_equipartitionQuantitiesWithUpperBound_order
+    :: TokenMap -> TokenQuantity -> Property
+prop_equipartitionQuantitiesWithUpperBound_order m maxQuantity =
+    maxQuantity > TokenQuantity.zero ==>
+        inAscendingPartialOrder
+            (TokenMap.equipartitionQuantitiesWithUpperBound m maxQuantity)
+
+prop_equipartitionQuantitiesWithUpperBound_sum
+    :: TokenMap -> TokenQuantity -> Property
+prop_equipartitionQuantitiesWithUpperBound_sum m maxQuantity =
+    maxQuantity > TokenQuantity.zero ==>
+        F.fold (TokenMap.equipartitionQuantitiesWithUpperBound m maxQuantity)
+            === m
+
 --------------------------------------------------------------------------------
 -- JSON serialization tests
 --------------------------------------------------------------------------------

From 928ca163fa89710d8daf5185e6a8e72e17f39c45 Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Tue, 2 Mar 2021 05:29:39 +0000
Subject: [PATCH 08/24] Move `equipartitionTokenBundleWithMaxQuantity` to
 `TokenBundle`.

---
 .../Primitive/CoinSelection/MA/RoundRobin.hs  | 27 +-------
 .../Wallet/Primitive/Types/TokenBundle.hs     | 29 +++++++++
 .../CoinSelection/MA/RoundRobinSpec.hs        | 46 -------------
 .../Wallet/Primitive/Types/TokenBundleSpec.hs | 64 ++++++++++++++++++-
 4 files changed, 95 insertions(+), 71 deletions(-)

diff --git a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
index ebd30affd32..04413c4499d 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
@@ -54,7 +54,6 @@ module Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     , assignCoinsToChangeMaps
 
     -- * Partitioning
-    , equipartitionTokenBundleWithMaxQuantity
     , equipartitionTokenBundlesWithMaxQuantity
 
     -- * Grouping and ungrouping
@@ -116,8 +115,6 @@ import Data.Maybe
     ( fromMaybe )
 import Data.Ord
     ( comparing )
-import Data.Ratio
-    ( (%) )
 import Data.Set
     ( Set )
 import Data.Word
@@ -1206,26 +1203,8 @@ makeChangeForCoin targets excess =
 -- Equipartitioning according to a maximum token quantity
 --------------------------------------------------------------------------------
 
--- | Computes the equipartition of a token bundle into 'n' smaller bundles,
---   according to the given maximum token quantity.
---
--- The value 'n' is computed automatically, and is the minimum value required
--- to achieve the goal that no token quantity in any of the resulting bundles
--- exceeds the maximum allowable token quantity.
---
-equipartitionTokenBundleWithMaxQuantity
-    :: TokenBundle
-    -> TokenQuantity
-    -- ^ Maximum allowable token quantity.
-    -> NonEmpty TokenBundle
-    -- ^ The partitioned bundles.
-equipartitionTokenBundleWithMaxQuantity (TokenBundle c m) maxQuantity =
-    NE.zipWith TokenBundle cs ms
-  where
-    cs = Coin.equipartition c ms
-    ms = TokenMap.equipartitionQuantitiesWithUpperBound m maxQuantity
-
--- | Applies 'equipartitionTokenBundleWithMaxQuantity' to a list of bundles.
+-- | Applies 'TokenBundle.equipartitionQuantitiesWithUpperBound' to a list of
+--   bundles.
 --
 -- Only token bundles containing quantities that exceed the maximum token
 -- quantity will be partitioned.
@@ -1241,7 +1220,7 @@ equipartitionTokenBundlesWithMaxQuantity
     -> NonEmpty TokenBundle
     -- ^ The partitioned bundles.
 equipartitionTokenBundlesWithMaxQuantity bs maxQuantity =
-    (`equipartitionTokenBundleWithMaxQuantity` maxQuantity) =<< bs
+    (`TokenBundle.equipartitionQuantitiesWithUpperBound` maxQuantity) =<< bs
 
 --------------------------------------------------------------------------------
 -- Grouping and ungrouping
diff --git a/lib/core/src/Cardano/Wallet/Primitive/Types/TokenBundle.hs b/lib/core/src/Cardano/Wallet/Primitive/Types/TokenBundle.hs
index f4c7b10bcac..ab38a4e4144 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/Types/TokenBundle.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/Types/TokenBundle.hs
@@ -53,6 +53,9 @@ module Cardano.Wallet.Primitive.Types.TokenBundle
     , adjustQuantity
     , removeQuantity
 
+    -- * Partitioning
+    , equipartitionQuantitiesWithUpperBound
+
     -- * Policies
     , hasPolicy
 
@@ -110,6 +113,7 @@ import GHC.TypeLits
 
 import qualified Cardano.Wallet.Primitive.Types.Coin as Coin
 import qualified Cardano.Wallet.Primitive.Types.TokenMap as TokenMap
+import qualified Data.List.NonEmpty as NE
 
 --------------------------------------------------------------------------------
 -- Types
@@ -365,6 +369,31 @@ adjustQuantity b a f = b { tokens = TokenMap.adjustQuantity (tokens b) a f }
 removeQuantity :: TokenBundle -> AssetId -> TokenBundle
 removeQuantity b a = b { tokens = TokenMap.removeQuantity (tokens b) a }
 
+--------------------------------------------------------------------------------
+-- Partitioning
+--------------------------------------------------------------------------------
+
+-- | Partitions a token bundle into 'n' smaller bundles, where the quantity of
+--   each token is equipartitioned across the resultant bundles, with the goal
+--   that no token quantity in any of the resultant bundles exceeds the given
+--   upper bound.
+--
+-- The value 'n' is computed automatically, and is the minimum value required
+-- to achieve the goal that no token quantity in any of the resulting bundles
+-- exceeds the maximum allowable token quantity.
+--
+equipartitionQuantitiesWithUpperBound
+    :: TokenBundle
+    -> TokenQuantity
+    -- ^ Maximum allowable token quantity.
+    -> NonEmpty TokenBundle
+    -- ^ The partitioned bundles.
+equipartitionQuantitiesWithUpperBound (TokenBundle c m) maxQuantity =
+    NE.zipWith TokenBundle cs ms
+  where
+    cs = Coin.equipartition c ms
+    ms = TokenMap.equipartitionQuantitiesWithUpperBound m maxQuantity
+
 --------------------------------------------------------------------------------
 -- Policies
 --------------------------------------------------------------------------------
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 aca94d14dd6..0e14abf1a1a 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,6 @@ import Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     , assignCoinsToChangeMaps
     , balanceMissing
     , coinSelectionLens
-    , equipartitionTokenBundleWithMaxQuantity
     , equipartitionTokenBundlesWithMaxQuantity
     , fullBalance
     , groupByKey
@@ -321,15 +320,6 @@ spec = describe "Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobinSpec" $
 
     parallel $ describe "Equipartitioning token bundles by max quantity" $ do
 
-        describe "Individual token bundles" $ do
-
-            it "prop_equipartitionTokenBundleWithMaxQuantity_length" $
-                property prop_equipartitionTokenBundleWithMaxQuantity_length
-            it "prop_equipartitionTokenBundleWithMaxQuantity_order" $
-                property prop_equipartitionTokenBundleWithMaxQuantity_order
-            it "prop_equipartitionTokenBundleWithMaxQuantity_sum" $
-                property prop_equipartitionTokenBundleWithMaxQuantity_sum
-
         describe "Lists of token bundles" $ do
 
             it "prop_equipartitionTokenBundlesWithMaxQuantity_length" $
@@ -1837,42 +1827,6 @@ unit_makeChangeForUserSpecifiedAsset =
     assetC :: AssetId
     assetC = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "2")
 
---------------------------------------------------------------------------------
--- Equipartitioning token bundles according to a maximum quantity
---------------------------------------------------------------------------------
-
--- | Computes the number of parts that 'equipartitionTokenBundleWithMaxQuantity'
---   should return.
---
-equipartitionTokenBundleWithMaxQuantity_expectedLength
-    :: TokenBundle -> TokenQuantity -> Int
-equipartitionTokenBundleWithMaxQuantity_expectedLength
-    (TokenBundle _ m) (TokenQuantity maxQuantity) =
-        max 1 $ ceiling $ currentMaxQuantity % maxQuantity
-  where
-    TokenQuantity currentMaxQuantity = TokenMap.maximumQuantity m
-
-prop_equipartitionTokenBundleWithMaxQuantity_length
-    :: TokenBundle -> TokenQuantity -> Property
-prop_equipartitionTokenBundleWithMaxQuantity_length m maxQuantity =
-    maxQuantity > TokenQuantity.zero ==>
-        length (equipartitionTokenBundleWithMaxQuantity m maxQuantity)
-            === equipartitionTokenBundleWithMaxQuantity_expectedLength
-                m maxQuantity
-
-prop_equipartitionTokenBundleWithMaxQuantity_order
-    :: TokenBundle -> TokenQuantity -> Property
-prop_equipartitionTokenBundleWithMaxQuantity_order m maxQuantity =
-    maxQuantity > TokenQuantity.zero ==>
-        inAscendingPartialOrder
-            (equipartitionTokenBundleWithMaxQuantity m maxQuantity)
-
-prop_equipartitionTokenBundleWithMaxQuantity_sum
-    :: TokenBundle -> TokenQuantity -> Property
-prop_equipartitionTokenBundleWithMaxQuantity_sum m maxQuantity =
-    maxQuantity > TokenQuantity.zero ==>
-        F.fold (equipartitionTokenBundleWithMaxQuantity m maxQuantity) === m
-
 --------------------------------------------------------------------------------
 -- Equipartitioning lists of token bundles according to a maximum quantity
 --------------------------------------------------------------------------------
diff --git a/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenBundleSpec.hs b/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenBundleSpec.hs
index b1550316dc9..06c5a996836 100644
--- a/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenBundleSpec.hs
+++ b/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenBundleSpec.hs
@@ -12,9 +12,17 @@ import Prelude hiding
 import Algebra.PartialOrd
     ( leq )
 import Cardano.Wallet.Primitive.Types.TokenBundle
-    ( TokenBundle, add, difference, isCoin, subtract, unsafeSubtract )
+    ( TokenBundle (..), add, difference, isCoin, subtract, unsafeSubtract )
+import Cardano.Numeric.Util
+    ( inAscendingPartialOrder )
 import Cardano.Wallet.Primitive.Types.TokenBundle.Gen
     ( genTokenBundleSmallRange, shrinkTokenBundleSmallRange )
+import Cardano.Wallet.Primitive.Types.TokenQuantity
+    ( TokenQuantity (..) )
+import Cardano.Wallet.Primitive.Types.TokenQuantity.Gen
+    ( genTokenQuantitySmallPositive, shrinkTokenQuantitySmallPositive )
+import Data.Ratio
+    ( (%) )
 import Test.Hspec
     ( Spec, describe, it )
 import Test.Hspec.Core.QuickCheck
@@ -36,6 +44,11 @@ import Test.Utils.Laws
 import Test.Utils.Laws.PartialOrd
     ( partialOrdLaws )
 
+import qualified Cardano.Wallet.Primitive.Types.TokenBundle as TokenBundle
+import qualified Cardano.Wallet.Primitive.Types.TokenMap as TokenMap
+import qualified Cardano.Wallet.Primitive.Types.TokenQuantity as TokenQuantity
+import qualified Data.Foldable as F
+
 spec :: Spec
 spec =
     describe "Token bundle properties" $
@@ -66,6 +79,14 @@ spec =
         it "prop_difference_equality" $
             property prop_difference_equality
 
+    describe "Partitioning quantities with an upper bound" $ do
+        it "prop_equipartitionQuantitiesWithUpperBound_length" $
+            property prop_equipartitionQuantitiesWithUpperBound_length
+        it "prop_equipartitionQuantitiesWithUpperBound_order" $
+            property prop_equipartitionQuantitiesWithUpperBound_order
+        it "prop_equipartitionQuantitiesWithUpperBound_sum" $
+            property prop_equipartitionQuantitiesWithUpperBound_sum
+
 --------------------------------------------------------------------------------
 -- Arithmetic properties
 --------------------------------------------------------------------------------
@@ -115,6 +136,43 @@ prop_difference_equality x y = checkCoverage $
     xExcess = x `difference` y
     yExcess = y `difference` x
 
+--------------------------------------------------------------------------------
+-- Partitioning quantities according to an upper bound
+--------------------------------------------------------------------------------
+
+-- | Computes the number of parts that 'equipartitionQuantitiesWithUpperBound'
+--   should return.
+--
+equipartitionQuantitiesWithUpperBound_expectedLength
+    :: TokenBundle -> TokenQuantity -> Int
+equipartitionQuantitiesWithUpperBound_expectedLength
+    (TokenBundle _ m) (TokenQuantity maxQuantity) =
+        max 1 $ ceiling $ currentMaxQuantity % maxQuantity
+  where
+    TokenQuantity currentMaxQuantity = TokenMap.maximumQuantity m
+
+prop_equipartitionQuantitiesWithUpperBound_length
+    :: TokenBundle -> TokenQuantity -> Property
+prop_equipartitionQuantitiesWithUpperBound_length m maxQuantity =
+    maxQuantity > TokenQuantity.zero ==>
+        length (TokenBundle.equipartitionQuantitiesWithUpperBound m maxQuantity)
+            === equipartitionQuantitiesWithUpperBound_expectedLength
+                m maxQuantity
+
+prop_equipartitionQuantitiesWithUpperBound_order
+    :: TokenBundle -> TokenQuantity -> Property
+prop_equipartitionQuantitiesWithUpperBound_order m maxQuantity =
+    maxQuantity > TokenQuantity.zero ==>
+        inAscendingPartialOrder
+            (TokenBundle.equipartitionQuantitiesWithUpperBound m maxQuantity)
+
+prop_equipartitionQuantitiesWithUpperBound_sum
+    :: TokenBundle -> TokenQuantity -> Property
+prop_equipartitionQuantitiesWithUpperBound_sum m maxQuantity =
+    maxQuantity > TokenQuantity.zero ==>
+        F.fold (TokenBundle.equipartitionQuantitiesWithUpperBound m maxQuantity)
+            === m
+
 --------------------------------------------------------------------------------
 -- Arbitrary instances
 --------------------------------------------------------------------------------
@@ -122,3 +180,7 @@ prop_difference_equality x y = checkCoverage $
 instance Arbitrary TokenBundle where
     arbitrary = genTokenBundleSmallRange
     shrink = shrinkTokenBundleSmallRange
+
+instance Arbitrary TokenQuantity where
+    arbitrary = genTokenQuantitySmallPositive
+    shrink = shrinkTokenQuantitySmallPositive

From 59a039b33dd98ec5c58e66fc5c38418b1fa154d2 Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Tue, 2 Mar 2021 05:41:20 +0000
Subject: [PATCH 09/24] Rename `equipartitionTokenBundlesWithMaxQuantity`.

Rename to `splitBundlesWithExcessiveTokenQuantities`.
---
 .../Primitive/CoinSelection/MA/RoundRobin.hs  | 41 ++++---------------
 .../CoinSelection/MA/RoundRobinSpec.hs        | 30 ++++++--------
 2 files changed, 21 insertions(+), 50 deletions(-)

diff --git a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
index 04413c4499d..bbbd5b690d7 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
@@ -53,8 +53,8 @@ module Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     , makeChangeForNonUserSpecifiedAsset
     , assignCoinsToChangeMaps
 
-    -- * Partitioning
-    , equipartitionTokenBundlesWithMaxQuantity
+    -- * Splitting bundles
+    , splitBundlesWithExcessiveTokenQuantities
 
     -- * Grouping and ungrouping
     , groupByKey
@@ -888,7 +888,7 @@ makeChange minCoinFor requiredCost mExtraCoinSource inputBundles outputBundles
           where
             bundle (m, c) = TokenBundle c m
             unbundle (TokenBundle c m) = (m, c)
-            split = flip equipartitionTokenBundlesWithMaxQuantity
+            split = flip splitBundlesWithExcessiveTokenQuantities
                 maxTxOutTokenQuantity
 
     -- Change for user-specified assets: assets that were present in the
@@ -1176,50 +1176,25 @@ makeChangeForCoin targets excess =
     weights = coinToNatural <$> targets
 
 --------------------------------------------------------------------------------
--- Equipartitioning
+-- Splitting bundles
 --------------------------------------------------------------------------------
 
--- An /equipartition/ of a value 'v' (of some type) is a /partition/ of that
--- value into 'n' smaller values whose /sizes/ differ by no more than 1. The
--- the notion of /size/ is dependent on the type of value 'v'.
---
--- In this section, equipartitions have the following properties:
---
---  1.  The length is observed:
---      >>> length (equipartition v n) == n
---
---  2.  The sum is preserved:
---      >>> sum (equipartition v n) == v
---
---  3.  Each resulting value is less than or equal to the original value:
---      >>> all (`leq` v) (equipartition v n)
---
---  4.  The resultant list is sorted into ascending order when values are
---      compared with the 'leq' function.
---
---------------------------------------------------------------------------------
-
---------------------------------------------------------------------------------
--- Equipartitioning according to a maximum token quantity
---------------------------------------------------------------------------------
-
--- | Applies 'TokenBundle.equipartitionQuantitiesWithUpperBound' to a list of
---   bundles.
+-- | Splits bundles with excessive token quantities into smaller bundles.
 --
 -- Only token bundles containing quantities that exceed the maximum token
--- quantity will be partitioned.
+-- quantity will be split.
 --
 -- If none of the bundles in the given list contain a quantity that exceeds
 -- the maximum token quantity, this function will return the original list.
 --
-equipartitionTokenBundlesWithMaxQuantity
+splitBundlesWithExcessiveTokenQuantities
     :: NonEmpty TokenBundle
     -- ^ Token bundles.
     -> TokenQuantity
     -- ^ Maximum allowable token quantity.
     -> NonEmpty TokenBundle
     -- ^ The partitioned bundles.
-equipartitionTokenBundlesWithMaxQuantity bs maxQuantity =
+splitBundlesWithExcessiveTokenQuantities bs maxQuantity =
     (`TokenBundle.equipartitionQuantitiesWithUpperBound` maxQuantity) =<< bs
 
 --------------------------------------------------------------------------------
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 0e14abf1a1a..cc4f046bcf0 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,6 @@ import Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     , assignCoinsToChangeMaps
     , balanceMissing
     , coinSelectionLens
-    , equipartitionTokenBundlesWithMaxQuantity
     , fullBalance
     , groupByKey
     , makeChange
@@ -52,6 +51,7 @@ import Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     , runRoundRobin
     , runSelection
     , runSelectionStep
+    , splitBundlesWithExcessiveTokenQuantities
     , ungroupByKey
     )
 import Cardano.Wallet.Primitive.Types.Address
@@ -112,8 +112,6 @@ import Data.Map.Strict
     ( Map )
 import Data.Maybe
     ( isJust )
-import Data.Ratio
-    ( (%) )
 import Data.Set
     ( Set )
 import Data.Tuple
@@ -318,14 +316,12 @@ spec = describe "Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobinSpec" $
         unitTests "makeChangeForUserSpecifiedAsset"
             unit_makeChangeForUserSpecifiedAsset
 
-    parallel $ describe "Equipartitioning token bundles by max quantity" $ do
+    parallel $ describe "Splitting bundles with excessive token quantities" $ do
 
-        describe "Lists of token bundles" $ do
-
-            it "prop_equipartitionTokenBundlesWithMaxQuantity_length" $
-                property prop_equipartitionTokenBundlesWithMaxQuantity_length
-            it "prop_equipartitionTokenBundlesWithMaxQuantity_sum" $
-                property prop_equipartitionTokenBundlesWithMaxQuantity_sum
+        it "prop_splitBundlesWithExcessiveTokenQuantities_length" $
+            property prop_splitBundlesWithExcessiveTokenQuantities_length
+        it "prop_splitBundlesWithExcessiveTokenQuantities_sum" $
+            property prop_splitBundlesWithExcessiveTokenQuantities_sum
 
     parallel $ describe "Grouping and ungrouping" $ do
 
@@ -1828,12 +1824,12 @@ unit_makeChangeForUserSpecifiedAsset =
     assetC = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "2")
 
 --------------------------------------------------------------------------------
--- Equipartitioning lists of token bundles according to a maximum quantity
+-- Splitting bundles with excessive token quantities
 --------------------------------------------------------------------------------
 
-prop_equipartitionTokenBundlesWithMaxQuantity_length
+prop_splitBundlesWithExcessiveTokenQuantities_length
     :: NonEmpty TokenBundle -> TokenQuantity -> Property
-prop_equipartitionTokenBundlesWithMaxQuantity_length input maxQuantityAllowed =
+prop_splitBundlesWithExcessiveTokenQuantities_length input maxQuantityAllowed =
     maxQuantityAllowed > TokenQuantity.zero ==> checkCoverage $ property $
         cover 5 (lengthOutput > lengthInput)
             "length has increased" $
@@ -1858,13 +1854,13 @@ prop_equipartitionTokenBundlesWithMaxQuantity_length input maxQuantityAllowed =
     maxQuantityOutput =
         F.maximum (TokenMap.maximumQuantity . view #tokens <$> output)
     output =
-        equipartitionTokenBundlesWithMaxQuantity input maxQuantityAllowed
+        splitBundlesWithExcessiveTokenQuantities input maxQuantityAllowed
 
-prop_equipartitionTokenBundlesWithMaxQuantity_sum
+prop_splitBundlesWithExcessiveTokenQuantities_sum
     :: NonEmpty TokenBundle -> TokenQuantity -> Property
-prop_equipartitionTokenBundlesWithMaxQuantity_sum ms maxQuantity =
+prop_splitBundlesWithExcessiveTokenQuantities_sum ms maxQuantity =
     maxQuantity > TokenQuantity.zero ==>
-        F.fold (equipartitionTokenBundlesWithMaxQuantity ms maxQuantity)
+        F.fold (splitBundlesWithExcessiveTokenQuantities ms maxQuantity)
             === F.fold ms
 
 --------------------------------------------------------------------------------

From 6d94ff9e3c41ff0583582368934cbf05184fd7d5 Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Wed, 3 Mar 2021 05:10:59 +0000
Subject: [PATCH 10/24] Adjust `makeChange` to take a single record value.

The argument list for `makeChange` is already very long, and is almost
identical to the fields of the `MakeChange` record data type.

Adjusting `makeChange` to accept a record type has the following advantages:

    1.  It's less cumbersome to add another argument to `makeChange`, because
        the compiler will tell us that a field is missing in all the required
        call sites (both in implementation code and in test code).

    2.  We can remove some duplication and boilerplate from the test code.
---
 .../Primitive/CoinSelection/MA/RoundRobin.hs  | 90 ++++++++++++-------
 .../CoinSelection/MA/RoundRobinSpec.hs        | 60 +++++--------
 2 files changed, 81 insertions(+), 69 deletions(-)

diff --git a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
index bbbd5b690d7..0a8daf91997 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
@@ -47,6 +47,7 @@ module Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     , coinSelectionLens
 
     -- * Making change
+    , MakeChangeCriteria (..)
     , makeChange
     , makeChangeForCoin
     , makeChangeForUserSpecifiedAsset
@@ -407,7 +408,7 @@ performSelection
     -> SelectionCriteria
         -- ^ The selection goal to satisfy.
     -> m (Either SelectionError (SelectionResult TokenBundle))
-performSelection minCoinValueFor costFor criteria
+performSelection minCoinFor costFor criteria
     | not (balanceRequired `leq` balanceAvailable) =
         pure $ Left $ BalanceInsufficient $ BalanceInsufficientError
             { balanceAvailable, balanceRequired }
@@ -460,7 +461,7 @@ performSelection minCoinValueFor costFor criteria
                 Just $ InsufficientMinCoinValueError
                     { expectedMinCoinValue, outputWithInsufficientAda = o }
           where
-            expectedMinCoinValue = minCoinValueFor (view (#tokens . #tokens) o)
+            expectedMinCoinValue = minCoinFor (view (#tokens . #tokens) o)
 
     -- Given a UTxO index that corresponds to a valid selection covering
     -- 'outputsToCover', 'predictChange' yields a non-empty list of assets
@@ -496,12 +497,18 @@ performSelection minCoinValueFor costFor criteria
     predictChange inputsPreSelected = either
         (const $ invariantResultWithNoCost inputsPreSelected)
         (fmap (TokenMap.getAssets . view #tokens))
-        (makeChange noMinimumCoin noCost
-            extraCoinSource
-            (view #tokens . snd <$> mkInputsSelected inputsPreSelected)
-            (view #tokens <$> outputsToCover)
+        (makeChange MakeChangeCriteria
+            { minCoinFor = noMinimumCoin
+            , requiredCost = noCost
+            , extraCoinSource
+            , inputBundles
+            , outputBundles
+            }
         )
       where
+        inputBundles = view #tokens . snd <$> mkInputsSelected inputsPreSelected
+        outputBundles = view #tokens <$> outputsToCover
+
         noMinimumCoin :: TokenMap -> Coin
         noMinimumCoin = const (Coin 0)
 
@@ -553,10 +560,13 @@ performSelection minCoinValueFor costFor criteria
                     pure $ Left $ UnableToConstructChange changeErr
       where
         mChangeGenerated :: Either UnableToConstructChangeError [TokenBundle]
-        mChangeGenerated = makeChange minCoinValueFor cost
-            extraCoinSource
-            (view #tokens . snd <$> inputsSelected)
-            (view #tokens <$> outputsToCover)
+        mChangeGenerated = makeChange MakeChangeCriteria
+            { minCoinFor
+            , requiredCost
+            , extraCoinSource
+            , inputBundles =  view #tokens . snd <$> inputsSelected
+            , outputBundles = view #tokens <$> outputsToCover
+            }
 
         mkSelectionResult :: [TokenBundle] -> SelectionResult TokenBundle
         mkSelectionResult changeGenerated = SelectionResult
@@ -572,7 +582,7 @@ performSelection minCoinValueFor costFor criteria
 
         SelectionState {selected, leftover} = s
 
-        cost = costFor SelectionSkeleton
+        requiredCost = costFor SelectionSkeleton
             { inputsSkeleton  = selected
             , outputsSkeleton = NE.toList outputsToCover
             , changeSkeleton
@@ -788,6 +798,30 @@ runSelectionStep lens s
 -- Making change
 --------------------------------------------------------------------------------
 
+-- | Criteria for the 'makeChange' function.
+--
+data MakeChangeCriteria minCoinFor = MakeChangeCriteria
+    { minCoinFor :: minCoinFor
+      -- ^ A function that computes the minimum required ada quantity for a
+      -- particular output.
+    , requiredCost :: Coin
+      -- ^ The minimal (and optimal) delta between the total ada balance
+      -- of all input bundles and the total ada balance of all output and
+      -- change bundles, where:
+      --
+      --    delta = getCoin (fold inputBundles)
+      --          - getCoin (fold outputBundles)
+      --          - getCoin (fold changeBundles)
+      --
+      -- This typically captures fees plus key deposits.
+    , extraCoinSource :: Maybe Coin
+        -- ^ An optional extra source of ada.
+    , inputBundles :: NonEmpty TokenBundle
+        -- ^ Token bundles of selected inputs.
+    , outputBundles :: NonEmpty TokenBundle
+        -- ^ Token bundles of original outputs.
+    } deriving (Eq, Generic, Show)
+
 -- | Constructs change bundles for a set of selected inputs and outputs.
 --
 -- Returns 'Nothing' if the specified inputs do not provide enough ada to
@@ -808,29 +842,11 @@ runSelectionStep lens s
 -- to every output token bundle.
 --
 makeChange
-    :: (TokenMap -> Coin)
-        -- A function that computes the minimum required ada quantity for a
-        -- particular output.
-    -> Coin
-        -- ^ The minimal (and optimal) delta between the total ada balance
-        -- of all input bundles and the total ada balance of all output and
-        -- change bundles, where:
-        --
-        --    delta = getCoin (fold inputBundles)
-        --          - getCoin (fold outputBundles)
-        --          - getCoin (fold changeBundles)
-        --
-        -- This typically captures fees plus key deposits.
-        --
-    -> Maybe Coin
-        -- ^ An optional extra source of ada.
-    -> NonEmpty TokenBundle
-        -- ^ Token bundles of selected inputs.
-    -> NonEmpty TokenBundle
-        -- ^ Token bundles of original outputs.
+    :: MakeChangeCriteria (TokenMap -> Coin)
+        -- ^ Criteria for making change.
     -> Either UnableToConstructChangeError [TokenBundle]
         -- ^ Generated change bundles.
-makeChange minCoinFor requiredCost mExtraCoinSource inputBundles outputBundles
+makeChange criteria
     | not (totalOutputValue `leq` totalInputValue) =
         totalInputValueInsufficient
     | TokenBundle.getCoin totalOutputValue == Coin 0 =
@@ -841,6 +857,14 @@ makeChange minCoinFor requiredCost mExtraCoinSource inputBundles outputBundles
             assignCoinsToChangeMaps
                 adaAvailable minCoinFor changeMapOutputCoinPairs
   where
+    MakeChangeCriteria
+        { minCoinFor
+        , requiredCost
+        , extraCoinSource
+        , inputBundles
+        , outputBundles
+        } = criteria
+
     -- The following subtraction is safe, as we have already checked
     -- that the total input value is greater than the total output
     -- value:
@@ -940,7 +964,7 @@ makeChange minCoinFor requiredCost mExtraCoinSource inputBundles outputBundles
     totalInputValue :: TokenBundle
     totalInputValue = TokenBundle.add
         (F.fold inputBundles)
-        (maybe TokenBundle.empty TokenBundle.fromCoin mExtraCoinSource)
+        (maybe TokenBundle.empty TokenBundle.fromCoin extraCoinSource)
 
     totalOutputValue :: TokenBundle
     totalOutputValue = F.fold outputBundles
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 cc4f046bcf0..4860f89e73c 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
@@ -24,6 +24,7 @@ import Cardano.Numeric.Util
 import Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     ( AssetCount (..)
     , BalanceInsufficientError (..)
+    , MakeChangeCriteria (..)
     , InsufficientMinCoinValueError (..)
     , SelectionCriteria (..)
     , SelectionError (..)
@@ -1308,18 +1309,7 @@ linearCost SelectionSkeleton{inputsSkeleton, outputsSkeleton, changeSkeleton}
     + F.length outputsSkeleton
     + F.length changeSkeleton
 
-data MakeChangeData = MakeChangeData
-    { inputBundles
-        :: NonEmpty TokenBundle
-    , extraInputCoin
-        :: Maybe Coin
-    , outputBundles
-        :: NonEmpty TokenBundle
-    , cost
-        :: Coin
-    , minCoinValueDef
-        :: MinCoinValueFor
-    } deriving (Eq, Show)
+type MakeChangeData = MakeChangeCriteria MinCoinValueFor
 
 isValidMakeChangeData :: MakeChangeData -> Bool
 isValidMakeChangeData p = (&&)
@@ -1328,7 +1318,7 @@ isValidMakeChangeData p = (&&)
   where
     totalInputValue = TokenBundle.add
         (F.fold $ inputBundles p)
-        (maybe TokenBundle.empty TokenBundle.fromCoin (extraInputCoin p))
+        (maybe TokenBundle.empty TokenBundle.fromCoin (view #extraCoinSource p))
     totalOutputValue = F.fold $ outputBundles p
     totalOutputCoinValue = TokenBundle.getCoin totalOutputValue
 
@@ -1336,12 +1326,12 @@ genMakeChangeData :: Gen MakeChangeData
 genMakeChangeData = flip suchThat isValidMakeChangeData $ do
     outputBundleCount <- choose (0, 15)
     let inputBundleCount = outputBundleCount * 4
-    MakeChangeData
-        <$> genTokenBundles inputBundleCount
+    MakeChangeCriteria
+        <$> arbitrary
+        <*> genCoinSmall
         <*> oneof [pure Nothing, Just <$> genCoinSmallPositive]
+        <*> genTokenBundles inputBundleCount
         <*> genTokenBundles outputBundleCount
-        <*> genCoinSmall
-        <*> arbitrary
   where
     genTokenBundles :: Int -> Gen (NonEmpty TokenBundle)
     genTokenBundles count = (:|)
@@ -1351,16 +1341,14 @@ genMakeChangeData = flip suchThat isValidMakeChangeData $ do
 makeChangeWith
     :: MakeChangeData
     -> Either UnableToConstructChangeError [TokenBundle]
-makeChangeWith p = makeChange
-    (mkMinCoinValueFor $ minCoinValueDef p)
-    (cost p)
-    (extraInputCoin p) (inputBundles p)
-    (outputBundles p)
+makeChangeWith p = makeChange p
+    { minCoinFor = mkMinCoinValueFor $ minCoinFor p}
 
 prop_makeChange_identity
     :: NonEmpty TokenBundle -> Property
 prop_makeChange_identity bundles = (===)
-    (F.fold <$> makeChange (const (Coin 0)) (Coin 0) Nothing bundles bundles)
+    (F.fold <$> makeChange
+        (MakeChangeCriteria (const (Coin 0)) (Coin 0) Nothing bundles bundles))
     (Right TokenBundle.empty)
 
 prop_makeChange_length
@@ -1371,8 +1359,7 @@ prop_makeChange_length p =
         Left{} -> property False
         Right xs -> length xs === length (outputBundles p)
   where
-    change = makeChange noMinCoin noCost
-        (extraInputCoin p) (inputBundles p) (outputBundles p)
+    change = makeChange p {minCoinFor = noMinCoin, requiredCost = noCost}
 
 prop_makeChange
     :: MakeChangeData
@@ -1408,7 +1395,7 @@ prop_makeChange_success_delta p change =
             totalInputValue
             totalOutputWithChange
     in
-        (delta === TokenBundle.fromCoin (cost p))
+        (delta === TokenBundle.fromCoin (view #requiredCost p))
             & counterexample counterExampleText
   where
     counterExampleText = unlines
@@ -1421,7 +1408,7 @@ prop_makeChange_success_delta p change =
         ]
     totalInputValue = TokenBundle.add
         (F.fold (inputBundles p))
-        (maybe TokenBundle.empty TokenBundle.fromCoin (extraInputCoin p))
+        (maybe TokenBundle.empty TokenBundle.fromCoin (view #extraCoinSource p))
     totalInputCoin =
         TokenBundle.getCoin totalInputValue
     totalOutputValue =
@@ -1443,7 +1430,7 @@ prop_makeChange_success_minValueRespected p =
     F.foldr ((.&&.) . checkMinValue) (property True)
   where
     minCoinValueFor :: TokenMap -> Coin
-    minCoinValueFor = mkMinCoinValueFor (minCoinValueDef p)
+    minCoinValueFor = mkMinCoinValueFor (minCoinFor p)
 
     checkMinValue :: TokenBundle -> Property
     checkMinValue m@TokenBundle{coin,tokens} =
@@ -1472,12 +1459,12 @@ prop_makeChange_fail_costTooBig p =
             totalInputValue
             totalOutputValue
     in
-        deltaCoin < cost p
+        deltaCoin < view #requiredCost p
             & counterexample ("delta: " <> pretty deltaCoin)
   where
     totalInputValue = TokenBundle.add
         (F.fold (inputBundles p))
-        (maybe TokenBundle.empty TokenBundle.fromCoin (extraInputCoin p))
+        (maybe TokenBundle.empty TokenBundle.fromCoin (view #extraCoinSource p))
     totalOutputValue =
         F.fold $ outputBundles p
 
@@ -1490,7 +1477,7 @@ prop_makeChange_fail_minValueTooBig
     :: MakeChangeData
     -> Property
 prop_makeChange_fail_minValueTooBig p =
-    case makeChangeWith (p { cost = noCost, minCoinValueDef = NoMinCoin }) of
+    case makeChangeWith p {requiredCost = noCost, minCoinFor = NoMinCoin} of
         Left{} ->
             property False & counterexample "makeChange failed with no cost!"
         -- If 'makeChange' failed to generate change, we try to re-run it with
@@ -1501,8 +1488,8 @@ prop_makeChange_fail_minValueTooBig p =
         -- coins available to generate all change outputs.
         Right change ->
             conjoin
-                [ deltaCoin < (totalMinCoinDeposit `addCoin` cost p)
-                , deltaCoin >= cost p
+                [ deltaCoin < totalMinCoinDeposit `addCoin` view #requiredCost p
+                , deltaCoin >= view #requiredCost p
                 ]
                 & counterexample counterexampleText
           where
@@ -1518,20 +1505,21 @@ prop_makeChange_fail_minValueTooBig p =
                 totalInputValue
                 totalOutputValue
             minCoinValueFor =
-                mkMinCoinValueFor (minCoinValueDef p)
+                mkMinCoinValueFor (minCoinFor p)
             totalMinCoinDeposit = F.foldr addCoin (Coin 0)
                 (minCoinValueFor . view #tokens <$> change)
   where
     totalInputValue = TokenBundle.add
         (F.fold (inputBundles p))
-        (maybe TokenBundle.empty TokenBundle.fromCoin (extraInputCoin p))
+        (maybe TokenBundle.empty TokenBundle.fromCoin (view #extraCoinSource p))
     totalOutputValue =
         F.fold $ outputBundles p
 
 unit_makeChange
     :: [Expectation]
 unit_makeChange =
-    [ makeChange minCoinValueFor cost extraSource i o `shouldBe` expectation
+    [ makeChange (MakeChangeCriteria minCoinValueFor cost extraSource i o)
+        `shouldBe` expectation
     | (minCoinValueFor, cost, extraSource, i, o, expectation) <- matrix
     ]
   where

From 82f08f83e477e9e5712792763e01d6b92d8b698e Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Tue, 2 Mar 2021 09:32:26 +0000
Subject: [PATCH 11/24] Test `equipartitionQuantitiesWithUpperBound` only with
 positive upper bounds.

---
 .../Wallet/Primitive/Types/TokenMapSpec.hs    | 88 ++++++++++---------
 1 file changed, 47 insertions(+), 41 deletions(-)

diff --git a/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenMapSpec.hs b/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenMapSpec.hs
index 91ba2596094..5aedf1ae049 100644
--- a/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenMapSpec.hs
+++ b/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenMapSpec.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE NoMonomorphismRestriction #-}
 {-# LANGUAGE QuasiQuotes #-}
 {-# LANGUAGE TemplateHaskell #-}
@@ -35,7 +36,11 @@ import Cardano.Wallet.Primitive.Types.TokenPolicy.Gen
 import Cardano.Wallet.Primitive.Types.TokenQuantity
     ( TokenQuantity (..) )
 import Cardano.Wallet.Primitive.Types.TokenQuantity.Gen
-    ( genTokenQuantitySmall, shrinkTokenQuantitySmall )
+    ( genTokenQuantitySmall
+    , genTokenQuantitySmallPositive
+    , shrinkTokenQuantitySmall
+    , shrinkTokenQuantitySmallPositive
+    )
 import Data.Aeson
     ( FromJSON (..), ToJSON (..) )
 import Data.Aeson.QQ
@@ -592,60 +597,53 @@ equipartitionQuantitiesWithUpperBound_expectedLength
     TokenQuantity currentMaxQuantity = TokenMap.maximumQuantity m
 
 prop_equipartitionQuantitiesWithUpperBound_coverage
-    :: TokenMap -> TokenQuantity -> Property
-prop_equipartitionQuantitiesWithUpperBound_coverage m maxQuantity =
-    maxQuantity > TokenQuantity.zero ==>
-        checkCoverage $
-        cover 8 (maxQuantity == TokenQuantity 1)
-            "Maximum allowable quantity == 1" $
-        cover 8 (maxQuantity == TokenQuantity 2)
-            "Maximum allowable quantity == 2" $
-        cover 8 (maxQuantity >= TokenQuantity 3)
-            "Maximum allowable quantity >= 3" $
-        cover 8 (expectedLength == 1)
-            "Expected number of parts == 1" $
-        cover 8 (expectedLength == 2)
-            "Expected number of parts == 2" $
-        cover 8 (expectedLength >= 3)
-            "Expected number of parts >= 3" $
-        property $ expectedLength > 0
+    :: TokenMap -> Positive TokenQuantity -> Property
+prop_equipartitionQuantitiesWithUpperBound_coverage m (Positive maxQuantity) =
+    checkCoverage $
+    cover 8 (maxQuantity == TokenQuantity 1)
+        "Maximum allowable quantity == 1" $
+    cover 8 (maxQuantity == TokenQuantity 2)
+        "Maximum allowable quantity == 2" $
+    cover 8 (maxQuantity >= TokenQuantity 3)
+        "Maximum allowable quantity >= 3" $
+    cover 8 (expectedLength == 1)
+        "Expected number of parts == 1" $
+    cover 8 (expectedLength == 2)
+        "Expected number of parts == 2" $
+    cover 8 (expectedLength >= 3)
+        "Expected number of parts >= 3" $
+    property $ expectedLength > 0
   where
     expectedLength = equipartitionQuantitiesWithUpperBound_expectedLength
         m maxQuantity
 
 prop_equipartitionQuantitiesWithUpperBound_length
-    :: TokenMap -> TokenQuantity -> Property
-prop_equipartitionQuantitiesWithUpperBound_length m maxQuantity =
-    maxQuantity > TokenQuantity.zero ==>
-        length (TokenMap.equipartitionQuantitiesWithUpperBound m maxQuantity)
-            === equipartitionQuantitiesWithUpperBound_expectedLength
-                m maxQuantity
+    :: TokenMap -> Positive TokenQuantity -> Property
+prop_equipartitionQuantitiesWithUpperBound_length m (Positive maxQuantity) =
+    length (TokenMap.equipartitionQuantitiesWithUpperBound m maxQuantity)
+        === equipartitionQuantitiesWithUpperBound_expectedLength m maxQuantity
 
 prop_equipartitionQuantitiesWithUpperBound_max
-    :: TokenMap -> TokenQuantity -> Property
-prop_equipartitionQuantitiesWithUpperBound_max m maxQuantity =
-    maxQuantity > TokenQuantity.zero ==>
-        checkCoverage $
-        cover 10 (maxResultQuantity == maxQuantity)
-            "At least one resultant token map has a maximal quantity" $
-        property $ maxResultQuantity <= maxQuantity
+    :: TokenMap -> Positive TokenQuantity -> Property
+prop_equipartitionQuantitiesWithUpperBound_max m (Positive maxQuantity) =
+    checkCoverage $
+    cover 10 (maxResultQuantity == maxQuantity)
+        "At least one resultant token map has a maximal quantity" $
+    property $ maxResultQuantity <= maxQuantity
   where
     results = TokenMap.equipartitionQuantitiesWithUpperBound m maxQuantity
     maxResultQuantity = F.maximum (TokenMap.maximumQuantity <$> results)
 
 prop_equipartitionQuantitiesWithUpperBound_order
-    :: TokenMap -> TokenQuantity -> Property
-prop_equipartitionQuantitiesWithUpperBound_order m maxQuantity =
-    maxQuantity > TokenQuantity.zero ==>
-        inAscendingPartialOrder
-            (TokenMap.equipartitionQuantitiesWithUpperBound m maxQuantity)
+    :: TokenMap -> Positive TokenQuantity -> Property
+prop_equipartitionQuantitiesWithUpperBound_order m (Positive maxQuantity) =
+    property $ inAscendingPartialOrder
+        (TokenMap.equipartitionQuantitiesWithUpperBound m maxQuantity)
 
 prop_equipartitionQuantitiesWithUpperBound_sum
-    :: TokenMap -> TokenQuantity -> Property
-prop_equipartitionQuantitiesWithUpperBound_sum m maxQuantity =
-    maxQuantity > TokenQuantity.zero ==>
-        F.fold (TokenMap.equipartitionQuantitiesWithUpperBound m maxQuantity)
-            === m
+    :: TokenMap -> Positive TokenQuantity -> Property
+prop_equipartitionQuantitiesWithUpperBound_sum m (Positive maxQuantity) =
+    F.fold (TokenMap.equipartitionQuantitiesWithUpperBound m maxQuantity) === m
 
 --------------------------------------------------------------------------------
 -- JSON serialization tests
@@ -816,6 +814,10 @@ tokenPolicyIdHexStringLength = 56
 -- Arbitrary instances
 --------------------------------------------------------------------------------
 
+newtype Positive a = Positive
+    { getPositive :: a }
+    deriving (Eq, Show)
+
 instance Arbitrary a => Arbitrary (Flat a) where
     arbitrary = Flat <$> arbitrary
     shrink = fmap Flat . shrink . getFlat
@@ -855,3 +857,7 @@ instance Arbitrary TokenQuantity where
     -- zero-valued tokens) is maintained.
     arbitrary = genTokenQuantitySmall
     shrink = shrinkTokenQuantitySmall
+
+instance Arbitrary (Positive TokenQuantity) where
+    arbitrary = Positive <$> genTokenQuantitySmallPositive
+    shrink = fmap Positive . shrinkTokenQuantitySmallPositive . getPositive

From 5a03b820f9726a5fa34759806edd488b616f0ba0 Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Tue, 2 Mar 2021 06:16:31 +0000
Subject: [PATCH 12/24] Add function `TokenMap.equipartitionAssets`.

---
 .../Wallet/Primitive/Types/TokenMap.hs        | 40 ++++++++++
 .../Wallet/Primitive/Types/TokenMapSpec.hs    | 77 +++++++++++++++++++
 2 files changed, 117 insertions(+)

diff --git a/lib/core/src/Cardano/Wallet/Primitive/Types/TokenMap.hs b/lib/core/src/Cardano/Wallet/Primitive/Types/TokenMap.hs
index b20c6786cb3..1a4761261bc 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/Types/TokenMap.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/Types/TokenMap.hs
@@ -72,6 +72,7 @@ module Cardano.Wallet.Primitive.Types.TokenMap
     , maximumQuantity
 
     -- * Partitioning
+    , equipartitionAssets
     , equipartitionQuantities
     , equipartitionQuantitiesWithUpperBound
 
@@ -95,6 +96,8 @@ import Prelude hiding
 
 import Algebra.PartialOrd
     ( PartialOrd (..) )
+import Cardano.Numeric.Util
+    ( equipartitionNatural )
 import Cardano.Wallet.Primitive.Types.TokenPolicy
     ( TokenName, TokenPolicyId )
 import Cardano.Wallet.Primitive.Types.TokenQuantity
@@ -133,6 +136,8 @@ import GHC.Generics
     ( Generic )
 import GHC.TypeLits
     ( ErrorMessage (..), TypeError )
+import Numeric.Natural
+    ( Natural )
 import Quiet
     ( Quiet (..) )
 
@@ -650,6 +655,41 @@ maximumQuantity =
 -- Partitioning
 --------------------------------------------------------------------------------
 
+-- | Partitions a token map into 'n' smaller maps, where the asset sets of the
+--   resultant maps are disjoint.
+--
+-- In the resultant maps, the smallest asset set size and largest asset set
+-- size will differ by no more than 1.
+--
+-- The quantities of each asset are unchanged.
+--
+equipartitionAssets
+    :: TokenMap
+    -- ^ The token map to be partitioned.
+    -> NonEmpty a
+    -- ^ Represents the number of portions in which to partition the token map.
+    -> NonEmpty TokenMap
+    -- ^ The partitioned maps.
+equipartitionAssets m mapCount =
+    fromFlatList <$> NE.unfoldr generateChunk (assetCounts, toFlatList m)
+  where
+    -- The total number of assets.
+    assetCount :: Int
+    assetCount = Set.size $ getAssets m
+
+    -- How many asset quantities to include in each chunk.
+    assetCounts :: NonEmpty Int
+    assetCounts = fromIntegral @Natural @Int <$>
+        equipartitionNatural (fromIntegral @Int @Natural assetCount) mapCount
+
+    -- Generates a single chunk of asset quantities.
+    generateChunk :: (NonEmpty Int, [aq]) -> ([aq], Maybe (NonEmpty Int, [aq]))
+    generateChunk (c :| mcs, aqs) = case NE.nonEmpty mcs of
+        Just cs -> (prefix, Just (cs, suffix))
+        Nothing -> (aqs, Nothing)
+      where
+        (prefix, suffix) = L.splitAt c aqs
+
 -- | Partitions a token map into 'n' smaller maps, where the quantity of each
 --   token is equipartitioned across the resultant maps.
 --
diff --git a/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenMapSpec.hs b/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenMapSpec.hs
index 5aedf1ae049..768b3f6c2b8 100644
--- a/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenMapSpec.hs
+++ b/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenMapSpec.hs
@@ -20,6 +20,7 @@ import Cardano.Wallet.Primitive.Types.TokenMap
     ( AssetId (..), Flat (..), Nested (..), TokenMap, difference )
 import Cardano.Wallet.Primitive.Types.TokenMap.Gen
     ( AssetIdF (..)
+    , genAssetIdLargeRange
     , genAssetIdSmallRange
     , genTokenMapSmallRange
     , shrinkAssetIdSmallRange
@@ -41,6 +42,8 @@ import Cardano.Wallet.Primitive.Types.TokenQuantity.Gen
     , shrinkTokenQuantitySmall
     , shrinkTokenQuantitySmallPositive
     )
+import Control.Monad
+    ( replicateM )
 import Data.Aeson
     ( FromJSON (..), ToJSON (..) )
 import Data.Aeson.QQ
@@ -83,13 +86,17 @@ import Test.Hspec.Extra
     ( parallel )
 import Test.QuickCheck
     ( Arbitrary (..)
+    , Blind (..)
     , Fun
     , Property
     , applyFun
     , checkCoverage
+    , choose
     , counterexample
     , cover
+    , frequency
     , property
+    , (.||.)
     , (===)
     , (==>)
     )
@@ -217,6 +224,17 @@ spec =
         it "prop_maximumQuantity_all" $
             property prop_maximumQuantity_all
 
+    parallel $ describe "Partitioning assets" $ do
+
+        it "prop_equipartitionAssets_coverage" $
+            property prop_equipartitionAssets_coverage
+        it "prop_equipartitionAssets_length" $
+            property prop_equipartitionAssets_length
+        it "prop_equipartitionAssets_sizes" $
+            property prop_equipartitionAssets_sizes
+        it "prop_equipartitionAssets_sum" $
+            property prop_equipartitionAssets_sum
+
     parallel $ describe "Partitioning quantities" $ do
 
         it "prop_equipartitionQuantities_fair" $
@@ -532,6 +550,47 @@ prop_maximumQuantity_all b =
   where
     maxQ = TokenMap.maximumQuantity b
 
+--------------------------------------------------------------------------------
+-- Partitioning assets
+--------------------------------------------------------------------------------
+
+prop_equipartitionAssets_coverage
+    :: Blind (Large TokenMap) -> Property
+prop_equipartitionAssets_coverage m = checkCoverage $
+    cover 4 (assetCount == 0)
+        "asset count = 0" $
+    cover 4 (assetCount == 1)
+        "asset count = 1" $
+    cover 20 (2 <= assetCount && assetCount <= 31)
+        "2 <= asset count <= 31" $
+    cover 20 (32 <= assetCount && assetCount <= 63)
+        "32 <= asset count <= 63" $
+    True
+  where
+    assetCount = Set.size $ TokenMap.getAssets $ getLarge $ getBlind m
+
+prop_equipartitionAssets_length
+    :: Blind (Large TokenMap) -> NonEmpty () -> Property
+prop_equipartitionAssets_length (Blind (Large m)) count =
+    NE.length (TokenMap.equipartitionAssets m count) === NE.length count
+
+prop_equipartitionAssets_sizes
+    :: Blind (Large TokenMap) -> NonEmpty () -> Property
+prop_equipartitionAssets_sizes (Blind (Large m)) count = (.||.)
+    (assetCountDifference == 0)
+    (assetCountDifference == 1)
+  where
+    assetCounts = Set.size . TokenMap.getAssets <$> results
+    assetCountMin = F.minimum assetCounts
+    assetCountMax = F.maximum assetCounts
+    assetCountDifference = assetCountMax - assetCountMin
+    results = TokenMap.equipartitionAssets m count
+
+prop_equipartitionAssets_sum
+    :: Blind (Large TokenMap) -> NonEmpty () -> Property
+prop_equipartitionAssets_sum (Blind (Large m)) count =
+    F.fold (TokenMap.equipartitionAssets m count) === m
+
 --------------------------------------------------------------------------------
 -- Partitioning quantities
 --------------------------------------------------------------------------------
@@ -814,6 +873,10 @@ tokenPolicyIdHexStringLength = 56
 -- Arbitrary instances
 --------------------------------------------------------------------------------
 
+newtype Large a = Large
+    { getLarge :: a }
+    deriving (Eq, Show)
+
 newtype Positive a = Positive
     { getPositive :: a }
     deriving (Eq, Show)
@@ -838,6 +901,20 @@ instance Arbitrary TokenMap where
     arbitrary = genTokenMapSmallRange
     shrink = shrinkTokenMapSmallRange
 
+instance Arbitrary (Large TokenMap) where
+    arbitrary = Large <$> do
+        assetCount <- frequency
+            [ (1, pure 0)
+            , (1, pure 1)
+            , (8, choose (2, 63))
+            ]
+        TokenMap.fromFlatList <$> replicateM assetCount genAssetQuantity
+      where
+        genAssetQuantity = (,)
+            <$> genAssetIdLargeRange
+            <*> genTokenQuantitySmallPositive
+    -- No shrinking
+
 instance Arbitrary TokenName where
     arbitrary = genTokenNameSmallRange
     shrink = shrinkTokenNameSmallRange

From da8e996a07662e9f707c4c2ee109d16562e4ee28 Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Tue, 2 Mar 2021 06:23:46 +0000
Subject: [PATCH 13/24] Add function `TokenBundle.equipartitionAssets`.

---
 .../Wallet/Primitive/Types/TokenBundle.hs     | 24 +++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/lib/core/src/Cardano/Wallet/Primitive/Types/TokenBundle.hs b/lib/core/src/Cardano/Wallet/Primitive/Types/TokenBundle.hs
index ab38a4e4144..0c828d1af97 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/Types/TokenBundle.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/Types/TokenBundle.hs
@@ -54,6 +54,7 @@ module Cardano.Wallet.Primitive.Types.TokenBundle
     , removeQuantity
 
     -- * Partitioning
+    , equipartitionAssets
     , equipartitionQuantitiesWithUpperBound
 
     -- * Policies
@@ -373,6 +374,29 @@ removeQuantity b a = b { tokens = TokenMap.removeQuantity (tokens b) a }
 -- Partitioning
 --------------------------------------------------------------------------------
 
+-- | Partitions a token bundle into 'n' smaller bundles, where the asset sets
+--   of the resultant bundles are disjoint.
+--
+-- In the resultant bundles, the smallest asset set size and largest asset set
+-- size will differ by no more than 1.
+--
+-- The ada 'Coin' quantity is equipartitioned across the resulting bundles.
+--
+-- The quantities of each non-ada asset are unchanged.
+--
+equipartitionAssets
+    :: TokenBundle
+    -- ^ The token bundle to be partitioned.
+    -> NonEmpty a
+    -- ^ Represents the number of portions in which to partition the bundle.
+    -> NonEmpty TokenBundle
+    -- ^ The partitioned bundles.
+equipartitionAssets (TokenBundle c m) count =
+    NE.zipWith TokenBundle cs ms
+  where
+    cs = Coin.equipartition c count
+    ms = TokenMap.equipartitionAssets m count
+
 -- | Partitions a token bundle into 'n' smaller bundles, where the quantity of
 --   each token is equipartitioned across the resultant bundles, with the goal
 --   that no token quantity in any of the resultant bundles exceeds the given

From eae3da161174ec6ff083f4b4de3895b7bc826151 Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Tue, 2 Mar 2021 09:52:28 +0000
Subject: [PATCH 14/24] Add function `splitBundleIfAssetCountExcessive`.

---
 .../Primitive/CoinSelection/MA/RoundRobin.hs  | 21 ++++
 .../CoinSelection/MA/RoundRobinSpec.hs        | 98 ++++++++++++++++++-
 2 files changed, 117 insertions(+), 2 deletions(-)

diff --git a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
index 0a8daf91997..3e8606df200 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
@@ -55,6 +55,7 @@ module Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     , assignCoinsToChangeMaps
 
     -- * Splitting bundles
+    , splitBundleIfAssetCountExcessive
     , splitBundlesWithExcessiveTokenQuantities
 
     -- * Grouping and ungrouping
@@ -1203,6 +1204,26 @@ makeChangeForCoin targets excess =
 -- Splitting bundles
 --------------------------------------------------------------------------------
 
+-- | Splits a bundle into smaller bundles if its asset count is excessive when
+--   measured with the given 'isExcessive' indicator function.
+--
+-- Returns a list of smaller bundles for which 'isExcessive' returns 'False'.
+--
+splitBundleIfAssetCountExcessive
+    :: TokenBundle
+    -- ^ The token bundle suspected to have an excessive number of assets.
+    -> (TokenBundle -> Bool)
+    -- ^ A function that returns 'True' if (and only if) the asset count of
+    -- the given bundle is excessive.
+    -> NonEmpty TokenBundle
+splitBundleIfAssetCountExcessive b isExcessive
+    | isExcessive b =
+        splitInHalf b >>= flip splitBundleIfAssetCountExcessive isExcessive
+    | otherwise =
+        pure b
+  where
+    splitInHalf = flip TokenBundle.equipartitionAssets (() :| [()])
+
 -- | Splits bundles with excessive token quantities into smaller bundles.
 --
 -- Only token bundles containing quantities that exceed the maximum token
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 4860f89e73c..e80235e4745 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
@@ -52,6 +52,7 @@ import Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     , runRoundRobin
     , runSelection
     , runSelectionStep
+    , splitBundleIfAssetCountExcessive
     , splitBundlesWithExcessiveTokenQuantities
     , ungroupByKey
     )
@@ -60,7 +61,11 @@ import Cardano.Wallet.Primitive.Types.Address
 import Cardano.Wallet.Primitive.Types.Coin
     ( Coin (..), addCoin )
 import Cardano.Wallet.Primitive.Types.Coin.Gen
-    ( genCoinSmall, genCoinSmallPositive, shrinkCoinSmallPositive )
+    ( genCoinLargePositive
+    , genCoinSmall
+    , genCoinSmallPositive
+    , shrinkCoinSmallPositive
+    )
 import Cardano.Wallet.Primitive.Types.Hash
     ( Hash (..) )
 import Cardano.Wallet.Primitive.Types.TokenBundle
@@ -70,7 +75,8 @@ import Cardano.Wallet.Primitive.Types.TokenBundle.Gen
 import Cardano.Wallet.Primitive.Types.TokenMap
     ( AssetId (..), TokenMap )
 import Cardano.Wallet.Primitive.Types.TokenMap.Gen
-    ( genAssetIdSmallRange
+    ( genAssetIdLargeRange
+    , genAssetIdSmallRange
     , genTokenMapSmallRange
     , shrinkAssetIdSmallRange
     , shrinkTokenMapSmallRange
@@ -317,6 +323,17 @@ spec = describe "Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobinSpec" $
         unitTests "makeChangeForUserSpecifiedAsset"
             unit_makeChangeForUserSpecifiedAsset
 
+    parallel $ describe "Splitting bundles with excessive asset counts" $ do
+
+        it "prop_splitBundleIfAssetCountExcessive_length" $
+            property prop_splitBundleIfAssetCountExcessive_length
+        it "prop_splitBundleIfAssetCountExcessive_maximalSplitting" $
+            property prop_splitBundleIfAssetCountExcessive_maximalSplitting
+        it "prop_splitBundleIfAssetCountExcessive_postCondition" $
+            property prop_splitBundleIfAssetCountExcessive_postCondition
+        it "prop_splitBundleIfAssetCountExcessive_sum" $
+            property prop_splitBundleIfAssetCountExcessive_sum
+
     parallel $ describe "Splitting bundles with excessive token quantities" $ do
 
         it "prop_splitBundlesWithExcessiveTokenQuantities_length" $
@@ -1811,6 +1828,64 @@ unit_makeChangeForUserSpecifiedAsset =
     assetC :: AssetId
     assetC = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "2")
 
+--------------------------------------------------------------------------------
+-- Splitting bundles with excessive asset counts
+--------------------------------------------------------------------------------
+
+prop_splitBundleIfAssetCountExcessive_length
+    :: Blind (Large TokenBundle) -> Positive Int -> Property
+prop_splitBundleIfAssetCountExcessive_length
+    (Blind (Large b)) (Positive maxAssetCount) =
+        checkCoverage $ property $
+        cover 5 (resultLength == 1)
+            "length = 1" $
+        cover 5 (resultLength >= 2 && resultLength < 8)
+            "length >= 2 && length < 8" $
+        cover 5 (resultLength >= 8 && resultLength < 16)
+            "length >= 8 && length < 16" $
+        True
+  where
+    isExcessive = (> maxAssetCount) . Set.size . TokenBundle.getAssets
+    result = splitBundleIfAssetCountExcessive b isExcessive
+    resultLength = NE.length result
+
+prop_splitBundleIfAssetCountExcessive_maximalSplitting
+    :: Blind (Large TokenBundle) -> Property
+prop_splitBundleIfAssetCountExcessive_maximalSplitting (Blind (Large b)) =
+    checkCoverage $ property $
+    cover 5 (assetCount == 0)
+        "asset count = 0" $
+    cover 5 (assetCount == 1)
+        "asset count = 1" $
+    cover 5 (assetCount >= 2 && assetCount < 8)
+        "asset count >= 2 && asset count < 8" $
+    cover 5 (assetCount >= 8 && assetCount < 16)
+        "asset count >= 8 && asset count < 16" $
+    (.&&.)
+        (NE.length result === max 1 assetCount)
+        (F.all ((<= 1) . Set.size . TokenBundle.getAssets) result)
+  where
+    assetCount = Set.size $ TokenBundle.getAssets b
+    isExcessive = (> 1) . Set.size . TokenBundle.getAssets
+    result = splitBundleIfAssetCountExcessive b isExcessive
+
+prop_splitBundleIfAssetCountExcessive_postCondition
+    :: Blind (Large TokenBundle) -> Positive Int -> Property
+prop_splitBundleIfAssetCountExcessive_postCondition
+    (Blind (Large b)) (Positive maxAssetCount) =
+        property $ F.all (not . isExcessive) results
+  where
+    isExcessive = (> maxAssetCount) . Set.size . TokenBundle.getAssets
+    results = splitBundleIfAssetCountExcessive b isExcessive
+
+prop_splitBundleIfAssetCountExcessive_sum
+    :: Blind (Large TokenBundle) -> Positive Int -> Property
+prop_splitBundleIfAssetCountExcessive_sum
+    (Blind (Large b)) (Positive maxAssetCount) =
+        F.fold (splitBundleIfAssetCountExcessive b isExcessive) === b
+  where
+    isExcessive = (> maxAssetCount) . Set.size . TokenBundle.getAssets
+
 --------------------------------------------------------------------------------
 -- Splitting bundles with excessive token quantities
 --------------------------------------------------------------------------------
@@ -2075,6 +2150,25 @@ instance Arbitrary TokenBundle where
     arbitrary = genTokenBundleSmallRangePositive
     shrink = shrinkTokenBundleSmallRangePositive
 
+instance Arbitrary (Large TokenBundle) where
+    arbitrary = fmap Large $ TokenBundle
+        <$> genCoinLargePositive
+        <*> genTokenMapLarge
+    -- No shrinking
+
+genTokenMapLarge :: Gen TokenMap
+genTokenMapLarge = do
+    assetCount <- frequency
+        [ (1, pure 0)
+        , (1, pure 1)
+        , (8, choose (2, 63))
+        ]
+    TokenMap.fromFlatList <$> replicateM assetCount genAssetQuantity
+  where
+    genAssetQuantity = (,)
+        <$> genAssetIdLargeRange
+        <*> genTokenQuantitySmallPositive
+
 instance Arbitrary TokenMap where
     arbitrary = genTokenMapSmallRange
     shrink = shrinkTokenMapSmallRange

From 2637e70ea41829db761abd41186ed922e2a0e5d8 Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Tue, 2 Mar 2021 10:58:42 +0000
Subject: [PATCH 15/24] Add function `splitBundlesWithExcessiveAssetCounts`.

---
 .../Primitive/CoinSelection/MA/RoundRobin.hs  | 21 ++++++++++
 .../CoinSelection/MA/RoundRobinSpec.hs        | 42 +++++++++++++++++++
 2 files changed, 63 insertions(+)

diff --git a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
index 3e8606df200..485ba929674 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
@@ -56,6 +56,7 @@ module Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
 
     -- * Splitting bundles
     , splitBundleIfAssetCountExcessive
+    , splitBundlesWithExcessiveAssetCounts
     , splitBundlesWithExcessiveTokenQuantities
 
     -- * Grouping and ungrouping
@@ -1224,6 +1225,26 @@ splitBundleIfAssetCountExcessive b isExcessive
   where
     splitInHalf = flip TokenBundle.equipartitionAssets (() :| [()])
 
+-- | Splits bundles with excessive asset counts into smaller bundles.
+--
+-- Only token bundles where the 'isExcessive' indicator function returns 'True'
+-- will be split.
+--
+-- Returns a list of smaller bundles for which 'isExcessive' returns 'False'.
+--
+-- If none of the bundles in the given list has an excessive asset count,
+-- this function will return the original list.
+--
+splitBundlesWithExcessiveAssetCounts
+    :: NonEmpty TokenBundle
+    -- ^ Token bundles.
+    -> (TokenBundle -> Bool)
+    -- ^ A function that returns 'True' if (and only if) the asset count of
+    -- the given bundle is excessive.
+    -> NonEmpty TokenBundle
+splitBundlesWithExcessiveAssetCounts bs isExcessive =
+    (`splitBundleIfAssetCountExcessive` isExcessive) =<< bs
+
 -- | Splits bundles with excessive token quantities into smaller bundles.
 --
 -- Only token bundles containing quantities that exceed the maximum token
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 e80235e4745..c1d83ee73ac 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
@@ -53,6 +53,7 @@ import Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     , runSelection
     , runSelectionStep
     , splitBundleIfAssetCountExcessive
+    , splitBundlesWithExcessiveAssetCounts
     , splitBundlesWithExcessiveTokenQuantities
     , ungroupByKey
     )
@@ -333,6 +334,10 @@ spec = describe "Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobinSpec" $
             property prop_splitBundleIfAssetCountExcessive_postCondition
         it "prop_splitBundleIfAssetCountExcessive_sum" $
             property prop_splitBundleIfAssetCountExcessive_sum
+        it "prop_splitBundlesWithExcessiveAssetCounts_length" $
+            property prop_splitBundlesWithExcessiveAssetCounts_length
+        it "prop_splitBundlesWithExcessiveAssetCounts_sum" $
+            property prop_splitBundlesWithExcessiveAssetCounts_sum
 
     parallel $ describe "Splitting bundles with excessive token quantities" $ do
 
@@ -1886,6 +1891,43 @@ prop_splitBundleIfAssetCountExcessive_sum
   where
     isExcessive = (> maxAssetCount) . Set.size . TokenBundle.getAssets
 
+prop_splitBundlesWithExcessiveAssetCounts_length
+    :: Blind (NonEmpty TokenBundle) -> Positive Int -> Property
+prop_splitBundlesWithExcessiveAssetCounts_length
+    (Blind input) (Positive maxAssetCount) =
+        checkCoverage $ property $
+        cover 5 (lengthOutput > lengthInput)
+            "length has increased" $
+        cover 5 (lengthOutput == lengthInput)
+            "length has remained the same" $
+        case compare lengthOutput lengthInput of
+            GT -> (&&)
+                (F.any isExcessive input)
+                (F.all (not . isExcessive) output)
+            EQ -> (&&)
+                (F.all (not . isExcessive) input)
+                (input == output)
+            LT ->
+                error "length has unexpectedly decreased"
+  where
+    isExcessive =
+        (> maxAssetCount) . Set.size . TokenBundle.getAssets
+    lengthInput =
+        NE.length input
+    lengthOutput =
+        NE.length output
+    output =
+        splitBundlesWithExcessiveAssetCounts input isExcessive
+
+prop_splitBundlesWithExcessiveAssetCounts_sum
+    :: Blind (NonEmpty TokenBundle) -> Positive Int -> Property
+prop_splitBundlesWithExcessiveAssetCounts_sum
+    (Blind bundles) (Positive maxAssetCount) = (===)
+        (F.fold $ splitBundlesWithExcessiveAssetCounts bundles isExcessive)
+        (F.fold bundles)
+  where
+    isExcessive = (> maxAssetCount) . Set.size . TokenBundle.getAssets
+
 --------------------------------------------------------------------------------
 -- Splitting bundles with excessive token quantities
 --------------------------------------------------------------------------------

From 8c3ca7830245356109df7d38daba4aece07256b4 Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Wed, 3 Mar 2021 05:43:26 +0000
Subject: [PATCH 16/24] Add `assessBundleSize` field to `MakeChangeCriteria`.

---
 .../Primitive/CoinSelection/MA/RoundRobin.hs  | 14 +++--
 .../src/Cardano/Wallet/Primitive/Types/Tx.hs  | 13 +++++
 .../CoinSelection/MA/RoundRobinSpec.hs        | 52 +++++++++++++++----
 .../Wallet/Primitive/Types/TokenBundleSpec.hs |  4 +-
 4 files changed, 68 insertions(+), 15 deletions(-)

diff --git a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
index 485ba929674..01a323e5eeb 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
@@ -97,7 +97,7 @@ import Cardano.Wallet.Primitive.Types.TokenMap
 import Cardano.Wallet.Primitive.Types.TokenQuantity
     ( TokenQuantity (..) )
 import Cardano.Wallet.Primitive.Types.Tx
-    ( TxIn, TxOut, txOutCoin )
+    ( TokenBundleSizeAssessment (..), TxIn, TxOut, txOutCoin )
 import Cardano.Wallet.Primitive.Types.UTxOIndex
     ( SelectionFilter (..), UTxOIndex (..) )
 import Control.Monad.Random.Class
@@ -501,6 +501,7 @@ performSelection minCoinFor costFor criteria
         (fmap (TokenMap.getAssets . view #tokens))
         (makeChange MakeChangeCriteria
             { minCoinFor = noMinimumCoin
+            , assessBundleSize = const TokenBundleSizeWithinLimit
             , requiredCost = noCost
             , extraCoinSource
             , inputBundles
@@ -564,6 +565,9 @@ performSelection minCoinFor costFor criteria
         mChangeGenerated :: Either UnableToConstructChangeError [TokenBundle]
         mChangeGenerated = makeChange MakeChangeCriteria
             { minCoinFor
+            -- TODO: pass the implementation of this function in via
+            -- 'performSelection':
+            , assessBundleSize = const TokenBundleSizeWithinLimit
             , requiredCost
             , extraCoinSource
             , inputBundles =  view #tokens . snd <$> inputsSelected
@@ -802,10 +806,12 @@ runSelectionStep lens s
 
 -- | Criteria for the 'makeChange' function.
 --
-data MakeChangeCriteria minCoinFor = MakeChangeCriteria
+data MakeChangeCriteria minCoinFor assessBundleSize = MakeChangeCriteria
     { minCoinFor :: minCoinFor
       -- ^ A function that computes the minimum required ada quantity for a
       -- particular output.
+    , assessBundleSize :: assessBundleSize
+        -- ^ A function to assess the size of a token bundle.
     , requiredCost :: Coin
       -- ^ The minimal (and optimal) delta between the total ada balance
       -- of all input bundles and the total ada balance of all output and
@@ -844,7 +850,9 @@ data MakeChangeCriteria minCoinFor = MakeChangeCriteria
 -- to every output token bundle.
 --
 makeChange
-    :: MakeChangeCriteria (TokenMap -> Coin)
+    :: MakeChangeCriteria
+        (TokenMap -> Coin)
+        (TokenBundle -> TokenBundleSizeAssessment)
         -- ^ Criteria for making change.
     -> Either UnableToConstructChangeError [TokenBundle]
         -- ^ Generated change bundles.
diff --git a/lib/core/src/Cardano/Wallet/Primitive/Types/Tx.hs b/lib/core/src/Cardano/Wallet/Primitive/Types/Tx.hs
index 040ded12da6..937d336dc08 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/Types/Tx.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/Types/Tx.hs
@@ -28,6 +28,7 @@ module Cardano.Wallet.Primitive.Types.Tx
     , UnsignedTx (..)
     , TransactionInfo (..)
     , Direction (..)
+    , TokenBundleSizeAssessment (..)
 
     -- * Functions
     , fromTransactionInfo
@@ -430,3 +431,15 @@ txMetadataIsNull (TxMetadata md) = Map.null md
 toTxHistory :: TransactionInfo -> (Tx, TxMeta)
 toTxHistory info =
     (fromTransactionInfo info, txInfoMeta info)
+
+-- | Indicates the size of a token bundle relative to the upper limit of what
+--   can be included in a single transaction output, defined by the protocol.
+--
+data TokenBundleSizeAssessment
+    = TokenBundleSizeWithinLimit
+    -- ^ Indicates that the size of a token bundle does not exceed the maximum
+    -- size that can be included in a transaction output.
+    | TokenBundleSizeExceedsLimit
+    -- ^ Indicates that the size of a token bundle exceeds the maximum size
+    -- that can be included in a transaction output.
+    deriving (Eq, Generic, Show)
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 c1d83ee73ac..7593cf09dc9 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
@@ -24,8 +24,8 @@ import Cardano.Numeric.Util
 import Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     ( AssetCount (..)
     , BalanceInsufficientError (..)
-    , MakeChangeCriteria (..)
     , InsufficientMinCoinValueError (..)
+    , MakeChangeCriteria (..)
     , SelectionCriteria (..)
     , SelectionError (..)
     , SelectionInsufficientError (..)
@@ -91,7 +91,7 @@ import Cardano.Wallet.Primitive.Types.TokenQuantity
 import Cardano.Wallet.Primitive.Types.TokenQuantity.Gen
     ( genTokenQuantitySmallPositive, shrinkTokenQuantitySmallPositive )
 import Cardano.Wallet.Primitive.Types.Tx
-    ( TxIn (..), TxOut (..), txOutCoin )
+    ( TokenBundleSizeAssessment (..), TxIn (..), TxOut (..), txOutCoin )
 import Cardano.Wallet.Primitive.Types.Tx.Gen
     ( genTxOutSmallRange, shrinkTxOutSmallRange )
 import Cardano.Wallet.Primitive.Types.UTxOIndex
@@ -1331,7 +1331,16 @@ linearCost SelectionSkeleton{inputsSkeleton, outputsSkeleton, changeSkeleton}
     + F.length outputsSkeleton
     + F.length changeSkeleton
 
-type MakeChangeData = MakeChangeCriteria MinCoinValueFor
+type MakeChangeData = MakeChangeCriteria MinCoinValueFor BundleSizeAssessor
+
+data BundleSizeAssessor
+    = NoBundleSizeLimit
+    deriving (Eq, Show)
+
+mkBundleSizeAssessor
+    :: BundleSizeAssessor
+    -> (TokenBundle -> TokenBundleSizeAssessment)
+mkBundleSizeAssessor NoBundleSizeLimit = const TokenBundleSizeWithinLimit
 
 isValidMakeChangeData :: MakeChangeData -> Bool
 isValidMakeChangeData p = (&&)
@@ -1350,6 +1359,7 @@ genMakeChangeData = flip suchThat isValidMakeChangeData $ do
     let inputBundleCount = outputBundleCount * 4
     MakeChangeCriteria
         <$> arbitrary
+        <*> pure NoBundleSizeLimit
         <*> genCoinSmall
         <*> oneof [pure Nothing, Just <$> genCoinSmallPositive]
         <*> genTokenBundles inputBundleCount
@@ -1364,14 +1374,24 @@ makeChangeWith
     :: MakeChangeData
     -> Either UnableToConstructChangeError [TokenBundle]
 makeChangeWith p = makeChange p
-    { minCoinFor = mkMinCoinValueFor $ minCoinFor p}
+    { minCoinFor = mkMinCoinValueFor $ minCoinFor p
+    , assessBundleSize = mkBundleSizeAssessor $ assessBundleSize p
+    }
 
 prop_makeChange_identity
     :: NonEmpty TokenBundle -> Property
 prop_makeChange_identity bundles = (===)
-    (F.fold <$> makeChange
-        (MakeChangeCriteria (const (Coin 0)) (Coin 0) Nothing bundles bundles))
+    (F.fold <$> makeChange criteria)
     (Right TokenBundle.empty)
+  where
+    criteria = MakeChangeCriteria
+        { minCoinFor = const (Coin 0)
+        , requiredCost = Coin 0
+        , extraCoinSource = Nothing
+        , assessBundleSize = mkBundleSizeAssessor NoBundleSizeLimit
+        , inputBundles = bundles
+        , outputBundles = bundles
+        }
 
 prop_makeChange_length
     :: MakeChangeData
@@ -1381,7 +1401,11 @@ prop_makeChange_length p =
         Left{} -> property False
         Right xs -> length xs === length (outputBundles p)
   where
-    change = makeChange p {minCoinFor = noMinCoin, requiredCost = noCost}
+    change = makeChange p
+        { minCoinFor = noMinCoin
+        , requiredCost = noCost
+        , assessBundleSize = mkBundleSizeAssessor NoBundleSizeLimit
+        }
 
 prop_makeChange
     :: MakeChangeData
@@ -1540,11 +1564,19 @@ prop_makeChange_fail_minValueTooBig p =
 unit_makeChange
     :: [Expectation]
 unit_makeChange =
-    [ makeChange (MakeChangeCriteria minCoinValueFor cost extraSource i o)
-        `shouldBe` expectation
-    | (minCoinValueFor, cost, extraSource, i, o, expectation) <- matrix
+    [ makeChange criteria `shouldBe` expectation
+    | (minCoinFor, requiredCost, extraCoinSource, i, o, expectation) <- matrix
+    , let criteria = MakeChangeCriteria
+              { minCoinFor
+              , requiredCost
+              , extraCoinSource
+              , assessBundleSize
+              , inputBundles = i
+              , outputBundles = o
+              }
     ]
   where
+    assessBundleSize = mkBundleSizeAssessor NoBundleSizeLimit
     matrix =
         -- Simple, only ada, should construct a single change output with 1 ada.
         [ ( noMinCoin, noCost
diff --git a/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenBundleSpec.hs b/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenBundleSpec.hs
index 06c5a996836..3ff0129063f 100644
--- a/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenBundleSpec.hs
+++ b/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenBundleSpec.hs
@@ -11,10 +11,10 @@ import Prelude hiding
 
 import Algebra.PartialOrd
     ( leq )
-import Cardano.Wallet.Primitive.Types.TokenBundle
-    ( TokenBundle (..), add, difference, isCoin, subtract, unsafeSubtract )
 import Cardano.Numeric.Util
     ( inAscendingPartialOrder )
+import Cardano.Wallet.Primitive.Types.TokenBundle
+    ( TokenBundle (..), add, difference, isCoin, subtract, unsafeSubtract )
 import Cardano.Wallet.Primitive.Types.TokenBundle.Gen
     ( genTokenBundleSmallRange, shrinkTokenBundleSmallRange )
 import Cardano.Wallet.Primitive.Types.TokenQuantity

From 9efaf80eaa5a6c36e4082c25e31ab3f00a2df20a Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Wed, 3 Mar 2021 10:46:45 +0000
Subject: [PATCH 17/24] Add `assessBundleSize` argument to `performSelection`.

---
 lib/core/src/Cardano/Wallet.hs                   |  4 ++++
 .../Primitive/CoinSelection/MA/RoundRobin.hs     | 10 +++++-----
 .../Primitive/CoinSelection/MA/RoundRobinSpec.hs | 16 ++++++++++++----
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/lib/core/src/Cardano/Wallet.hs b/lib/core/src/Cardano/Wallet.hs
index b19bdca9704..6570d80d8c5 100644
--- a/lib/core/src/Cardano/Wallet.hs
+++ b/lib/core/src/Cardano/Wallet.hs
@@ -248,6 +248,7 @@ import Cardano.Wallet.Primitive.AddressDiscovery.Sequential
 import Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     ( SelectionError (..)
     , SelectionResult (..)
+    , TokenBundleSizeAssessment (..)
     , UnableToConstructChangeError (..)
     , emptySkeleton
     , performSelection
@@ -1403,6 +1404,9 @@ selectAssets ctx (utxo, cp, pending) tx outs transform = do
     sel <- performSelection
         (calcMinimumCoinValue tl pp)
         (calcMinimumCost tl pp tx)
+        -- TODO: Pass in the real implementation of this function here,
+        -- as determined by the protocol:
+        (const TokenBundleSizeWithinLimit)
         (initSelectionCriteria tl pp tx utxo outs)
     liftIO $ traceWith tr $ MsgSelectionDone sel
     withExceptT ErrSelectAssetsSelectionError $ except $
diff --git a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
index 01a323e5eeb..78849bca821 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
@@ -407,10 +407,12 @@ performSelection
         -- ^ A function that computes the extra cost corresponding to a given
         -- selection. This function must not depend on the magnitudes of
         -- individual asset quantities held within each change output.
+    -> (TokenBundle -> TokenBundleSizeAssessment)
+        -- ^ A function that assesses the size of a token bundle.
     -> SelectionCriteria
         -- ^ The selection goal to satisfy.
     -> m (Either SelectionError (SelectionResult TokenBundle))
-performSelection minCoinFor costFor criteria
+performSelection minCoinFor costFor assessBundleSize criteria
     | not (balanceRequired `leq` balanceAvailable) =
         pure $ Left $ BalanceInsufficient $ BalanceInsufficientError
             { balanceAvailable, balanceRequired }
@@ -501,7 +503,7 @@ performSelection minCoinFor costFor criteria
         (fmap (TokenMap.getAssets . view #tokens))
         (makeChange MakeChangeCriteria
             { minCoinFor = noMinimumCoin
-            , assessBundleSize = const TokenBundleSizeWithinLimit
+            , assessBundleSize
             , requiredCost = noCost
             , extraCoinSource
             , inputBundles
@@ -565,9 +567,7 @@ performSelection minCoinFor costFor criteria
         mChangeGenerated :: Either UnableToConstructChangeError [TokenBundle]
         mChangeGenerated = makeChange MakeChangeCriteria
             { minCoinFor
-            -- TODO: pass the implementation of this function in via
-            -- 'performSelection':
-            , assessBundleSize = const TokenBundleSizeWithinLimit
+            , assessBundleSize
             , requiredCost
             , extraCoinSource
             , inputBundles =  view #tokens . snd <$> inputsSelected
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 7593cf09dc9..e7e9b67b8b8 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
@@ -629,10 +629,11 @@ prop_performSelection minCoinValueFor costFor (Blind criteria) coverage =
             , "selectionLimit:"
             , show selectionLimit
             ]
-        result <- run (performSelection
+        result <- run $ performSelection
             (mkMinCoinValueFor minCoinValueFor)
             (mkCostFor costFor)
-            criteria)
+            (mkBundleSizeAssessor NoBundleSizeLimit)
+            (criteria)
         monitor (coverage result)
         either onFailure onSuccess result
   where
@@ -776,7 +777,11 @@ prop_performSelection minCoinValueFor costFor (Blind criteria) coverage =
         monitor $ counterexample $ show e
         assert (shortfall e > Coin 0)
         let criteria' = criteria { selectionLimit = NoLimit }
-        run (performSelection noMinCoin (const noCost) criteria') >>= \case
+        let assessBundleSize =
+                mkBundleSizeAssessor NoBundleSizeLimit
+        let performSelection' = performSelection
+                noMinCoin (const noCost) assessBundleSize criteria'
+        run performSelection' >>= \case
             Left e' -> do
                 monitor $ counterexample $ unlines
                     [ "Failed to re-run selection with no cost!"
@@ -1117,7 +1122,10 @@ type BoundaryTestEntry = (Coin, [(AssetId, TokenQuantity)])
 mkBoundaryTestExpectation :: BoundaryTestData -> Expectation
 mkBoundaryTestExpectation (BoundaryTestData criteria expectedResult) = do
     actualResult <- performSelection
-        noMinCoin (mkCostFor NoCost) (encodeBoundaryTestCriteria criteria)
+        (noMinCoin)
+        (mkCostFor NoCost)
+        (mkBundleSizeAssessor NoBundleSizeLimit)
+        (encodeBoundaryTestCriteria criteria)
     fmap decodeBoundaryTestResult actualResult `shouldBe` Right expectedResult
 
 encodeBoundaryTestCriteria :: BoundaryTestCriteria -> SelectionCriteria

From 3acd58e9965114c61f55629d3d971bf3dcce5cf5 Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Wed, 3 Mar 2021 06:01:28 +0000
Subject: [PATCH 18/24] Split up change maps with excessive asset counts in
 `makeChange`.

---
 lib/core/src/Cardano/Wallet.hs                |  2 +-
 .../Primitive/CoinSelection/MA/RoundRobin.hs  | 28 +++++++++++++++----
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/lib/core/src/Cardano/Wallet.hs b/lib/core/src/Cardano/Wallet.hs
index 6570d80d8c5..935412a80db 100644
--- a/lib/core/src/Cardano/Wallet.hs
+++ b/lib/core/src/Cardano/Wallet.hs
@@ -248,7 +248,6 @@ import Cardano.Wallet.Primitive.AddressDiscovery.Sequential
 import Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     ( SelectionError (..)
     , SelectionResult (..)
-    , TokenBundleSizeAssessment (..)
     , UnableToConstructChangeError (..)
     , emptySkeleton
     , performSelection
@@ -314,6 +313,7 @@ import Cardano.Wallet.Primitive.Types.TokenBundle
 import Cardano.Wallet.Primitive.Types.Tx
     ( Direction (..)
     , SealedTx (..)
+    , TokenBundleSizeAssessment (..)
     , TransactionInfo (..)
     , Tx
     , TxChange (..)
diff --git a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
index 78849bca821..2e6d73385c4 100644
--- a/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
+++ b/lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs
@@ -830,6 +830,18 @@ data MakeChangeCriteria minCoinFor assessBundleSize = MakeChangeCriteria
         -- ^ Token bundles of original outputs.
     } deriving (Eq, Generic, Show)
 
+-- | Indicates 'True' if and only if a token bundle exceeds the maximum size
+--   that can be included in a transaction output.
+--
+tokenBundleSizeExceedsLimit
+    :: (TokenBundle -> TokenBundleSizeAssessment) -> TokenBundle -> Bool
+tokenBundleSizeExceedsLimit calculateBundleSize b =
+    case calculateBundleSize b of
+        TokenBundleSizeWithinLimit->
+            False
+        TokenBundleSizeExceedsLimit ->
+            True
+
 -- | Constructs change bundles for a set of selected inputs and outputs.
 --
 -- Returns 'Nothing' if the specified inputs do not provide enough ada to
@@ -869,6 +881,7 @@ makeChange criteria
   where
     MakeChangeCriteria
         { minCoinFor
+        , assessBundleSize
         , requiredCost
         , extraCoinSource
         , inputBundles
@@ -908,13 +921,13 @@ makeChange criteria
         -- asset count:
         & NE.zipWith (\m1 (m2, c) -> (m1 <> m2, c))
             changeForNonUserSpecifiedAssets
-        -- Finally, if there are any maps with excessive token quantities, then
+        -- Finally, if there are any maps that are oversized (in any way), then
         -- split these maps up along with their corresponding output coins:
-        & splitMapsWithExcessiveQuantities
+        & splitOversizedMaps
       where
-        splitMapsWithExcessiveQuantities
+        splitOversizedMaps
             :: NonEmpty (TokenMap, Coin) -> NonEmpty (TokenMap, Coin)
-        splitMapsWithExcessiveQuantities =
+        splitOversizedMaps =
             -- For the sake of convenience when splitting up change maps and
             -- output coins (which are treated as weights), treat each change
             -- map and its corresponding output coin as a token bundle.
@@ -922,8 +935,11 @@ makeChange criteria
           where
             bundle (m, c) = TokenBundle c m
             unbundle (TokenBundle c m) = (m, c)
-            split = flip splitBundlesWithExcessiveTokenQuantities
-                maxTxOutTokenQuantity
+            split b = b
+                & flip splitBundlesWithExcessiveAssetCounts
+                    (tokenBundleSizeExceedsLimit assessBundleSize)
+                & flip splitBundlesWithExcessiveTokenQuantities
+                    maxTxOutTokenQuantity
 
     -- Change for user-specified assets: assets that were present in the
     -- original set of user-specified outputs ('outputsToCover').

From 801b5e6188efa686f719676d5f191667e81e15c4 Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Thu, 4 Mar 2021 03:43:36 +0000
Subject: [PATCH 19/24] Add `bundleSizeAssessor` field to
 `BoundaryTestCriteria`.

---
 .../Primitive/CoinSelection/MA/RoundRobinSpec.hs       | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

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 e7e9b67b8b8..d12722d5fd6 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
@@ -1102,7 +1102,9 @@ data BoundaryTestData = BoundaryTestData
     deriving (Eq, Show)
 
 data BoundaryTestCriteria = BoundaryTestCriteria
-    { boundaryTestOutputs
+    { boundaryTestBundleSizeAssessor
+        :: BundleSizeAssessor
+    , boundaryTestOutputs
         :: [BoundaryTestEntry]
     , boundaryTestUTxO
         :: [BoundaryTestEntry]
@@ -1124,7 +1126,7 @@ mkBoundaryTestExpectation (BoundaryTestData criteria expectedResult) = do
     actualResult <- performSelection
         (noMinCoin)
         (mkCostFor NoCost)
-        (mkBundleSizeAssessor NoBundleSizeLimit)
+        (mkBundleSizeAssessor $ boundaryTestBundleSizeAssessor criteria)
         (encodeBoundaryTestCriteria criteria)
     fmap decodeBoundaryTestResult actualResult `shouldBe` Right expectedResult
 
@@ -1182,6 +1184,7 @@ boundaryTest1 = BoundaryTestData
   where
     assetA = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "1")
     (q1, q2) = (TokenQuantity 1, TokenQuantity.pred maxTxOutTokenQuantity)
+    boundaryTestBundleSizeAssessor = NoBundleSizeLimit
     boundaryTestOutputs =
       [ (Coin 1_500_000, []) ]
     boundaryTestUTxO =
@@ -1211,6 +1214,7 @@ boundaryTest2 = BoundaryTestData
   where
     assetA = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "1")
     q1 :| [q2] = TokenQuantity.equipartition maxTxOutTokenQuantity (() :| [()])
+    boundaryTestBundleSizeAssessor = NoBundleSizeLimit
     boundaryTestOutputs =
       [ (Coin 1_500_000, []) ]
     boundaryTestUTxO =
@@ -1241,6 +1245,7 @@ boundaryTest3 = BoundaryTestData
     assetA = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "1")
     q1 :| [q2] = TokenQuantity.equipartition
         (TokenQuantity.succ maxTxOutTokenQuantity) (() :| [()])
+    boundaryTestBundleSizeAssessor = NoBundleSizeLimit
     boundaryTestOutputs =
       [ (Coin 1_500_000, []) ]
     boundaryTestUTxO =
@@ -1271,6 +1276,7 @@ boundaryTest4 = BoundaryTestData
     }
   where
     assetA = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "1")
+    boundaryTestBundleSizeAssessor = NoBundleSizeLimit
     boundaryTestOutputs =
       [ (Coin 1_500_000, []) ]
     boundaryTestUTxO =

From f87a280aab50deb1feb588103f881757c8a61478 Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Thu, 4 Mar 2021 03:53:23 +0000
Subject: [PATCH 20/24] Add `BundleAssetCountUpperLimit` constructor to
 `BundleSizeAssessor`.

---
 .../Primitive/CoinSelection/MA/RoundRobinSpec.hs  | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

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 d12722d5fd6..48a5ad5c098 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
@@ -1349,12 +1349,25 @@ type MakeChangeData = MakeChangeCriteria MinCoinValueFor BundleSizeAssessor
 
 data BundleSizeAssessor
     = NoBundleSizeLimit
+      -- ^ Indicates that there is no limit on a token bundle's size.
+    | BundleAssetCountUpperLimit Int
+      -- ^ Indicates an inclusive upper bound on the number of assets in a
+      -- token bundle.
     deriving (Eq, Show)
 
 mkBundleSizeAssessor
     :: BundleSizeAssessor
     -> (TokenBundle -> TokenBundleSizeAssessment)
-mkBundleSizeAssessor NoBundleSizeLimit = const TokenBundleSizeWithinLimit
+mkBundleSizeAssessor = \case
+    NoBundleSizeLimit ->
+        const TokenBundleSizeWithinLimit
+    BundleAssetCountUpperLimit upperLimit ->
+        \bundle ->
+            let assetCount = Set.size $ TokenBundle.getAssets bundle in
+            case assetCount `compare` upperLimit of
+                LT -> TokenBundleSizeWithinLimit
+                EQ -> TokenBundleSizeWithinLimit
+                GT -> TokenBundleSizeExceedsLimit
 
 isValidMakeChangeData :: MakeChangeData -> Bool
 isValidMakeChangeData p = (&&)

From 3928d55332d0b5f7b5b7d84ce352968d148a8660 Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Thu, 4 Mar 2021 04:42:21 +0000
Subject: [PATCH 21/24] Add boundary tests for `BundleAssetCountUpperLimit`.

---
 .../CoinSelection/MA/RoundRobinSpec.hs        | 197 ++++++++++++++++--
 1 file changed, 179 insertions(+), 18 deletions(-)

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 48a5ad5c098..0f125a985cc 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
@@ -277,8 +277,10 @@ spec = describe "Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobinSpec" $
 
     parallel $ describe "Boundary tests" $ do
 
-        unitTests "testBoundaries"
-            unit_testBoundaries
+        unit_testBoundaries "Large token quantities"
+            boundaryTestMatrix_largeTokenQuantities
+        unit_testBoundaries "Large asset counts"
+            boundaryTestMatrix_largeAssetCounts
 
     parallel $ describe "Making change" $ do
 
@@ -1090,8 +1092,8 @@ prop_coinSelectionLens_givesPriorityToCoins (Blind (Small u)) =
 -- Boundary tests
 --------------------------------------------------------------------------------
 
-unit_testBoundaries :: [Expectation]
-unit_testBoundaries = mkBoundaryTestExpectation <$> boundaryTestMatrix
+unit_testBoundaries :: String -> [BoundaryTestData] -> SpecWith ()
+unit_testBoundaries title = unitTests title . fmap mkBoundaryTestExpectation
 
 data BoundaryTestData = BoundaryTestData
     { boundaryTestCriteria
@@ -1160,12 +1162,16 @@ decodeBoundaryTestResult r = BoundaryTestResult
         TokenBundle.toFlatList <$> view #changeGenerated r
     }
 
-boundaryTestMatrix :: [BoundaryTestData]
-boundaryTestMatrix =
-    [ boundaryTest1
-    , boundaryTest2
-    , boundaryTest3
-    , boundaryTest4
+--------------------------------------------------------------------------------
+-- Boundary tests: handling of large token quantities
+--------------------------------------------------------------------------------
+
+boundaryTestMatrix_largeTokenQuantities :: [BoundaryTestData]
+boundaryTestMatrix_largeTokenQuantities =
+    [ boundaryTest_largeTokenQuantities_1
+    , boundaryTest_largeTokenQuantities_2
+    , boundaryTest_largeTokenQuantities_3
+    , boundaryTest_largeTokenQuantities_4
     ]
 
 -- Reach (but do not exceed) the maximum token quantity by selecting inputs
@@ -1176,8 +1182,8 @@ boundaryTestMatrix =
 --
 -- We expect no splitting of token bundles.
 --
-boundaryTest1 :: BoundaryTestData
-boundaryTest1 = BoundaryTestData
+boundaryTest_largeTokenQuantities_1 :: BoundaryTestData
+boundaryTest_largeTokenQuantities_1 = BoundaryTestData
     { boundaryTestCriteria = BoundaryTestCriteria {..}
     , boundaryTestExpectedResult = BoundaryTestResult {..}
     }
@@ -1206,8 +1212,8 @@ boundaryTest1 = BoundaryTestData
 --
 -- We expect no splitting of token bundles.
 --
-boundaryTest2 :: BoundaryTestData
-boundaryTest2 = BoundaryTestData
+boundaryTest_largeTokenQuantities_2 :: BoundaryTestData
+boundaryTest_largeTokenQuantities_2 = BoundaryTestData
     { boundaryTestCriteria = BoundaryTestCriteria {..}
     , boundaryTestExpectedResult = BoundaryTestResult {..}
     }
@@ -1236,8 +1242,8 @@ boundaryTest2 = BoundaryTestData
 --
 -- We expect splitting of change bundles.
 --
-boundaryTest3 :: BoundaryTestData
-boundaryTest3 = BoundaryTestData
+boundaryTest_largeTokenQuantities_3 :: BoundaryTestData
+boundaryTest_largeTokenQuantities_3 = BoundaryTestData
     { boundaryTestCriteria = BoundaryTestCriteria {..}
     , boundaryTestExpectedResult = BoundaryTestResult {..}
     }
@@ -1269,8 +1275,8 @@ boundaryTest3 = BoundaryTestData
 --
 -- We expect splitting of change bundles.
 --
-boundaryTest4 :: BoundaryTestData
-boundaryTest4 = BoundaryTestData
+boundaryTest_largeTokenQuantities_4 :: BoundaryTestData
+boundaryTest_largeTokenQuantities_4 = BoundaryTestData
     { boundaryTestCriteria = BoundaryTestCriteria {..}
     , boundaryTestExpectedResult = BoundaryTestResult {..}
     }
@@ -1292,6 +1298,161 @@ boundaryTest4 = BoundaryTestData
       , (Coin 250_000, [(assetA, maxTxOutTokenQuantity)])
       ]
 
+--------------------------------------------------------------------------------
+-- Boundary tests: handling of large asset counts
+--------------------------------------------------------------------------------
+
+boundaryTestMatrix_largeAssetCounts :: [BoundaryTestData]
+boundaryTestMatrix_largeAssetCounts =
+    [ boundaryTest_largeAssetCounts_1
+    , boundaryTest_largeAssetCounts_2
+    , boundaryTest_largeAssetCounts_3
+    , boundaryTest_largeAssetCounts_4
+    ]
+
+-- Reach (but do not exceed) the maximum per-bundle asset count.
+--
+-- We expect no splitting of change bundles.
+--
+boundaryTest_largeAssetCounts_1 :: BoundaryTestData
+boundaryTest_largeAssetCounts_1 = BoundaryTestData
+    { boundaryTestCriteria = BoundaryTestCriteria {..}
+    , boundaryTestExpectedResult = BoundaryTestResult {..}
+    }
+  where
+    assetA = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "1")
+    assetB = AssetId (UnsafeTokenPolicyId $ Hash "B") (UnsafeTokenName "1")
+    assetC = AssetId (UnsafeTokenPolicyId $ Hash "C") (UnsafeTokenName "1")
+    assetD = AssetId (UnsafeTokenPolicyId $ Hash "D") (UnsafeTokenName "1")
+    boundaryTestBundleSizeAssessor = BundleAssetCountUpperLimit 4
+    boundaryTestOutputs =
+      [ (Coin 1_000_000, []) ]
+    boundaryTestUTxO =
+      [ (Coin 500_000, [(assetA, TokenQuantity 1)])
+      , (Coin 500_000, [(assetB, TokenQuantity 1)])
+      , (Coin 500_000, [(assetC, TokenQuantity 1)])
+      , (Coin 500_000, [(assetD, TokenQuantity 1)])
+      ]
+    -- Expect that all entries will be selected:
+    boundaryTestInputs = boundaryTestUTxO
+    boundaryTestChange =
+      [ ( Coin 1_000_000
+        , [ (assetA, TokenQuantity 1)
+          , (assetB, TokenQuantity 1)
+          , (assetC, TokenQuantity 1)
+          , (assetD, TokenQuantity 1)
+          ]
+        )
+      ]
+
+-- Exceed the maximum per-bundle asset count of 3.
+--
+-- We expect splitting of change bundles.
+--
+boundaryTest_largeAssetCounts_2 :: BoundaryTestData
+boundaryTest_largeAssetCounts_2 = BoundaryTestData
+    { boundaryTestCriteria = BoundaryTestCriteria {..}
+    , boundaryTestExpectedResult = BoundaryTestResult {..}
+    }
+  where
+    assetA = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "1")
+    assetB = AssetId (UnsafeTokenPolicyId $ Hash "B") (UnsafeTokenName "1")
+    assetC = AssetId (UnsafeTokenPolicyId $ Hash "C") (UnsafeTokenName "1")
+    assetD = AssetId (UnsafeTokenPolicyId $ Hash "D") (UnsafeTokenName "1")
+    boundaryTestBundleSizeAssessor = BundleAssetCountUpperLimit 3
+    boundaryTestOutputs =
+      [ (Coin 1_000_000, []) ]
+    boundaryTestUTxO =
+      [ (Coin 500_000, [(assetA, TokenQuantity 1)])
+      , (Coin 500_000, [(assetB, TokenQuantity 1)])
+      , (Coin 500_000, [(assetC, TokenQuantity 1)])
+      , (Coin 500_000, [(assetD, TokenQuantity 1)])
+      ]
+    -- Expect that all entries will be selected:
+    boundaryTestInputs = boundaryTestUTxO
+    boundaryTestChange =
+      [ ( Coin 500_000
+        , [ (assetA, TokenQuantity 1)
+          , (assetB, TokenQuantity 1)
+          ]
+        )
+      , ( Coin 500_000
+        , [ (assetC, TokenQuantity 1)
+          , (assetD, TokenQuantity 1)
+          ]
+        )
+      ]
+
+-- Exceed the maximum per-bundle asset count of 2.
+--
+-- We expect splitting of change bundles.
+--
+boundaryTest_largeAssetCounts_3 :: BoundaryTestData
+boundaryTest_largeAssetCounts_3 = BoundaryTestData
+    { boundaryTestCriteria = BoundaryTestCriteria {..}
+    , boundaryTestExpectedResult = BoundaryTestResult {..}
+    }
+  where
+    assetA = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "1")
+    assetB = AssetId (UnsafeTokenPolicyId $ Hash "B") (UnsafeTokenName "1")
+    assetC = AssetId (UnsafeTokenPolicyId $ Hash "C") (UnsafeTokenName "1")
+    assetD = AssetId (UnsafeTokenPolicyId $ Hash "D") (UnsafeTokenName "1")
+    boundaryTestBundleSizeAssessor = BundleAssetCountUpperLimit 2
+    boundaryTestOutputs =
+      [ (Coin 1_000_000, []) ]
+    boundaryTestUTxO =
+      [ (Coin 500_000, [(assetA, TokenQuantity 1)])
+      , (Coin 500_000, [(assetB, TokenQuantity 1)])
+      , (Coin 500_000, [(assetC, TokenQuantity 1)])
+      , (Coin 500_000, [(assetD, TokenQuantity 1)])
+      ]
+    -- Expect that all entries will be selected:
+    boundaryTestInputs = boundaryTestUTxO
+    boundaryTestChange =
+      [ ( Coin 500_000
+        , [ (assetA, TokenQuantity 1)
+          , (assetB, TokenQuantity 1)
+          ]
+        )
+      , ( Coin 500_000
+        , [ (assetC, TokenQuantity 1)
+          , (assetD, TokenQuantity 1)
+          ]
+        )
+      ]
+
+-- Exceed the maximum per-bundle asset count of 1.
+--
+-- We expect splitting of change bundles.
+--
+boundaryTest_largeAssetCounts_4 :: BoundaryTestData
+boundaryTest_largeAssetCounts_4 = BoundaryTestData
+    { boundaryTestCriteria = BoundaryTestCriteria {..}
+    , boundaryTestExpectedResult = BoundaryTestResult {..}
+    }
+  where
+    assetA = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "1")
+    assetB = AssetId (UnsafeTokenPolicyId $ Hash "B") (UnsafeTokenName "1")
+    assetC = AssetId (UnsafeTokenPolicyId $ Hash "C") (UnsafeTokenName "1")
+    assetD = AssetId (UnsafeTokenPolicyId $ Hash "D") (UnsafeTokenName "1")
+    boundaryTestBundleSizeAssessor = BundleAssetCountUpperLimit 1
+    boundaryTestOutputs =
+      [ (Coin 1_000_000, []) ]
+    boundaryTestUTxO =
+      [ (Coin 500_000, [(assetA, TokenQuantity 1)])
+      , (Coin 500_000, [(assetB, TokenQuantity 1)])
+      , (Coin 500_000, [(assetC, TokenQuantity 1)])
+      , (Coin 500_000, [(assetD, TokenQuantity 1)])
+      ]
+    -- Expect that all entries will be selected:
+    boundaryTestInputs = boundaryTestUTxO
+    boundaryTestChange =
+      [ (Coin 250_000, [(assetA, TokenQuantity 1)])
+      , (Coin 250_000, [(assetB, TokenQuantity 1)])
+      , (Coin 250_000, [(assetC, TokenQuantity 1)])
+      , (Coin 250_000, [(assetD, TokenQuantity 1)])
+      ]
+
 --------------------------------------------------------------------------------
 -- Making change
 --------------------------------------------------------------------------------

From c6d733aa0f42584b25631b17f764545c7beeded4 Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Thu, 4 Mar 2021 07:18:07 +0000
Subject: [PATCH 22/24] Use `mockAsset` and `mockAssetQuantity` to simplify
 boundary tests.

---
 .../CoinSelection/MA/RoundRobinSpec.hs        | 144 +++++++-----------
 1 file changed, 58 insertions(+), 86 deletions(-)

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 0f125a985cc..4b4e56cab24 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
@@ -106,6 +106,8 @@ import Control.Monad
     ( forM_, replicateM )
 import Data.Bifunctor
     ( bimap, second )
+import Data.ByteString
+    ( ByteString )
 import Data.Function
     ( on, (&) )
 import Data.Functor.Identity
@@ -1188,21 +1190,20 @@ boundaryTest_largeTokenQuantities_1 = BoundaryTestData
     , boundaryTestExpectedResult = BoundaryTestResult {..}
     }
   where
-    assetA = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "1")
     (q1, q2) = (TokenQuantity 1, TokenQuantity.pred maxTxOutTokenQuantity)
     boundaryTestBundleSizeAssessor = NoBundleSizeLimit
     boundaryTestOutputs =
       [ (Coin 1_500_000, []) ]
     boundaryTestUTxO =
-      [ (Coin 1_000_000, [(assetA, q1)])
-      , (Coin 1_000_000, [(assetA, q2)])
+      [ (Coin 1_000_000, [(mockAsset "A", q1)])
+      , (Coin 1_000_000, [(mockAsset "A", q2)])
       ]
     boundaryTestInputs =
-      [ (Coin 1_000_000, [(assetA, q1)])
-      , (Coin 1_000_000, [(assetA, q2)])
+      [ (Coin 1_000_000, [(mockAsset "A", q1)])
+      , (Coin 1_000_000, [(mockAsset "A", q2)])
       ]
     boundaryTestChange =
-      [ (Coin 500_000, [(assetA, maxTxOutTokenQuantity)]) ]
+      [ (Coin 500_000, [(mockAsset "A", maxTxOutTokenQuantity)]) ]
 
 -- Reach (but do not exceed) the maximum token quantity by selecting inputs
 -- with the following quantities:
@@ -1218,21 +1219,20 @@ boundaryTest_largeTokenQuantities_2 = BoundaryTestData
     , boundaryTestExpectedResult = BoundaryTestResult {..}
     }
   where
-    assetA = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "1")
     q1 :| [q2] = TokenQuantity.equipartition maxTxOutTokenQuantity (() :| [()])
     boundaryTestBundleSizeAssessor = NoBundleSizeLimit
     boundaryTestOutputs =
       [ (Coin 1_500_000, []) ]
     boundaryTestUTxO =
-      [ (Coin 1_000_000, [(assetA, q1)])
-      , (Coin 1_000_000, [(assetA, q2)])
+      [ (Coin 1_000_000, [(mockAsset "A", q1)])
+      , (Coin 1_000_000, [(mockAsset "A", q2)])
       ]
     boundaryTestInputs =
-      [ (Coin 1_000_000, [(assetA, q1)])
-      , (Coin 1_000_000, [(assetA, q2)])
+      [ (Coin 1_000_000, [(mockAsset "A", q1)])
+      , (Coin 1_000_000, [(mockAsset "A", q2)])
       ]
     boundaryTestChange =
-      [ (Coin 500_000, [(assetA, maxTxOutTokenQuantity)]) ]
+      [ (Coin 500_000, [(mockAsset "A", maxTxOutTokenQuantity)]) ]
 
 -- Slightly exceed the maximum token quantity by selecting inputs with the
 -- following quantities:
@@ -1248,23 +1248,22 @@ boundaryTest_largeTokenQuantities_3 = BoundaryTestData
     , boundaryTestExpectedResult = BoundaryTestResult {..}
     }
   where
-    assetA = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "1")
     q1 :| [q2] = TokenQuantity.equipartition
         (TokenQuantity.succ maxTxOutTokenQuantity) (() :| [()])
     boundaryTestBundleSizeAssessor = NoBundleSizeLimit
     boundaryTestOutputs =
       [ (Coin 1_500_000, []) ]
     boundaryTestUTxO =
-      [ (Coin 1_000_000, [(assetA, TokenQuantity 1)])
-      , (Coin 1_000_000, [(assetA, maxTxOutTokenQuantity)])
+      [ (Coin 1_000_000, [(mockAsset "A", TokenQuantity 1)])
+      , (Coin 1_000_000, [(mockAsset "A", maxTxOutTokenQuantity)])
       ]
     boundaryTestInputs =
-      [ (Coin 1_000_000, [(assetA, TokenQuantity 1)])
-      , (Coin 1_000_000, [(assetA, maxTxOutTokenQuantity)])
+      [ (Coin 1_000_000, [(mockAsset "A", TokenQuantity 1)])
+      , (Coin 1_000_000, [(mockAsset "A", maxTxOutTokenQuantity)])
       ]
     boundaryTestChange =
-      [ (Coin 250_000, [(assetA, q1)])
-      , (Coin 250_000, [(assetA, q2)])
+      [ (Coin 250_000, [(mockAsset "A", q1)])
+      , (Coin 250_000, [(mockAsset "A", q2)])
       ]
 
 -- Reach (but do not exceed) exactly twice the maximum token quantity by
@@ -1281,21 +1280,20 @@ boundaryTest_largeTokenQuantities_4 = BoundaryTestData
     , boundaryTestExpectedResult = BoundaryTestResult {..}
     }
   where
-    assetA = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "1")
     boundaryTestBundleSizeAssessor = NoBundleSizeLimit
     boundaryTestOutputs =
       [ (Coin 1_500_000, []) ]
     boundaryTestUTxO =
-      [ (Coin 1_000_000, [(assetA, maxTxOutTokenQuantity)])
-      , (Coin 1_000_000, [(assetA, maxTxOutTokenQuantity)])
+      [ (Coin 1_000_000, [(mockAsset "A", maxTxOutTokenQuantity)])
+      , (Coin 1_000_000, [(mockAsset "A", maxTxOutTokenQuantity)])
       ]
     boundaryTestInputs =
-      [ (Coin 1_000_000, [(assetA, maxTxOutTokenQuantity)])
-      , (Coin 1_000_000, [(assetA, maxTxOutTokenQuantity)])
+      [ (Coin 1_000_000, [(mockAsset "A", maxTxOutTokenQuantity)])
+      , (Coin 1_000_000, [(mockAsset "A", maxTxOutTokenQuantity)])
       ]
     boundaryTestChange =
-      [ (Coin 250_000, [(assetA, maxTxOutTokenQuantity)])
-      , (Coin 250_000, [(assetA, maxTxOutTokenQuantity)])
+      [ (Coin 250_000, [(mockAsset "A", maxTxOutTokenQuantity)])
+      , (Coin 250_000, [(mockAsset "A", maxTxOutTokenQuantity)])
       ]
 
 --------------------------------------------------------------------------------
@@ -1320,27 +1318,23 @@ boundaryTest_largeAssetCounts_1 = BoundaryTestData
     , boundaryTestExpectedResult = BoundaryTestResult {..}
     }
   where
-    assetA = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "1")
-    assetB = AssetId (UnsafeTokenPolicyId $ Hash "B") (UnsafeTokenName "1")
-    assetC = AssetId (UnsafeTokenPolicyId $ Hash "C") (UnsafeTokenName "1")
-    assetD = AssetId (UnsafeTokenPolicyId $ Hash "D") (UnsafeTokenName "1")
     boundaryTestBundleSizeAssessor = BundleAssetCountUpperLimit 4
     boundaryTestOutputs =
       [ (Coin 1_000_000, []) ]
     boundaryTestUTxO =
-      [ (Coin 500_000, [(assetA, TokenQuantity 1)])
-      , (Coin 500_000, [(assetB, TokenQuantity 1)])
-      , (Coin 500_000, [(assetC, TokenQuantity 1)])
-      , (Coin 500_000, [(assetD, TokenQuantity 1)])
+      [ (Coin 500_000, [mockAssetQuantity "A" 1])
+      , (Coin 500_000, [mockAssetQuantity "B" 1])
+      , (Coin 500_000, [mockAssetQuantity "C" 1])
+      , (Coin 500_000, [mockAssetQuantity "D" 1])
       ]
     -- Expect that all entries will be selected:
     boundaryTestInputs = boundaryTestUTxO
     boundaryTestChange =
       [ ( Coin 1_000_000
-        , [ (assetA, TokenQuantity 1)
-          , (assetB, TokenQuantity 1)
-          , (assetC, TokenQuantity 1)
-          , (assetD, TokenQuantity 1)
+        , [ mockAssetQuantity "A" 1
+          , mockAssetQuantity "B" 1
+          , mockAssetQuantity "C" 1
+          , mockAssetQuantity "D" 1
           ]
         )
       ]
@@ -1355,32 +1349,20 @@ boundaryTest_largeAssetCounts_2 = BoundaryTestData
     , boundaryTestExpectedResult = BoundaryTestResult {..}
     }
   where
-    assetA = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "1")
-    assetB = AssetId (UnsafeTokenPolicyId $ Hash "B") (UnsafeTokenName "1")
-    assetC = AssetId (UnsafeTokenPolicyId $ Hash "C") (UnsafeTokenName "1")
-    assetD = AssetId (UnsafeTokenPolicyId $ Hash "D") (UnsafeTokenName "1")
     boundaryTestBundleSizeAssessor = BundleAssetCountUpperLimit 3
     boundaryTestOutputs =
       [ (Coin 1_000_000, []) ]
     boundaryTestUTxO =
-      [ (Coin 500_000, [(assetA, TokenQuantity 1)])
-      , (Coin 500_000, [(assetB, TokenQuantity 1)])
-      , (Coin 500_000, [(assetC, TokenQuantity 1)])
-      , (Coin 500_000, [(assetD, TokenQuantity 1)])
+      [ (Coin 500_000, [mockAssetQuantity "A" 1])
+      , (Coin 500_000, [mockAssetQuantity "B" 1])
+      , (Coin 500_000, [mockAssetQuantity "C" 1])
+      , (Coin 500_000, [mockAssetQuantity "D" 1])
       ]
     -- Expect that all entries will be selected:
     boundaryTestInputs = boundaryTestUTxO
     boundaryTestChange =
-      [ ( Coin 500_000
-        , [ (assetA, TokenQuantity 1)
-          , (assetB, TokenQuantity 1)
-          ]
-        )
-      , ( Coin 500_000
-        , [ (assetC, TokenQuantity 1)
-          , (assetD, TokenQuantity 1)
-          ]
-        )
+      [ (Coin 500_000, [mockAssetQuantity "A" 1, mockAssetQuantity "B" 1])
+      , (Coin 500_000, [mockAssetQuantity "C" 1, mockAssetQuantity "D" 1])
       ]
 
 -- Exceed the maximum per-bundle asset count of 2.
@@ -1393,32 +1375,20 @@ boundaryTest_largeAssetCounts_3 = BoundaryTestData
     , boundaryTestExpectedResult = BoundaryTestResult {..}
     }
   where
-    assetA = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "1")
-    assetB = AssetId (UnsafeTokenPolicyId $ Hash "B") (UnsafeTokenName "1")
-    assetC = AssetId (UnsafeTokenPolicyId $ Hash "C") (UnsafeTokenName "1")
-    assetD = AssetId (UnsafeTokenPolicyId $ Hash "D") (UnsafeTokenName "1")
     boundaryTestBundleSizeAssessor = BundleAssetCountUpperLimit 2
     boundaryTestOutputs =
       [ (Coin 1_000_000, []) ]
     boundaryTestUTxO =
-      [ (Coin 500_000, [(assetA, TokenQuantity 1)])
-      , (Coin 500_000, [(assetB, TokenQuantity 1)])
-      , (Coin 500_000, [(assetC, TokenQuantity 1)])
-      , (Coin 500_000, [(assetD, TokenQuantity 1)])
+      [ (Coin 500_000, [mockAssetQuantity "A" 1])
+      , (Coin 500_000, [mockAssetQuantity "B" 1])
+      , (Coin 500_000, [mockAssetQuantity "C" 1])
+      , (Coin 500_000, [mockAssetQuantity "D" 1])
       ]
     -- Expect that all entries will be selected:
     boundaryTestInputs = boundaryTestUTxO
     boundaryTestChange =
-      [ ( Coin 500_000
-        , [ (assetA, TokenQuantity 1)
-          , (assetB, TokenQuantity 1)
-          ]
-        )
-      , ( Coin 500_000
-        , [ (assetC, TokenQuantity 1)
-          , (assetD, TokenQuantity 1)
-          ]
-        )
+      [ (Coin 500_000, [mockAssetQuantity "A" 1, mockAssetQuantity "B" 1])
+      , (Coin 500_000, [mockAssetQuantity "C" 1, mockAssetQuantity "D" 1])
       ]
 
 -- Exceed the maximum per-bundle asset count of 1.
@@ -1431,26 +1401,22 @@ boundaryTest_largeAssetCounts_4 = BoundaryTestData
     , boundaryTestExpectedResult = BoundaryTestResult {..}
     }
   where
-    assetA = AssetId (UnsafeTokenPolicyId $ Hash "A") (UnsafeTokenName "1")
-    assetB = AssetId (UnsafeTokenPolicyId $ Hash "B") (UnsafeTokenName "1")
-    assetC = AssetId (UnsafeTokenPolicyId $ Hash "C") (UnsafeTokenName "1")
-    assetD = AssetId (UnsafeTokenPolicyId $ Hash "D") (UnsafeTokenName "1")
     boundaryTestBundleSizeAssessor = BundleAssetCountUpperLimit 1
     boundaryTestOutputs =
       [ (Coin 1_000_000, []) ]
     boundaryTestUTxO =
-      [ (Coin 500_000, [(assetA, TokenQuantity 1)])
-      , (Coin 500_000, [(assetB, TokenQuantity 1)])
-      , (Coin 500_000, [(assetC, TokenQuantity 1)])
-      , (Coin 500_000, [(assetD, TokenQuantity 1)])
+      [ (Coin 500_000, [mockAssetQuantity "A" 1])
+      , (Coin 500_000, [mockAssetQuantity "B" 1])
+      , (Coin 500_000, [mockAssetQuantity "C" 1])
+      , (Coin 500_000, [mockAssetQuantity "D" 1])
       ]
     -- Expect that all entries will be selected:
     boundaryTestInputs = boundaryTestUTxO
     boundaryTestChange =
-      [ (Coin 250_000, [(assetA, TokenQuantity 1)])
-      , (Coin 250_000, [(assetB, TokenQuantity 1)])
-      , (Coin 250_000, [(assetC, TokenQuantity 1)])
-      , (Coin 250_000, [(assetD, TokenQuantity 1)])
+      [ (Coin 250_000, [mockAssetQuantity "A" 1])
+      , (Coin 250_000, [mockAssetQuantity "B" 1])
+      , (Coin 250_000, [mockAssetQuantity "C" 1])
+      , (Coin 250_000, [mockAssetQuantity "D" 1])
       ]
 
 --------------------------------------------------------------------------------
@@ -2376,6 +2342,12 @@ addExtraSource extraSource =
     TokenBundle.add
         (maybe TokenBundle.empty TokenBundle.fromCoin extraSource)
 
+mockAsset :: ByteString -> AssetId
+mockAsset a = AssetId (UnsafeTokenPolicyId $ Hash a) (UnsafeTokenName "1")
+
+mockAssetQuantity :: ByteString -> Natural -> (AssetId, TokenQuantity)
+mockAssetQuantity a q = (mockAsset a, TokenQuantity q)
+
 unitTests :: String -> [Expectation] -> SpecWith ()
 unitTests lbl cases =
     forM_ (zip [1..] cases) $ \(i, test) ->

From d744df3995e3dc98a6c02ebd71a5dec6d4b567cf Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Fri, 5 Mar 2021 06:30:43 +0000
Subject: [PATCH 23/24] Add token-bundle size assessment to the transaction
 layer.

---
 lib/core/src/Cardano/Wallet.hs                |  4 +--
 lib/core/src/Cardano/Wallet/Transaction.hs    | 15 ++++++++-
 .../Cardano/Wallet/Shelley/Compatibility.hs   | 31 +++++++++++++++++++
 .../src/Cardano/Wallet/Shelley/Transaction.hs | 11 ++++++-
 4 files changed, 56 insertions(+), 5 deletions(-)

diff --git a/lib/core/src/Cardano/Wallet.hs b/lib/core/src/Cardano/Wallet.hs
index 935412a80db..fbbbaedf60a 100644
--- a/lib/core/src/Cardano/Wallet.hs
+++ b/lib/core/src/Cardano/Wallet.hs
@@ -1404,9 +1404,7 @@ selectAssets ctx (utxo, cp, pending) tx outs transform = do
     sel <- performSelection
         (calcMinimumCoinValue tl pp)
         (calcMinimumCost tl pp tx)
-        -- TODO: Pass in the real implementation of this function here,
-        -- as determined by the protocol:
-        (const TokenBundleSizeWithinLimit)
+        (assessTokenBundleSize tl)
         (initSelectionCriteria tl pp tx utxo outs)
     liftIO $ traceWith tr $ MsgSelectionDone sel
     withExceptT ErrSelectAssetsSelectionError $ except $
diff --git a/lib/core/src/Cardano/Wallet/Transaction.hs b/lib/core/src/Cardano/Wallet/Transaction.hs
index eb949a7c323..55501302120 100644
--- a/lib/core/src/Cardano/Wallet/Transaction.hs
+++ b/lib/core/src/Cardano/Wallet/Transaction.hs
@@ -48,10 +48,17 @@ import Cardano.Wallet.Primitive.Types.Coin
     ( Coin (..) )
 import Cardano.Wallet.Primitive.Types.RewardAccount
     ( RewardAccount )
+import Cardano.Wallet.Primitive.Types.TokenBundle
+    ( TokenBundle )
 import Cardano.Wallet.Primitive.Types.TokenMap
     ( TokenMap )
 import Cardano.Wallet.Primitive.Types.Tx
-    ( SealedTx (..), Tx (..), TxMetadata, TxOut )
+    ( SealedTx (..)
+    , TokenBundleSizeAssessment (..)
+    , Tx (..)
+    , TxMetadata
+    , TxOut
+    )
 import Cardano.Wallet.Primitive.Types.UTxOIndex
     ( UTxOIndex )
 import Data.ByteString
@@ -117,6 +124,12 @@ data TransactionLayer k = TransactionLayer
         -> Coin
         -- ^ The minimum ada value needed in a UTxO carrying the asset bundle
 
+    , assessTokenBundleSize
+        :: TokenBundle
+        -- ^ A token bundle
+        -> TokenBundleSizeAssessment
+        -- ^ An assessment of the token bundle's size.
+
     , decodeSignedTx
         :: AnyCardanoEra
         -> ByteString
diff --git a/lib/shelley/src/Cardano/Wallet/Shelley/Compatibility.hs b/lib/shelley/src/Cardano/Wallet/Shelley/Compatibility.hs
index 9c88b2ab9ff..c2e72add9fb 100644
--- a/lib/shelley/src/Cardano/Wallet/Shelley/Compatibility.hs
+++ b/lib/shelley/src/Cardano/Wallet/Shelley/Compatibility.hs
@@ -73,6 +73,9 @@ module Cardano.Wallet.Shelley.Compatibility
     , toCardanoValue
     , fromCardanoValue
 
+      -- ** Tests
+    , assessTokenBundleSize
+
       -- ** Stake pools
     , fromPoolId
     , fromPoolDistr
@@ -260,6 +263,7 @@ import Type.Reflection
 import qualified Cardano.Address.Style.Shelley as CA
 import qualified Cardano.Api.Shelley as Cardano
 import qualified Cardano.Api.Typed as Cardano
+import qualified Cardano.Binary as Binary
 import qualified Cardano.Byron.Codec.Cbor as CBOR
 import qualified Cardano.Chain.Common as Byron
 import qualified Cardano.Ledger.Core as SL.Core
@@ -1216,6 +1220,33 @@ toStakePoolDlgCert xpub (W.PoolId pid) =
     cred = SL.KeyHash $ UnsafeHash $ toShort $ blake2b224 $ xpubPublicKey xpub
     pool = SL.KeyHash $ UnsafeHash $ toShort pid
 
+{-------------------------------------------------------------------------------
+                                 Tests
+-------------------------------------------------------------------------------}
+
+-- | Assesses a token bundle size in relation to the maximum size that can be
+--   included in a transaction output.
+--
+assessTokenBundleSize :: TokenBundle.TokenBundle -> W.TokenBundleSizeAssessment
+assessTokenBundleSize tb
+    | BS.length serialRepresentation <= serialRepresentationMaxLengthBytes =
+        W.TokenBundleSizeWithinLimit
+    | otherwise =
+        W.TokenBundleSizeExceedsLimit
+  where
+    serialRepresentation :: ByteString
+    serialRepresentation =
+        Binary.serialize' $ Cardano.toMaryValue $ toCardanoValue tb
+
+    -- NOTE: This hard-coded limit may change in future. Ideally, we should
+    -- delegate the assessment of whether a token bundle is too large to a
+    -- function exported by Cardano API.
+    --
+    -- See: https://jira.iohk.io/projects/ADP/issues/ADP-779
+    --
+    serialRepresentationMaxLengthBytes :: Int
+    serialRepresentationMaxLengthBytes = 4096
+
 {-------------------------------------------------------------------------------
                       Address Encoding / Decoding
 -------------------------------------------------------------------------------}
diff --git a/lib/shelley/src/Cardano/Wallet/Shelley/Transaction.hs b/lib/shelley/src/Cardano/Wallet/Shelley/Transaction.hs
index d1630ae1d6f..c7f36a2e40b 100644
--- a/lib/shelley/src/Cardano/Wallet/Shelley/Transaction.hs
+++ b/lib/shelley/src/Cardano/Wallet/Shelley/Transaction.hs
@@ -86,7 +86,13 @@ import Cardano.Wallet.Primitive.Types.TokenMap
 import Cardano.Wallet.Primitive.Types.TokenPolicy
     ( TokenName (..) )
 import Cardano.Wallet.Primitive.Types.Tx
-    ( SealedTx (..), Tx (..), TxIn (..), TxOut (..), txOutCoin )
+    ( SealedTx (..)
+    , TokenBundleSizeAssessment (..)
+    , Tx (..)
+    , TxIn (..)
+    , TxOut (..)
+    , txOutCoin
+    )
 import Cardano.Wallet.Shelley.Compatibility
     ( AllegraEra
     , CardanoEra (MaryEra)
@@ -148,6 +154,7 @@ import qualified Cardano.Crypto.Wallet as Crypto.HD
 import qualified Cardano.Ledger.Core as SL
 import qualified Cardano.Wallet.Primitive.Types.TokenBundle as TokenBundle
 import qualified Cardano.Wallet.Primitive.Types.UTxOIndex as UTxOIndex
+import qualified Cardano.Wallet.Shelley.Compatibility as Compatibility
 import qualified Codec.CBOR.Encoding as CBOR
 import qualified Codec.CBOR.Write as CBOR
 import qualified Data.ByteString as BS
@@ -345,6 +352,8 @@ newTransactionLayer networkId = TransactionLayer
     , calcMinimumCoinValue =
         _calcMinimumCoinValue
 
+    , assessTokenBundleSize = Compatibility.assessTokenBundleSize
+
     , decodeSignedTx =
         _decodeSignedTx
     }

From 7d1a6326485e9e62c7f6a40dfb635b55628e59ac Mon Sep 17 00:00:00 2001
From: Jonathan Knowles <jonathan.knowles@iohk.io>
Date: Fri, 5 Mar 2021 07:01:13 +0000
Subject: [PATCH 24/24] Fix hlint errors.

---
 .../Wallet/Primitive/CoinSelection/MA/RoundRobinSpec.hs       | 2 +-
 .../unit/Cardano/Wallet/Primitive/Types/TokenBundleSpec.hs    | 2 +-
 .../test/unit/Cardano/Wallet/Primitive/Types/TokenMapSpec.hs  | 4 ++--
 3 files changed, 4 insertions(+), 4 deletions(-)

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 4b4e56cab24..6d15c63e166 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
@@ -2033,7 +2033,7 @@ prop_splitBundleIfAssetCountExcessive_length
         cover 5 (resultLength >= 2 && resultLength < 8)
             "length >= 2 && length < 8" $
         cover 5 (resultLength >= 8 && resultLength < 16)
-            "length >= 8 && length < 16" $
+            "length >= 8 && length < 16"
         True
   where
     isExcessive = (> maxAssetCount) . Set.size . TokenBundle.getAssets
diff --git a/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenBundleSpec.hs b/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenBundleSpec.hs
index 3ff0129063f..096489334bb 100644
--- a/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenBundleSpec.hs
+++ b/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenBundleSpec.hs
@@ -1,6 +1,6 @@
 {-# LANGUAGE TypeApplications #-}
-
 {-# OPTIONS_GHC -fno-warn-orphans #-}
+{- HLINT ignore "Use camelCase" -}
 
 module Cardano.Wallet.Primitive.Types.TokenBundleSpec
     ( spec
diff --git a/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenMapSpec.hs b/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenMapSpec.hs
index 768b3f6c2b8..5661281cb49 100644
--- a/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenMapSpec.hs
+++ b/lib/core/test/unit/Cardano/Wallet/Primitive/Types/TokenMapSpec.hs
@@ -3,8 +3,8 @@
 {-# LANGUAGE QuasiQuotes #-}
 {-# LANGUAGE TemplateHaskell #-}
 {-# LANGUAGE TypeApplications #-}
-
 {-# OPTIONS_GHC -fno-warn-orphans #-}
+{- HLINT ignore "Use camelCase" -}
 
 module Cardano.Wallet.Primitive.Types.TokenMapSpec
     ( spec
@@ -564,7 +564,7 @@ prop_equipartitionAssets_coverage m = checkCoverage $
     cover 20 (2 <= assetCount && assetCount <= 31)
         "2 <= asset count <= 31" $
     cover 20 (32 <= assetCount && assetCount <= 63)
-        "32 <= asset count <= 63" $
+        "32 <= asset count <= 63"
     True
   where
     assetCount = Set.size $ TokenMap.getAssets $ getLarge $ getBlind m