Merge #2848

2848: Move important UTxO state transition functions to the top-level and test them r=jonathanknowles a=sevanspowell

This work is intended to allow future modifications to the UTxO state transition function (i.e. collateral) to be tested in isolation, at the unit level.

My changes are primarily to the nested function "applyTx". I have tried to keep my changes minimal:
  - Use more verbose variable names in that function to make the purpose of each variable clear.
  - More clearly delineate between transaction outputs that are unspent and transaction outputs "we know about" but may possibly be spent.
  - Make some effort to separate the ideas of "knowledge of a UTxO" from "ownership of a UTxO" to make testing easier.
  - Extract out the core state transition function from the existing "do lots" "applyTx" function and move it to the top-level.
  - Move other useful functions to the top-level for the same reason.
  - Test important functions in wallet model spec.
  - Re-implement hasKnownInput/hasKnownOutput to make their implementations a little more clear and symmetric.

### Comments

These changes make testing collateral much easier, we can just test the "applyTxToUTxO" function at the unit level, and adjust existing model tests as necessary.

### Issue Number

ADP-1092


Co-authored-by: Samuel Evans-Powell <[email protected]>
Co-authored-by: Jonathan Knowles <[email protected]>

lib/core/cardano-wallet-core.cabal

@@ -397,13 +397,15 @@ test-suite unit
       Cardano.Wallet.Primitive.Slotting.Legacy
       Cardano.Wallet.Primitive.SlottingSpec
       Cardano.Wallet.Primitive.SyncProgressSpec
+      Cardano.Wallet.Primitive.Types.AddressSpec
       Cardano.Wallet.Primitive.Types.CoinSpec
       Cardano.Wallet.Primitive.Types.HashSpec
       Cardano.Wallet.Primitive.Types.TokenBundleSpec
       Cardano.Wallet.Primitive.Types.TokenMapSpec
       Cardano.Wallet.Primitive.Types.TokenMapSpec.TypeErrorSpec
       Cardano.Wallet.Primitive.Types.TokenPolicySpec
       Cardano.Wallet.Primitive.Types.TokenQuantitySpec
+      Cardano.Wallet.Primitive.Types.UTxOSpec
       Cardano.Wallet.Primitive.Types.UTxOIndexSpec
       Cardano.Wallet.Primitive.Types.UTxOIndex.TypeErrorSpec
       Cardano.Wallet.Primitive.TypesSpec

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

@@ -39,6 +39,9 @@ module Cardano.Wallet.Primitive.Model
     , applyBlock
     , applyBlocks
     , unsafeInitWallet
+    , applyTxToUTxO
+    , utxoFromTx
+    , spendTx
 
     -- * Accessors
     , currentTip
@@ -74,21 +77,20 @@ import Cardano.Wallet.Primitive.Types.Tx
     , Tx (..)
     , TxIn (..)
     , TxMeta (..)
-    , TxOut (..)
     , TxStatus (..)
     , inputs
     , txOutCoin
     )
 import Cardano.Wallet.Primitive.Types.UTxO
-    ( Dom (..), UTxO (..), balance, excluding, restrictedBy )
+    ( Dom (..), UTxO (..), balance, excluding, filterByAddressM )
 import Control.DeepSeq
     ( NFData (..), deepseq )
 import Control.Monad
-    ( foldM, forM )
+    ( foldM )
 import Control.Monad.Extra
     ( mapMaybeM )
 import Control.Monad.Trans.State.Strict
-    ( State, evalState, runState, state )
+    ( State, StateT, evalState, runState, state )
 import Data.Functor
     ( (<&>) )
 import Data.Generics.Internal.VL.Lens
@@ -98,7 +100,7 @@ import Data.Generics.Labels
 import Data.List.NonEmpty
     ( NonEmpty (..) )
 import Data.Maybe
-    ( catMaybes, isJust )
+    ( isJust )
 import Data.Set
     ( Set )
 import Fmt
