Move equipartitionTokenMap into TokenMap.

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
@@ -86,7 +85,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
@@ -142,7 +141,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
@@ -1199,29 +1197,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
 --------------------------------------------------------------------------------
@@ -1282,7 +1257,7 @@ equipartitionTokenMapWithMaxQuantity m (TokenQuantity maxQuantity)
     | currentMaxQuantity <= maxQuantity =
         m :| []
     | otherwise =
-        equipartitionTokenMap m (() :| replicate extraPartCount ())
+        TokenMap.equipartitionQuantities m (() :| replicate extraPartCount ())
   where
     TokenQuantity currentMaxQuantity = TokenMap.maximumQuantity m
 

lib/core/src/Cardano/Wallet/Primitive/Types/TokenMap.hs

@@ -70,6 +70,9 @@ module Cardano.Wallet.Primitive.Types.TokenMap
     , removeQuantity
     , maximumQuantity
 
+    -- * Partitioning
+    , equipartitionQuantities
+
     -- * Policies
     , hasPolicy
 
@@ -615,6 +618,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
 --------------------------------------------------------------------------------

lib/core/test/unit/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobinSpec.hs

@@ -39,7 +39,6 @@ import Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     , coinSelectionLens
     , equipartitionTokenBundleWithMaxQuantity
     , equipartitionTokenBundlesWithMaxQuantity
-    , equipartitionTokenMap
     , equipartitionTokenMapWithMaxQuantity
     , fullBalance
     , groupByKey
@@ -320,17 +319,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
@@ -1837,55 +1825,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
 --------------------------------------------------------------------------------

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 )
 import Cardano.Wallet.Primitive.Types.TokenMap.Gen
@@ -102,7 +104,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" $
@@ -192,6 +193,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
@@ -435,6 +447,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
 --------------------------------------------------------------------------------