Add function TokenMap.equipartitionAssets.

cardano-foundation · Mar 2, 2021 · ead92ac · ead92ac
1 parent 05d9201
commit ead92ac
Show file tree

Hide file tree

Showing 2 changed files with 117 additions and 0 deletions.
diff --git a/lib/core/src/Cardano/Wallet/Primitive/Types/TokenMap.hs b/lib/core/src/Cardano/Wallet/Primitive/Types/TokenMap.hs
@@ -71,6 +71,7 @@ module Cardano.Wallet.Primitive.Types.TokenMap
     , maximumQuantity
 
     -- * Partitioning
+    , equipartitionAssets
     , equipartitionQuantities
     , equipartitionQuantitiesWithUpperBound
 
@@ -94,6 +95,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
@@ -132,6 +135,8 @@ import GHC.Generics
     ( Generic )
 import GHC.TypeLits
     ( ErrorMessage (..), TypeError )
+import Numeric.Natural
+    ( Natural )
 import Quiet
     ( Quiet (..) )
 
@@ -625,6 +630,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
@@ -20,6 +20,7 @@ import Cardano.Wallet.Primitive.Types.TokenMap
     ( AssetId (..), Flat (..), Nested (..), TokenMap )
 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,12 +86,16 @@ import Test.Hspec.Extra
     ( parallel )
 import Test.QuickCheck
     ( Arbitrary (..)
+    , Blind (..)
     , Fun
     , Property
     , applyFun
     , checkCoverage
+    , choose
     , cover
+    , frequency
     , property
+    , (.||.)
     , (===)
     , (==>)
     )
@@ -200,6 +207,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" $
@@ -467,6 +485,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
 --------------------------------------------------------------------------------
@@ -749,6 +808,10 @@ tokenPolicyIdHexStringLength = 56
 -- Arbitrary instances
 --------------------------------------------------------------------------------
 
+newtype Large a = Large
+    { getLarge :: a }
+    deriving (Eq, Show)
+
 newtype Positive a = Positive
     { getPositive :: a }
     deriving (Eq, Show)
@@ -773,6 +836,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