Move important UTxO state transition functions to the top-level and test them

lib/core/cardano-wallet-core.cabal

@@ -396,13 +396,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
@@ -320,6 +322,56 @@ totalUTxO
 totalUTxO pending wallet@(Wallet _ _ s) =
     availableUTxO pending wallet <> changeUTxO pending s
 
+-- | 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 giventransaction.
+--
+-- 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
 -------------------------------------------------------------------------------}
@@ -386,16 +438,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
@@ -423,7 +507,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
@@ -433,10 +517,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
 --
@@ -451,21 +535,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,12 @@
 module Cardano.Wallet.Primitive.Types.Address.Gen
-    ( genAddress
+    (
+      -- * Generators and shrinkers
+      genAddress
     , shrinkAddress
+
+      -- * Indicator functions on addresses
+    , addressParity
+    , Parity (..)
     )
     where
 
@@ -11,6 +17,8 @@ import Cardano.Wallet.Primitive.Types.Address
 import Test.QuickCheck
     ( Gen, elements, sized )
 
+import qualified Data.Bits as Bits
+import qualified Data.ByteString as BS
 import qualified Data.ByteString.Char8 as B8
 
 --------------------------------------------------------------------------------
@@ -30,6 +38,47 @@ shrinkAddress a
 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/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