Use indexed UTxO type within Wallet.

This change provides a new operation `availableUTxOIndex`, which is similar to `availableUTxO`, but returns an indexed set. This makes it possible for a multi-asset coin selection algorithm to efficiently search for entries containing a particular asset, without having to traverse the entire UTxO set.
cardano-foundation · Jan 11, 2021 · c962fe5 · c962fe5
1 parent 8777920
commit c962fe5
Show file tree

Hide file tree

Showing 5 changed files with 70 additions and 23 deletions.
diff --git a/lib/core/src/Cardano/Wallet/DB/Model.hs b/lib/core/src/Cardano/Wallet/DB/Model.hs
@@ -100,8 +100,6 @@ import Cardano.Wallet.Primitive.Types.Tx
     , TxMeta (..)
     , TxStatus (..)
     )
-import Cardano.Wallet.Primitive.Types.UTxO
-    ( UTxO (..) )
 import Control.Monad
     ( when )
 import Data.Bifunctor
@@ -129,6 +127,7 @@ import GHC.Generics
 import Numeric.Natural
     ( Natural )
 
+import qualified Cardano.Wallet.Primitive.Types.UTxOIndex as UTxOIndex
 import qualified Data.Map.Strict as Map
 
 {-------------------------------------------------------------------------------
@@ -454,7 +453,7 @@ mReadTxHistory ti wid minWithdrawal order range mstatus db@(Database wallets txs
         , txInfoFee =
             fee tx
         , txInfoInputs =
-            (\(inp, amt) -> (inp, amt, Map.lookup inp $ getUTxO $ utxo cp))
+            (\(inp, amt) -> (inp, amt, UTxOIndex.lookup inp $ utxo cp))
                 <$> resolvedInputs tx
         , txInfoOutputs =
             outputs tx

diff --git a/lib/core/src/Cardano/Wallet/DB/Sqlite.hs b/lib/core/src/Cardano/Wallet/DB/Sqlite.hs
@@ -226,6 +226,7 @@ import qualified Cardano.Wallet.Primitive.Types.Hash as W
 import qualified Cardano.Wallet.Primitive.Types.TokenBundle as TokenBundle
 import qualified Cardano.Wallet.Primitive.Types.Tx as W
 import qualified Cardano.Wallet.Primitive.Types.UTxO as W
+import qualified Cardano.Wallet.Primitive.Types.UTxOIndex as UTxOIndex
 import qualified Data.List as L
 import qualified Data.Map.Strict as Map
 import qualified Data.Text as T
@@ -1518,15 +1519,15 @@ mkCheckpointEntity wid wal =
         }
     utxo =
         [ UTxO wid sl (TxId input) ix addr (TokenBundle.getCoin tokens)
-        | (W.TxIn input ix, W.TxOut addr tokens) <- utxoMap
+        | (W.TxIn input ix, W.TxOut addr tokens) <- utxoEntries
         ]
     utxoTokens =
         [ UTxOToken wid sl (TxId input) ix policy token quantity
-        | (W.TxIn input ix, W.TxOut {tokens}) <- utxoMap
+        | (W.TxIn input ix, W.TxOut {tokens}) <- utxoEntries
         , let tokenList = snd (TokenBundle.toFlatList tokens)
         , (AssetId policy token, quantity) <- tokenList
         ]
-    utxoMap = Map.assocs (W.getUTxO (W.utxo wal))
+    utxoEntries = UTxOIndex.toList (W.utxo wal)
 
 -- note: TxIn records must already be sorted by order
 -- and TxOut records must already by sorted by index.

diff --git a/lib/core/src/Cardano/Wallet/Primitive/Model.hs b/lib/core/src/Cardano/Wallet/Primitive/Model.hs
@@ -10,6 +10,7 @@
 {-# LANGUAGE OverloadedLabels #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
 
 -- |
 -- Copyright: © 2018-2020 IOHK
@@ -47,6 +48,7 @@ module Cardano.Wallet.Primitive.Model
     , totalBalance
     , totalUTxO
     , availableUTxO
+    , availableUTxOIndex
     , utxo
     ) where
 
@@ -80,6 +82,8 @@ import Cardano.Wallet.Primitive.Types.Tx
     )
 import Cardano.Wallet.Primitive.Types.UTxO
     ( Dom (..), UTxO (..), balance, excluding, restrictedBy )
+import Cardano.Wallet.Primitive.Types.UTxOIndex
+    ( UTxOIndex )
 import Control.DeepSeq
     ( NFData (..), deepseq )
 import Control.Monad
@@ -88,6 +92,8 @@ import Control.Monad.Extra
     ( mapMaybeM )
 import Control.Monad.Trans.State.Strict
     ( State, evalState, runState, state )
+import Data.Function
+    ( (&) )
 import Data.Functor
     ( (<&>) )
 import Data.Generics.Internal.VL.Lens