@@ -340,6 +342,56 @@ totalUTxO pending (Wallet u _ s) =
     entriesToExcludeForTx :: Tx -> Set TxIn
     entriesToExcludeForTx tx = Set.fromList $ fst <$> tx ^. #resolvedInputs
 
+-- | Applies a transaction to a UTxO, moving it from one state from another.
+-- When applying a transaction to a UTxO:
+--   1. We need to remove any unspents that have been spent in the transaction.
+--   2. Add any unspents that we've received via the transaction.
+--
+-- We don't consider "ownership" here (is this address ours?), only "do we
+-- know about this address" (i.e. is it present in our UTxO?).
+--
+-- balance (applyTxToUTxO tx u) = balance u
+--                              + 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)
+-- applyTxToUTxO tx u = spend tx u <> utxoFromTx tx
+-- applyTxToUTxO tx u = spend tx (u <> utxoFromTx tx)
+applyTxToUTxO
+    :: Tx
+    -> UTxO
+    -> UTxO
+applyTxToUTxO tx !u = spendTx tx u <> utxoFromTx tx
+
+-- | Remove unspents that have been consumed by the given transaction.
+--
+-- 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 (u <> utxoFromTx tx) = spendTx tx u <> utxoFromTx tx
+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
+
+isOurAddress
+    :: forall s m
+     . (Monad m, IsOurs s Address)
+    => Address
+    -> StateT s m Bool
+isOurAddress = fmap isJust . state . isOurs
+
 {-------------------------------------------------------------------------------
                                Internals
 -------------------------------------------------------------------------------}
@@ -406,16 +458,48 @@ prefilterBlock b u0 = runState $ do
         => ([(Tx, TxMeta)], UTxO)
         -> Tx
         -> State s ([(Tx, TxMeta)], UTxO)
