Add spendTx

- Using "applyTxToUTxO" then "filterByAddress" filters the entire UTxO on every
  application of a Tx. By presuming that we've already filtered the previous
  UTxO, then only applying parts of the new Tx that we care about, we can get
  better performance characteristics.
  - Add new "spendTx" function to support this. It does part of what
  "applyTxToUTxO" does.
  - Redefine "applyTxToUTxO" in terms of "spendTx"
  - Replace call to "applyTxToUTxO" with the equivalent, and more efficient:
    "spendTx tx u <> filterByAddress (utxoFromTx tx)"
- Add property tests to prove equivalences.
- Re-jig and clean up existing property tests a bit.

lib/core/src/Cardano/Wallet/Primitive/Model.hs

@@ -41,6 +41,7 @@ module Cardano.Wallet.Primitive.Model
     , unsafeInitWallet
     , applyTxToUTxO
     , utxoFromTx
+    , spendTx
 
     -- * Accessors
     , currentTip
@@ -333,28 +334,30 @@ totalUTxO pending wallet@(Wallet _ _ s) =
 --                              + balance (utxoFromTx tx)
 --                              - balance (u `restrictedBy` inputs tx)
 -- unUTxO (applyTxToUTxO tx u) = unUTxO u
---   `Map.union` unUTxO (utxoFromTx tx)
---   `Map.difference` unUTxO (u `restrictedBy` inputs tx)
+--     `Map.union` unUTxO (utxoFromTx tx)
+--     `Map.difference` unUTxO (u `restrictedBy` inputs tx)
+-- applyTxToUTxO tx u = spend tx u <> utxoFromTx tx
+-- applyTxToUTxO tx u = spend tx (u <> utxoFromTx tx)
 applyTxToUTxO
     :: Tx
     -> UTxO
     -> UTxO
-applyTxToUTxO tx !u = do
-    let
-        -- When a transaction comes in, we now know about a set of new unspents.
-        newKnown = u <> utxoFromTx tx
-        -- But we also know that some of the new unspents we know about, as well
-        -- as (potentially) some old ones, may now be spent.
-        spentKnown =
-            Set.fromList (inputs tx) `Set.intersection` dom newKnown
+applyTxToUTxO tx !u = spendTx tx u <> utxoFromTx tx
 
-    newKnown `excluding` spentKnown
+-- spendTx tx u `isSubsetOf` u
+-- balance (spendTx tx u) <= balance u
+-- balance (spendTx tx u) = balance u - balance (u `restrictedBy` inputs tx)
+-- spendTx tx u = u `excluding` inputs tx
+-- spendTx tx (filterByAddress f u) = filterByAddress f (spendTx tx u)
+spendTx :: Tx -> UTxO -> UTxO
+spendTx tx !u = u `excluding` Set.fromList (inputs tx)
 
 -- | Construct a UTxO corresponding to a given transaction. It is important for
 -- the transaction outputs to be ordered correctly, since they become available
 -- inputs for the subsequent blocks.
 --
 -- balance (utxoFromTx tx) = foldMap tokens (outputs tx)
+-- utxoFromTx tx `excluding` Set.fromList (inputs tx) = utxoFrom tx
 utxoFromTx :: Tx -> UTxO
 utxoFromTx Tx {txId, outputs} =
     UTxO $ Map.fromList $ zip (TxIn txId <$> [0..]) outputs
