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

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

@@ -39,6 +39,8 @@ module Cardano.Wallet.Primitive.Model
     , applyBlock
     , applyBlocks
     , unsafeInitWallet
+    , applyTxToUTxO
+    , utxoFromTx
 
     -- * Accessors
     , currentTip
@@ -74,21 +76,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 +99,7 @@ import Data.Generics.Labels
 import Data.List.NonEmpty
     ( NonEmpty (..) )
 import Data.Maybe
-    ( catMaybes, isJust )
+    ( isJust )
 import Data.Set
     ( Set )
 import Fmt
@@ -320,6 +321,51 @@ 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
+    -> 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
+
+    newKnown `excluding` spentKnown
+
+-- | 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 -> UTxO
+utxoFromTx Tx {txId, outputs} =
+    UTxO $ Map.fromList $ zip (TxIn txId <$> [0..]) outputs
+
+isOurs' :: forall s m. (Monad m, IsOurs s Address) => Address -> StateT s m Bool
+isOurs' addr = do
+    mDerivationIndex <- state $ isOurs addr
+    case mDerivationIndex of
+        Nothing -> pure False
+        Just _  -> pure True
+
 {-------------------------------------------------------------------------------
                                Internals
 -------------------------------------------------------------------------------}
@@ -386,16 +432,47 @@ 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 <- filterByAddressM isOurs' (applyTxToUTxO tx prevUTxO)
+
         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:
+            let known = prevUTxO <> utxoFromTx tx
+            -- But not all those transaction inputs/outputs belong to us:
+            ownedAndKnown <- filterByAddressM isOurs' known
+            -- 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 =
+                Set.fromList (inputs tx)
+                    `Set.intersection` (Set.fromList ownedAndKnownTxIns)
+                /= mempty
+        -- A transaction has a known output if one of the transaction outputs
+        -- matches a transaction output we know about.
+        let hasKnownOutput =
+                Set.fromList (outputs tx)
+                    `Set.intersection` (Set.fromList ownedAndKnownTxOuts)
+                /= mempty
         let hasKnownWithdrawal = ourWithdrawals /= mempty
 
         -- NOTE 1: The only case where fees can be 'Nothing' is when dealing with
@@ -423,7 +500,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 +510,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 +528,4 @@ 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 isOurs' . utxoFromTx) (Set.toList pending)

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,29 @@ null (UTxO u) = Map.null u
 size :: UTxO -> Int
 size (UTxO u) = Map.size u
 
+-- | Limit a UTxO set to just the UTxOs that are ours, according to some
+-- given function.
+--
+-- 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
+    where
+        filterFunc :: TxIn -> TxOut -> f (Maybe TxOut)
+        filterFunc _txin txout = do
+            ours <- isOursF $ view #address txout
+            pure $ if ours then Just txout else Nothing
+
+filterByAddress :: (Address -> Bool) -> UTxO -> UTxO
+filterByAddress f = runIdentity . filterByAddressM (pure . f)
+
 data UTxOStatistics = UTxOStatistics
     { histogram :: ![HistogramBar]
     , allStakes :: !Word64