Make it possible to filter a UTxO set based on a given asset.

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

@@ -37,6 +37,7 @@ module Cardano.Wallet.Primitive.Types.Tx
     , txIns
     , txMetadataIsNull
     , txOutCoin
+    , txOutHasAsset
 
     ) where
 
@@ -57,7 +58,7 @@ import Cardano.Wallet.Primitive.Types.Hash
 import Cardano.Wallet.Primitive.Types.RewardAccount
     ( RewardAccount (..) )
 import Cardano.Wallet.Primitive.Types.TokenBundle
-    ( TokenBundle )
+    ( AssetId, TokenBundle )
 import Cardano.Wallet.Primitive.Types.TokenPolicy
     ( TokenName, TokenPolicyId )
 import Cardano.Wallet.Primitive.Types.TokenQuantity
@@ -200,6 +201,12 @@ data TxOut = TxOut
 txOutCoin :: TxOut -> Coin
 txOutCoin = TB.getCoin . view #tokens
 
+-- Returns true if (and only if) the given transaction output has a non-zero
+-- quantity of the given asset.
+--
+txOutHasAsset :: TxOut -> AssetId -> Bool
+txOutHasAsset (TxOut _ b) = TB.hasQuantity b
+
 -- Since the 'TokenBundle' type deliberately does not provide an 'Ord' instance
 -- (as that would lead to arithmetically invalid orderings), this means we can't
 -- automatically derive an 'Ord' instance for the 'TxOut' type.

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

@@ -27,6 +27,8 @@ module Cardano.Wallet.Primitive.Types.UTxO
     , computeStatistics
     , computeUtxoStatistics
     , excluding
+    , filterForAsset
+    , partitionForAsset
     , isSubsetOf
     , log10
     , pickRandom
@@ -39,10 +41,14 @@ import Prelude
 
 import Cardano.Wallet.Primitive.Types.Coin
     ( Coin (..) )
+import Cardano.Wallet.Primitive.Types.TokenBundle
+    ( AssetId )
 import Cardano.Wallet.Primitive.Types.Tx
-    ( TxIn, TxOut (..), txOutCoin )
+    ( TxIn, TxOut (..), txOutCoin, txOutHasAsset )
 import Control.DeepSeq
     ( NFData (..) )
+import Data.Bifunctor
+    ( bimap )
 import Data.List
     ( foldl' )
 import Data.List.NonEmpty
@@ -103,6 +109,22 @@ pickRandom (UTxO utxo)
         ix <- randomRIO (0, toEnum (Map.size utxo - 1))
         return (Just $ Map.elemAt ix utxo, UTxO $ Map.deleteAt ix utxo)
 
+-- | Filter the given UTxO set based on the given asset.
+--
+-- Every entry in the returned set has a non-zero quantity of the asset.
+--
+filterForAsset :: UTxO -> AssetId -> UTxO
+filterForAsset (UTxO m) a = UTxO $ Map.filter (`txOutHasAsset` a) m
+
+-- | Partition the given UTxO set based on the given asset.
+--
+-- All entries in the first set have a non-zero quantity of the asset.
+-- No entries in the second set have a non-zero quantity of the asset.
+--
+partitionForAsset :: UTxO -> AssetId -> (UTxO, UTxO)
+partitionForAsset (UTxO m) a = bimap UTxO UTxO $
+    Map.partition (`txOutHasAsset` a) m
+
 -- | Compute the balance of a UTxO
 balance :: UTxO -> Natural
 balance =