@@ -435,7 +438,8 @@ prefilterBlock b u0 = runState $ do
     applyTx (!txs, !prevUTxO) tx = do
         -- The next UTxO state (apply a state transition) (e.g. remove
         -- transaction outputs we've spent).
-        ourNextUTxO <- filterByAddressM isOurs' (applyTxToUTxO tx prevUTxO)
+        ourNextUTxO <-
+            (spendTx tx prevUTxO <>) <$> filterByAddressM isOurs' (utxoFromTx tx)
 
         ourWithdrawals <- Coin . sum . fmap (unCoin . snd) <$>
             mapMaybeM ourWithdrawal (Map.toList $ withdrawals tx)

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

@@ -148,15 +148,6 @@ size (UTxO u) = Map.size u
 --
 -- Returns the subset of UTxO entries that have addresses for which the given
 -- indicator function returns 'True'.
---
--- filterByAddressM (const $ pure True) u = u
--- filterByAddressM (const $ pure False) u = mempty
--- filterByAddressM f mempty = mempty
--- balance (filterByAddressM f (applyTxToUTxO tx mempty)) =
---   foldMap (\o -> do
---               ours <- f (address o)
---               if ours then tokens o else mempty
---            ) (outputs tx)
 filterByAddressM :: forall f. Monad f => (Address -> f Bool) -> UTxO -> f UTxO
 filterByAddressM isOursF (UTxO m) =
     UTxO <$> Map.traverseMaybeWithKey filterFunc m
@@ -166,6 +157,16 @@ filterByAddressM isOursF (UTxO m) =
             ours <- isOursF $ view #address txout
             pure $ if ours then Just txout else Nothing
 
+-- | Filters a 'UTxO' set with an indicator function on 'Address' values.
+--
+-- Returns the subset of UTxO entries that have addresses for which the given
+-- indicator function returns 'True'.
+--
+-- filterByAddress f u = runIdentity $ filterByAddressM (pure . f) u
+-- filterByAddress (const True) u = u
+-- filterByAddress (const False) u = mempty
+-- filterByAddress f mempty = mempty
+-- filterByAddress f u `isSubsetOf` u
 filterByAddress :: (Address -> Bool) -> UTxO -> UTxO
 filterByAddress f = runIdentity . filterByAddressM (pure . f)
 

lib/core/test/unit/Cardano/Wallet/Primitive/ModelSpec.hs

@@ -34,6 +34,7 @@ import Cardano.Wallet.Primitive.Model
     , currentTip
     , getState
     , initWallet
+    , spendTx
     , totalBalance
     , totalUTxO
     , unsafeInitWallet
@@ -100,7 +101,7 @@ import Data.List
 import Data.List.NonEmpty
     ( NonEmpty (..) )
 import Data.Maybe
-    ( catMaybes )
+    ( catMaybes, isJust )
 import Data.Quantity
     ( Quantity (..) )
 import Data.Set
@@ -190,9 +191,19 @@ spec = do
                 (property prop_applyTxToUTxO_entries)
             it "applyTxToUTxO then filterByAddress"
                 (property prop_filterByAddress_balance_applyTxToUTxO)
+            it "spendTx/utxoFromTx"
+                (property prop_applyTxToUTxO_spendTx_utxoFromTx)
 
-        describe "utxoFromTx" $
+        describe "utxoFromTx" $ do
             it "has expected balance" (property prop_utxoFromTx_balance)
+            it "is unspent" (property prop_utxoFromTx_is_unspent)
+
+        describe "spendTx" $ do
+            it "is subset of UTxO" (property prop_spendTx_isSubset)
+            it "balance is <= balance of UTxO" (property prop_spendTx_balance_inequality)
+            it "has expected balance" (property prop_spendTx_balance)
+            it "definition" (property prop_spendTx)
+            it "commutative with filterByAddress" (property prop_spendTx_filterByAddress)
 
     parallel $ describe "Available UTxO" $ do
         it "prop_availableUTxO_isSubmap" $
@@ -1399,3 +1410,67 @@ prop_utxoFromTx_balance :: Property
 prop_utxoFromTx_balance =
     forAllShrink genTx shrinkTx $ \tx ->
         balance (utxoFromTx tx) === foldMap tokens (outputs tx)
+
+prop_utxoFromTx_is_unspent :: Property
+prop_utxoFromTx_is_unspent =
+    forAllShrink genTx shrinkTx $ \tx ->
+        utxoFromTx tx `excluding` Set.fromList (inputs tx)
+        === utxoFromTx tx
+
+-- spendTx tx u `isSubsetOf` u
+prop_spendTx_isSubset :: Property
+prop_spendTx_isSubset =
+    forAllShrink genTx shrinkTx $ \tx ->
+    forAllShrink genUTxO shrinkUTxO $ \u ->
+        spendTx tx u `UTxO.isSubsetOf` u
+
+-- balance (spendTx tx u) <= balance u
+prop_spendTx_balance_inequality :: Property
+prop_spendTx_balance_inequality =
+    forAllShrink genTx shrinkTx $ \tx ->
+    forAllShrink genUTxO shrinkUTxO $ \u ->
+        let
+            lhs = balance (spendTx tx u)
+            rhs = balance u
+        in
+            isJust (rhs `TokenBundle.subtract` lhs)
+            & counterexample ("balance (spendTx tx u) = " <> show lhs)
+            & counterexample ("balance u = " <> show rhs)
+
+prop_spendTx_balance :: Property
+prop_spendTx_balance =
+    forAllShrink genTx shrinkTx $ \tx ->
+    forAllShrink genUTxO shrinkUTxO $ \u ->
+        let
+            lhs = balance (spendTx tx u)
+            rhs =
+                balance u
+                    `TokenBundle.unsafeSubtract`
+                        balance (u `UTxO.restrictedBy` Set.fromList (inputs tx))
+        in
+            lhs === rhs
+
+prop_spendTx :: Property
+prop_spendTx =
+    forAllShrink genTx shrinkTx $ \tx ->
+    forAllShrink genUTxO shrinkUTxO $ \u ->
+        spendTx tx u === u `excluding` Set.fromList (inputs tx)
+
+prop_applyTxToUTxO_spendTx_utxoFromTx :: Property
+prop_applyTxToUTxO_spendTx_utxoFromTx =
+    forAllShrink genTx shrinkTx $ \tx ->
+    forAllShrink genUTxO shrinkUTxO $ \u ->
+        conjoin
+        [ applyTxToUTxO tx u === spendTx tx u <> utxoFromTx tx
+        , applyTxToUTxO tx u === spendTx tx (u <> utxoFromTx tx)
+        ]
+
+prop_spendTx_filterByAddress :: Bool -> Property
+prop_spendTx_filterByAddress b =
+    let
+        f = const b
+    in
+        forAllShrink genTx shrinkTx $ \tx ->
+        forAllShrink genUTxO shrinkUTxO $ \u ->
+            filterByAddress f (spendTx tx u)
+            === spendTx tx (filterByAddress f u)

lib/core/test/unit/Cardano/Wallet/Primitive/Types/UTxOSpec.hs

@@ -9,7 +9,7 @@ import Prelude
 import Cardano.Wallet.Primitive.Types.Address.Gen
     ( Parity (..), addressParity )
 import Cardano.Wallet.Primitive.Types.UTxO
-    ( UTxO (..), dom, filterByAddress, filterByAddressM )
+    ( UTxO (..), dom, filterByAddress, filterByAddressM, isSubsetOf )
 import Cardano.Wallet.Primitive.Types.UTxO.Gen
     ( genUTxO, shrinkUTxO )
 import Data.Functor.Identity
@@ -39,6 +39,8 @@ spec =
             property prop_filterByAddress_empty
         it "filterByAddress/filterByAddressM" $
             property prop_filterByAddress_filterByAddressM
+        it "filterByAddress is always subset" $
+            property prop_filterByAddress_isSubset
 
 prop_filterByAddress_matchAll :: Property
 prop_filterByAddress_matchAll =
@@ -90,3 +92,11 @@ prop_filterByAddress_filterByAddressM b =
     in
         forAllShrink genUTxO shrinkUTxO $ \u ->
             filterByAddress f u === runIdentity (filterByAddressM (pure . f) u)
+
+prop_filterByAddress_isSubset :: Bool -> Property
+prop_filterByAddress_isSubset b =
+    let
+        f = const b
+    in
+        forAllShrink genUTxO shrinkUTxO $ \u ->
+            filterByAddress f u `isSubsetOf` u