@@ -102,13 +108,17 @@ import Data.Quantity
     ( Quantity (..) )
 import Data.Set
     ( Set )
+import Data.Word
+    ( Word64 )
 import Fmt
     ( Buildable (..), indentF )
 import GHC.Generics
     ( Generic )
 import Numeric.Natural
     ( Natural )
 
+import qualified Cardano.Wallet.Primitive.Types.TokenBundle as TokenBundle
+import qualified Cardano.Wallet.Primitive.Types.UTxOIndex as UTxOIndex
 import qualified Data.List.NonEmpty as NE
 import qualified Data.Map as Map
 import qualified Data.Set as Set
@@ -150,7 +160,7 @@ import qualified Data.Set as Set
 -- @
 data Wallet s = Wallet
     { -- | Unspent tx outputs belonging to this wallet
-      utxo :: UTxO
+      utxo :: UTxOIndex
 
       -- | Header of the latest applied block (current tip)
     , currentTip :: BlockHeader
@@ -169,7 +179,7 @@ instance NFData s => NFData (Wallet s) where
 instance Buildable s => Buildable (Wallet s) where
     build (Wallet u tip s) = "Wallet s\n"
         <> indentF 4 ("Tip: " <> build tip)
-        <> indentF 4 ("UTxO:\n" <> indentF 4 (build u))
+        <> indentF 4 ("UTxO:\n" <> indentF 4 (build $ UTxOIndex.toUTxO u))
         <> indentF 4 (build s)
 
 {-------------------------------------------------------------------------------
@@ -188,7 +198,7 @@ initWallet
     -> ([(Tx, TxMeta)], Wallet s)
 initWallet block s =
     let
-        ((FilteredBlock _ txs, u), s') = prefilterBlock block mempty s
+        ((FilteredBlock _ txs, u), s') = prefilterBlock block UTxOIndex.empty s
     in
         (txs, Wallet u (header block) s')
 
@@ -205,7 +215,7 @@ unsafeInitWallet
     -> s
     -- ^ Address discovery state
     -> Wallet s
-unsafeInitWallet = Wallet
+unsafeInitWallet = Wallet . UTxOIndex.fromUTxO
 
 -- | Update the state of an existing Wallet model
 updateState
@@ -279,8 +289,12 @@ applyBlocks (block0 :| blocks) cp =
 
 -- | Available balance = 'balance' . 'availableUTxO'
 availableBalance :: Set Tx -> Wallet s -> Natural
-availableBalance pending =
-    balance . availableUTxO pending
+availableBalance pending w = w
+    & availableUTxOIndex pending
+    & UTxOIndex.balance
+    & TokenBundle.getCoin
+    & unCoin
+    & fromIntegral @Word64 @Natural
 
 -- | Total balance = 'balance' . 'totalUTxO' +? rewards
 totalBalance
@@ -307,7 +321,19 @@ availableUTxO
     -> Wallet s
     -> UTxO
 availableUTxO pending (Wallet u _ _) =
-    u  `excluding` txIns pending
+    UTxOIndex.toUTxO u `excluding` txIns pending
+
+-- | Similar to `availableUTxO`, but returns an indexed UTxO set.
+--
+-- The index makes it possible to efficiently search for UTxO entries that
+-- contain a particular asset, without having to traverse the entire UTxO set.
+--
+availableUTxOIndex
+    :: Set Tx
+    -> Wallet s
+    -> UTxOIndex
+availableUTxOIndex pending (Wallet u _ _) =
+    UTxOIndex.deleteMany (txIns pending) u
 
 -- | Total UTxO = 'availableUTxO' @<>@ 'changeUTxO'
 totalUTxO
@@ -346,9 +372,9 @@ totalUTxO pending wallet@(Wallet _ _ s) =
 prefilterBlock
     :: (IsOurs s Address, IsOurs s RewardAccount)
     => Block
-    -> UTxO
+    -> UTxOIndex
     -> s
-    -> ((FilteredBlock, UTxO), s)
+    -> ((FilteredBlock, UTxOIndex), s)
 prefilterBlock b u0 = runState $ do
     delegations <- mapMaybeM ourDelegation (b ^. #delegations)
     (transactions, ourU) <- foldM applyTx (mempty, u0) (b ^. #transactions)
@@ -379,19 +405,26 @@ prefilterBlock b u0 = runState $ do
         , amount = Quantity amt
         , expiry = Nothing
         }
+
     applyTx
         :: (IsOurs s Address, IsOurs s RewardAccount)
-        => ([(Tx, TxMeta)], UTxO)
+        => ([(Tx, TxMeta)], UTxOIndex)
         -> Tx
-        -> State s ([(Tx, TxMeta)], UTxO)
+        -> State s ([(Tx, TxMeta)], UTxOIndex)
     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
+        let ourIns = Set.intersection
+                (Set.fromList $ inputs tx)
+                (dom $ UTxOIndex.toUTxO u <> ourU)
+        let u' = u
+                & UTxOIndex.insertMany (Map.toList $ getUTxO ourU)
+                & UTxOIndex.deleteMany ourIns
         ourWithdrawals <- fmap (fromIntegral . unCoin . snd) <$>
             mapMaybeM ourWithdrawal (Map.toList $ withdrawals tx)
         let received = balance ourU
-        let spent = balance (u `restrictedBy` ourIns) + sum ourWithdrawals
+        let spent
+                = balance (UTxOIndex.toUTxO u `restrictedBy` ourIns)
+                + sum ourWithdrawals
         let hasKnownInput = ourIns /= mempty
         let hasKnownOutput = ourU /= mempty
         let hasKnownWithdrawal = ourWithdrawals /= mempty

diff --git a/lib/core/test/unit/Cardano/Wallet/DB/Arbitrary.hs b/lib/core/test/unit/Cardano/Wallet/DB/Arbitrary.hs
@@ -199,6 +199,7 @@ import Test.Utils.Time
 import qualified Cardano.Wallet.Primitive.AddressDerivation.Byron as Byron
 import qualified Cardano.Wallet.Primitive.AddressDerivation.Shelley as Shelley
 import qualified Cardano.Wallet.Primitive.AddressDiscovery.Sequential as Seq
+import qualified Cardano.Wallet.Primitive.Types.UTxOIndex as UTxOIndex
 import qualified Data.ByteArray as BA
 import qualified Data.ByteString as BS
 import qualified Data.ByteString.Char8 as B8
@@ -302,7 +303,7 @@ instance GenState s => Arbitrary (InitialCheckpoint s) where
     arbitrary = do
         cp <- arbitrary @(Wallet s)
         pure $ InitialCheckpoint $ unsafeInitWallet
-            (utxo cp)
+            (UTxOIndex.toUTxO $ utxo cp)
             (block0 ^. #header)
             (getState cp)
 
@@ -313,7 +314,7 @@ instance GenState s => Arbitrary (InitialCheckpoint s) where
 instance GenState s => Arbitrary (Wallet s) where
     shrink w =
         [ unsafeInitWallet u (currentTip w) s
-        | (u, s) <- shrink (utxo w, getState w) ]
+        | (u, s) <- shrink (UTxOIndex.toUTxO (utxo w), getState w) ]
     arbitrary = unsafeInitWallet
         <$> arbitrary
         <*> arbitrary

diff --git a/lib/core/test/unit/Cardano/Wallet/DB/StateMachine.hs b/lib/core/test/unit/Cardano/Wallet/DB/StateMachine.hs
@@ -136,7 +136,7 @@ import Cardano.Wallet.Primitive.Types.RewardAccount
 import Cardano.Wallet.Primitive.Types.TokenBundle
     ( TokenBundle )
 import Cardano.Wallet.Primitive.Types.TokenMap
-    ( TokenMap )
+    ( AssetId, TokenMap )
 import Cardano.Wallet.Primitive.Types.TokenPolicy
     ( TokenName, TokenPolicyId )
 import Cardano.Wallet.Primitive.Types.TokenQuantity
@@ -154,6 +154,8 @@ import Cardano.Wallet.Primitive.Types.Tx
     )
 import Cardano.Wallet.Primitive.Types.UTxO
     ( UTxO (..) )
+import Cardano.Wallet.Primitive.Types.UTxOIndex
+    ( UTxOIndex )
 import Control.Foldl
     ( Fold (..) )
 import Control.Monad
@@ -184,6 +186,8 @@ import Data.Quantity
     ( Percentage (..), Quantity (..) )
 import Data.Set
     ( Set )
+import Data.Set.Strict.NonEmptySet
+    ( NonEmptySet )
 import Data.TreeDiff
     ( ToExpr (..), defaultExprViaShow, genericToExpr )
 import Data.Word
@@ -861,12 +865,18 @@ instance ToExpr s => ToExpr (Mock s) where
 instance (ToExpr k, ToExpr v) => ToExpr (NonEmptyMap k v) where
     toExpr = genericToExpr
 
+instance (ToExpr a) => ToExpr (NonEmptySet a) where
+    toExpr = genericToExpr
+
 instance ToExpr WalletId where
     toExpr = defaultExprViaShow
 
 instance ToExpr s => ToExpr (Wallet s) where
     toExpr = genericToExpr
 
+instance ToExpr UTxOIndex where
+    toExpr = genericToExpr
+
 instance ToExpr BlockHeader where
     toExpr = genericToExpr
 
@@ -943,6 +953,9 @@ instance ToExpr TokenPolicyId where
 instance ToExpr TokenQuantity where
     toExpr = genericToExpr
 
+instance ToExpr AssetId where
+    toExpr = genericToExpr
+
 instance ToExpr Address where
     toExpr = genericToExpr