diffAssets tests

Some improvements to the diffAssets

Author: v0d1ch

hydra-cluster/src/Hydra/Cluster/Scenarios.hs

@@ -115,7 +115,7 @@ import Hydra.Node.DepositPeriod (DepositPeriod (..))
 import Hydra.Options (CardanoChainConfig (..), ChainBackendOptions (..), DirectOptions (..), RunOptions (..), startChainFrom)
 import Hydra.Tx (HeadId, IsTx (balance), Party, txId)
 import Hydra.Tx.ContestationPeriod qualified as CP
-import Hydra.Tx.Deposit (capUTxO, filterAssets)
+import Hydra.Tx.Deposit (capUTxO, diffAssets)
 import Hydra.Tx.Utils (dummyValidatorScript, verificationKeyToOnChainId)
 import HydraNode (
   HydraClient (..),
@@ -1347,7 +1347,8 @@ canDepositPartially tracer workDir blockTime backend hydraScriptsTxId =
           let tokenAssetValue = assetsToValue tokenAssets
           let quantityMoreThan20 = (> 20)
           let partialTokenAssets = Map.map (\(CAPI.PolicyAssets policyAssetMap) -> CAPI.PolicyAssets $ Map.filter quantityMoreThan20 policyAssetMap) tokenAssets
-          let tokenDiff = filterAssets tokenAssets (Map.toList partialTokenAssets)
+          -- NOTE: using diffAssets here implicitly tests that function too.
+          let tokenDiff = assetsToValue $ Map.fromList $ diffAssets tokenAssets partialTokenAssets
           let partialTokenAssetValue = assetsToValue partialTokenAssets
           let tokenAssetValueWithoutAda = assetsToValue $ valueToPolicyAssets partialTokenAssetValue
           let seedAmount = 5_000_000

hydra-tx/src/Hydra/Tx/Deposit.hs

@@ -59,7 +59,7 @@ depositTx networkId headId commitBlueprintTx upperSlot deadline amount tokens =
   utxoToDeposit = mergeUTxO utxoToDeposit' tokensToDepositUTxO
 
   returnToUser =
-    let returnToUserUTxO = leftoverUTxO `filterExistingAssets` tokensToDepositUTxO
+    let returnToUserUTxO = leftoverUTxO `diffExistingAssets` tokensToDepositUTxO
      in if UTxO.null returnToUserUTxO
           then StrictSeq.empty
           else
@@ -72,37 +72,54 @@ depositTx networkId headId commitBlueprintTx upperSlot deadline amount tokens =
 
   depositInputs = (,BuildTxWith $ KeyWitness KeyWitnessForSpending) <$> depositInputsList
 
--- | Filter the first argument UTxO's non ADA assets in case any asset exists in the second UTxO argument.
--- Asset quantities will be subtracted if they are found.
-filterExistingAssets :: UTxO -> UTxO -> UTxO
-filterExistingAssets utxoToFilter utxoToLookup =
+-- | Find the difference between the first argument UTxO's non ADA assets in
+-- and the second UTxO argument. Matching asset quantities will be subtracted
+-- if they are found.
+diffExistingAssets :: UTxO -> UTxO -> UTxO
+diffExistingAssets utxoToFilter utxoToLookup =
   UTxO.fromList $ findAssets <$> UTxO.toList utxoToFilter
  where
   findAssets (i, TxOut a val d r) =
     let assets = valueToPolicyAssets val
         originalLovelace = selectLovelace val
+        filteredAssets' = diffAssets assets forLookup
         filteredAssets =
-         foldMap (uncurry policyAssetsToValue) $
-          filterAssets assets forLookup
-     in (i, TxOut a (lovelaceToValue originalLovelace <> filteredAssets) d r)
-  forLookup =
-     -- NOTE: Uses a list to store all policies, preserving multiple entries
-     -- with the same policyId but different assets. A Map would silently
-     -- overwrite duplicates.
-     concatMap (Map.toList . valueToPolicyAssets . txOutValue . snd) $ UTxO.toList utxoToLookup
-
- -- | Filter the first argument map of assets in case any asset exists in the second argument.
-filterAssets :: Map PolicyId PolicyAssets -> [(PolicyId, PolicyAssets)] -> [(PolicyId, PolicyAssets)]
-filterAssets assets forLookup =
-     Map.assocs $
-       Map.mapWithKey
-         ( \pid (PolicyAssets x) ->
-             let samePid = List.filter (\(pid', _) -> pid' == pid) forLookup
-              in case List.lookup pid samePid of
-                   Nothing -> PolicyAssets x
-                   Just (PolicyAssets foundAssets) -> PolicyAssets $ x `Map.difference` foundAssets
-         )
-         assets
+          if null (snd <$> filteredAssets')
+            then mempty
+            else filteredAssets'
+        filteredValue =
+          foldMap (uncurry policyAssetsToValue) filteredAssets
+     in (i, TxOut a (lovelaceToValue originalLovelace <> filteredValue) d r)
+  forLookup = valueToPolicyAssets $ UTxO.totalValue utxoToLookup
+
+-- | Diff the first argument map of assets in case any asset exists in the second argument.
+diffAssets :: Map PolicyId PolicyAssets -> Map PolicyId PolicyAssets -> [(PolicyId, PolicyAssets)]
+diffAssets assets forLookup =
+  if null forLookup
+    then Map.toList assets
+    else
+      Map.assocs $
+        Map.foldrWithKey
+          ( \pid (PolicyAssets existing) result ->
+              case Map.lookup pid forLookup of
+                Nothing -> result
+                Just foundAsset -> result `Map.union` Map.singleton pid (PolicyAssets $ go existing foundAsset)
+          )
+          Map.empty
+          assets
+ where
+  go :: Map AssetName Quantity -> PolicyAssets -> Map AssetName Quantity
+  go existing (PolicyAssets found) =
+    Map.differenceWith
+      checkQuantities
+      existing
+      found
+
+  checkQuantities :: Quantity -> Quantity -> Maybe Quantity
+  checkQuantities existing wanted =
+    if existing > wanted
+      then Just $ existing - wanted
+      else Nothing
 
 -- | Merges the two 'UTxO' favoring data coming from the first argument 'UTxO'.
 -- In case the same 'TxIn' was found in the first 'UTxO' - second 'UTxO' value

hydra-tx/test/Hydra/Tx/DepositSpec.hs

@@ -7,7 +7,7 @@ import Cardano.Api.UTxO qualified as UTxO
 import Data.Map.Strict qualified as Map
 import Data.Set qualified as Set
 import Hydra.Cardano.Api (AssetId (..), AssetName, Coin (..), PolicyAssets (..), PolicyId, Quantity (..), UTxO, selectLovelace, txOutValue, valueToPolicyAssets)
-import Hydra.Tx.Deposit (capUTxO, pickTokensToDeposit, splitTokens)
+import Hydra.Tx.Deposit (capUTxO, diffAssets, pickTokensToDeposit, splitTokens)
 import Test.Hydra.Tx.Fixture (testPolicyId)
 import Test.Hydra.Tx.Gen (genUTxOSized, genUTxOWithAssetsSized)
 import Test.QuickCheck (Property, chooseInteger, counterexample, cover, elements, forAll, frequency, listOf, oneof, property, (===), (==>))
@@ -240,6 +240,77 @@ spec =
         prop "monotonic with respect to target" propMonotonicTarget
         prop "no UTxO loss" propNoUTxOLoss
 
+    describe "diffAssets" $ do
+      describe "tests" $ do
+        it "returns empty assets when both inputs are empty" $ do
+          let assets = diffAssets mempty mempty
+          assets `shouldBe` mempty
+        it "returns existing assets when lookup input is empty" $
+          forAll arbitrary $ \assetMap -> do
+            let assets = diffAssets assetMap mempty
+            assets `shouldBe` Map.toList assetMap
+        it "returns empty assets if asset input is empty" $
+          forAll arbitrary $ \lookupMap -> do
+            let assets = diffAssets mempty lookupMap
+            assets `shouldBe` mempty
+        it "subracts found values" $
+          forAll arbitrary $ \aPolicy -> do
+            let policyAssets =
+                  Map.fromList [("SomeTokenA", 100), ("SomeTokenB", 200)]
+            let policyAssets' =
+                  Map.fromList [("SomeTokenA", 30), ("SomeTokenB", 50)]
+            let expectedAssets =
+                  Map.fromList [("SomeTokenA", 70), ("SomeTokenB", 150)]
+            let expectedResult =
+                  Map.fromList [(aPolicy, PolicyAssets expectedAssets)]
+
+            let a = Map.fromList [(aPolicy, PolicyAssets policyAssets)]
+            let b = Map.fromList [(aPolicy, PolicyAssets policyAssets')]
+            let assets = diffAssets a b
+            assets `shouldBe` Map.toList expectedResult
+        it "keeps assets not found in the lookup map" $
+          forAll arbitrary $ \aPolicy -> do
+            let policyAssets =
+                  Map.fromList [("SomeTokenA", 100), ("SomeTokenC", 2)]
+            let policyAssets' =
+                  Map.fromList [("SomeTokenA", 30), ("SomeTokenB", 50)]
+            let expectedAssets =
+                  Map.fromList [("SomeTokenA", 70), ("SomeTokenC", 2)]
+            let expectedResult =
+                  Map.fromList [(aPolicy, PolicyAssets expectedAssets)]
+
+            let a = Map.fromList [(aPolicy, PolicyAssets policyAssets)]
+            let b = Map.fromList [(aPolicy, PolicyAssets policyAssets')]
+            let assets = diffAssets a b
+            assets `shouldBe` Map.toList expectedResult
+        it "ignores extra assets in the lookup map" $
+          forAll arbitrary $ \aPolicy -> do
+            let policyAssets =
+                  Map.fromList [("SomeTokenA", 100), ("SomeTokenB", 85)]
+            let policyAssets' =
+                  Map.fromList [("SomeTokenA", 30), ("SomeTokenB", 50), ("SomeTokenC", 400)]
+            let expectedAssets =
+                  Map.fromList [("SomeTokenA", 70), ("SomeTokenB", 35)]
+            let expectedResult =
+                  Map.fromList [(aPolicy, PolicyAssets expectedAssets)]
+
+            let a = Map.fromList [(aPolicy, PolicyAssets policyAssets)]
+            let b = Map.fromList [(aPolicy, PolicyAssets policyAssets')]
+            let assets = diffAssets a b
+            assets `shouldBe` Map.toList expectedResult
+        it "ignores assets with too low values" $
+          forAll arbitrary $ \aPolicy -> do
+            let policyAssets =
+                  Map.fromList [("SomeTokenA", 100), ("SomeTokenB", 85)]
+            let policyAssets' =
+                  Map.fromList [("SomeTokenA", 100), ("SomeTokenB", 86)]
+            let expectedResult = [(aPolicy, mempty)]
+
+            let a = Map.fromList [(aPolicy, PolicyAssets policyAssets)]
+            let b = Map.fromList [(aPolicy, PolicyAssets policyAssets')]
+            let assets = diffAssets a b
+            assets `shouldBe` expectedResult
+
 -- | Property: The sum of selected and leftover values equals the input value
 propPreservesTotalValue :: UTxO -> Coin -> Property
 propPreservesTotalValue utxo target =