-    applyTx (!txs, !u) tx = do
-        ourU <- state $ utxoOurs tx
-        let ourIns = Set.fromList (inputs tx) `Set.intersection` dom (u <> ourU)
-        let u' = (u <> ourU) `excluding` ourIns
+    applyTx (!txs, !prevUTxO) tx = do
+        -- The next UTxO state (apply a state transition) (e.g. remove
+        -- transaction outputs we've spent).
+        ourNextUTxO <-
+            (spendTx tx prevUTxO <>)
+            <$> filterByAddressM isOurAddress (utxoFromTx tx)
+
         ourWithdrawals <- Coin . sum . fmap (unCoin . snd) <$>
             mapMaybeM ourWithdrawal (Map.toList $ withdrawals tx)
-        let received = balance ourU
-        let spent = balance (u `restrictedBy` ourIns) `TB.add` TB.fromCoin ourWithdrawals
-        let hasKnownInput = ourIns /= mempty
-        let hasKnownOutput = ourU /= mempty
+
+        let received = balance (ourNextUTxO `excluding` dom prevUTxO)
+        let spent =
+                balance (prevUTxO `excluding` dom ourNextUTxO)
+                `TB.add` TB.fromCoin ourWithdrawals
+
+        (ownedAndKnownTxIns, ownedAndKnownTxOuts) <- do
+            -- A new transaction expands the set of transaction inputs/outputs
+            -- we know about, but not all those transaction inputs/outputs
+            -- belong to us, so we filter any new inputs/outputs, presuming that
+            -- the previous UTxO has already been filtered:
+            ownedAndKnown <-
+                (prevUTxO <>) <$> filterByAddressM isOurAddress (utxoFromTx tx)
+            -- Also, the new transaction may spend some transaction
+            -- inputs/outputs. But we don't want to apply that logic yet. If we
+            -- do, any spent transaction input/output will be removed from our
+            -- knowledge base.
+            -- Therefore, because this is not technically an "Unspent TxO" set,
+            -- let's just return the TxIns and TxOuts, as the type "UTxO" will
+            -- create expectations which we explicitly aren't fulfilling:
+            let m = unUTxO ownedAndKnown
+            pure (Map.keys m, Map.elems m)
+
+        -- A transaction has a known input if one of the transaction inputs
+        -- matches a transaction input we know about.
+        let hasKnownInput = not $ Set.disjoint
+                (Set.fromList $ inputs tx)
+                (Set.fromList ownedAndKnownTxIns)
+        -- A transaction has a known output if one of the transaction outputs
+        -- matches a transaction output we know about.
+        let hasKnownOutput = not $ Set.disjoint
+                (Set.fromList $ outputs tx)
+                (Set.fromList ownedAndKnownTxOuts)
         let hasKnownWithdrawal = ourWithdrawals /= mempty
 
         -- NOTE 1: The only case where fees can be 'Nothing' is when dealing with
@@ -443,7 +527,7 @@ prefilterBlock b u0 = runState $ do
         return $ if hasKnownOutput && not hasKnownInput then
             let dir = Incoming in
             ( (tx { fee = actualFee dir }, mkTxMeta (TB.getCoin received) dir) : txs
-            , u'
+            , ourNextUTxO
             )
         else if hasKnownInput || hasKnownWithdrawal then
             let
@@ -453,10 +537,10 @@ prefilterBlock b u0 = runState $ do
                 amount = distance adaSpent adaReceived
             in
                 ( (tx { fee = actualFee dir }, mkTxMeta amount dir) : txs
-                , u'
+                , ourNextUTxO
                 )
         else
-            (txs, u)
+            (txs, prevUTxO)
 
 -- | Get the change UTxO
 --
@@ -471,21 +555,5 @@ changeUTxO
     -> s
     -> UTxO
 changeUTxO pending = evalState $
-    mconcat <$> mapM (state . utxoOurs) (Set.toList pending)
-
--- | Construct our _next_ UTxO (possible empty) from a transaction by selecting
--- outputs that are ours. It is important for the transaction outputs to be
--- ordered correctly, since they become available inputs for the subsequent
--- blocks.
-utxoOurs
-    :: IsOurs s Address
-    => Tx
-    -> s
-    -> (UTxO, s)
-utxoOurs tx = runState $ toUtxo <$> forM (zip [0..] (outputs tx)) filterOut
-  where
-    toUtxo = UTxO . Map.fromList . catMaybes
-    filterOut (ix, out) = do
-        state (isOurs $ address out) <&> \case
-            Just{}  -> Just (TxIn (txId tx) ix, out)
-            Nothing -> Nothing
+    mconcat
+    <$> mapM (filterByAddressM isOurAddress . utxoFromTx) (Set.toList pending)

lib/core/src/Cardano/Wallet/Primitive/Types/Address/Gen.hs

@@ -1,6 +1,13 @@
 module Cardano.Wallet.Primitive.Types.Address.Gen
-    ( genAddress
+    (
+      -- * Generators and shrinkers
+      genAddress
     , shrinkAddress
+    , coarbitraryAddress
+
+      -- * Indicator functions on addresses
+    , addressParity
+    , Parity (..)
     )
     where
 
@@ -9,8 +16,10 @@ import Prelude
 import Cardano.Wallet.Primitive.Types.Address
     ( Address (..) )
 import Test.QuickCheck
-    ( Gen, elements, sized )
+    ( Gen, coarbitrary, elements, sized )
 
+import qualified Data.Bits as Bits
+import qualified Data.ByteString as BS
 import qualified Data.ByteString.Char8 as B8
 
 --------------------------------------------------------------------------------
@@ -27,9 +36,53 @@ shrinkAddress a
   where
     simplest = head addresses
 
+coarbitraryAddress :: Address -> Gen a -> Gen a
+coarbitraryAddress = coarbitrary . BS.unpack . unAddress
+
 addresses :: [Address]
 addresses = mkAddress <$> ['0' ..]
 
+--------------------------------------------------------------------------------
+-- Indicator functions on addresses
+--------------------------------------------------------------------------------
+
+-- | Computes the parity of an address.
+--
+-- Parity is defined in the following way:
+--
+--    - even-parity address:
+--      an address with a pop count (Hamming weight) that is even.
+--
+--    - odd-parity address:
+--      an address with a pop count (Hamming weight) that is odd.
+--
+-- Examples of even-parity and odd-parity addresses:
+--
+--    - 0b00000000 : even (Hamming weight = 0)
+--    - 0b00000001 : odd  (Hamming weight = 1)
+--    - 0b00000010 : odd  (Hamming weight = 1)
+--    - 0b00000011 : even (Hamming weight = 2)
+--    - 0b00000100 : odd  (Hamming weight = 1)
+--    - ...
+--    - 0b11111110 : odd  (Hamming weight = 7)
+--    - 0b11111111 : even (Hamming weight = 8)
+--
+addressParity :: Address -> Parity
+addressParity = parity . addressPopCount
+  where
+    addressPopCount :: Address -> Int
+    addressPopCount = BS.foldl' (\acc -> (acc +) . Bits.popCount) 0 . unAddress
+
+    parity :: Integral a => a -> Parity
+    parity a
+        | even a    = Even
+        | otherwise = Odd
+
+-- | Represents the parity of a value (whether the value is even or odd).
+--
+data Parity = Even | Odd
+    deriving (Eq, Show)
+
 --------------------------------------------------------------------------------
 -- Internal utilities
 --------------------------------------------------------------------------------

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

@@ -59,6 +59,7 @@ import Test.QuickCheck
     , liftArbitrary2
     , liftShrink
     , liftShrink2
+    , listOf
     , listOf1
     , shrinkList
     , shrinkMapBy
@@ -102,7 +103,7 @@ genTxWithoutId = TxWithoutId
     <$> liftArbitrary genCoinPositive
     <*> listOf1 (liftArbitrary2 genTxIn genCoinPositive)
     <*> listOf1 (liftArbitrary2 genTxIn genCoinPositive)
-    <*> listOf1 genTxOut
+    <*> listOf genTxOut
     <*> liftArbitrary genNestedTxMetadata
     <*> genMapWith genRewardAccount genCoinPositive
 

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

@@ -3,6 +3,8 @@
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE OverloadedLabels #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeFamilies #-}
 
@@ -34,11 +36,15 @@ module Cardano.Wallet.Primitive.Types.UTxO
     , restrictedBy
     , restrictedTo
     , size
+    , filterByAddressM
+    , filterByAddress
     ) where
 
 import Prelude hiding
     ( null )
 
+import Cardano.Wallet.Primitive.Types.Address
+    ( Address )
 import Cardano.Wallet.Primitive.Types.Coin
     ( Coin (..) )
 import Cardano.Wallet.Primitive.Types.TokenBundle
@@ -49,6 +55,8 @@ import Control.DeepSeq
     ( NFData (..) )
 import Data.Bifunctor
     ( first )
+import Data.Functor.Identity
+    ( runIdentity )
 import Data.Generics.Internal.VL.Lens
     ( view )
 import Data.Kind
@@ -136,6 +144,32 @@ null (UTxO u) = Map.null u
 size :: UTxO -> Int
 size (UTxO u) = Map.size u
 
+-- | 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'.
+filterByAddressM :: forall f. Monad f => (Address -> f Bool) -> UTxO -> f UTxO
+filterByAddressM isOursF (UTxO m) =
+    UTxO <$> Map.traverseMaybeWithKey filterFunc m
+  where
+    filterFunc :: TxIn -> TxOut -> f (Maybe TxOut)
+    filterFunc _txin txout = do
+        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)
+
 data UTxOStatistics = UTxOStatistics
     { histogram :: ![HistogramBar]
     , allStakes :: !Word64