Resolve PartialTx inputs over LSQ (#4630)

- Add `getUTxOByTxIn` local state query
- Resolve `PartialTx` inputs over LSQ in `Cardano.Wallet.balanceTx`

### Comments

This is part of the fix for the [integration test in the 8.11
branch](https://buildkite.com/cardano-foundation/cardano-wallet/builds/5254#018fee3c-01ad-441e-bb7e-25acd78cb7f3);
if we have the UTxO of the reference input containing the script used
for minting in constructTx, then `Ledger.getMinFeeTxUtxo pp tx utxo`
will correctly account for minFeeRefScriptCoinsPerByte *
totalRefScriptSize.

Previous step: #4629
Next step: #4631

### Issue Number

ADP-3373

lib/api/src/Cardano/Wallet/Api/Http/Server/Error.hs

@@ -263,6 +263,7 @@ instance IsServerError WalletException where
         ExceptionSoftDerivationIndex e -> toServerError e
         ExceptionHardenedDerivationIndex e -> toServerError e
         ExceptionVoting e -> toServerError e
+        ExceptionInvalidTxOutInEra e -> toServerError e
 
 instance IsServerError ErrNoSuchWallet where
     toServerError = \case

lib/api/src/Cardano/Wallet/Api/Http/Shelley/Server.hs

@@ -1000,7 +1000,7 @@ import qualified Internal.Cardano.Write.Tx.Balance as Write
     ( PartialTx (PartialTx)
     )
 import qualified Internal.Cardano.Write.Tx.Sign as Write
-    ( TimelockKeyWitnessCounts (TimelockKeyWitnessCounts)
+    ( TimelockKeyWitnessCounts (..)
     , estimateMinWitnessRequiredPerInput
     )
 import qualified Network.Ntp as Ntp

lib/balance-tx/lib/internal/Internal/Cardano/Write/Tx.hs

@@ -99,6 +99,7 @@ module Internal.Cardano.Write.Tx
     , TxOutInRecentEra (..)
     , ErrInvalidTxOutInEra (..)
     , unwrapTxOutInRecentEra
+    , wrapTxOutInRecentEra
 
     , computeMinimumCoinForTxOut
     , isBelowMinimumCoinForTxOut
@@ -140,6 +141,7 @@ module Internal.Cardano.Write.Tx
     , Shelley.UTxO (..)
     , utxoFromTxOutsInRecentEra
     , unsafeUtxoFromTxOutsInRecentEra
+    , forceUTxOToEra
 
     -- * Policy and asset identifiers
     , type PolicyId
@@ -176,9 +178,11 @@ import Cardano.Ledger.Alonzo.UTxO
 import Cardano.Ledger.Api
     ( coinTxOutL
     , ppKeyDepositL
+    , upgradeTxOut
     )
 import Cardano.Ledger.Api.UTxO
     ( EraUTxO (ScriptsNeeded)
+    , UTxO (..)
     )
 import Cardano.Ledger.Babbage.TxBody
     ( BabbageTxOut (..)
@@ -188,6 +192,7 @@ import Cardano.Ledger.BaseTypes
     , StrictMaybe (..)
     , Version
     , maybeToStrictMaybe
+    , strictMaybeToMaybe
     )
 import Cardano.Ledger.Coin
     ( Coin (..)
@@ -221,7 +226,8 @@ import Cardano.Ledger.Val
     , modifyCoin
     )
 import Control.Arrow
-    ( (>>>)
+    ( second
+    , (>>>)
     )
 import Data.ByteString
     ( ByteString
@@ -591,6 +597,18 @@ data TxOutInRecentEra =
         (Maybe (AlonzoScript LatestLedgerEra))
         -- Same contents as 'TxOut LatestLedgerEra'.
 
+wrapTxOutInRecentEra
+    :: forall era. IsRecentEra era
+    => TxOut era
+    -> TxOutInRecentEra
+wrapTxOutInRecentEra out = case recentEra @era of
+    RecentEraConway ->
+        let
+            BabbageTxOut addr v d s = out
+        in
+            TxOutInRecentEra addr v d (strictMaybeToMaybe s)
+    RecentEraBabbage -> wrapTxOutInRecentEra @ConwayEra $ upgradeTxOut out
+
 data ErrInvalidTxOutInEra
     = InlinePlutusV3ScriptNotSupportedInBabbage
     deriving (Show, Eq)
@@ -716,6 +734,15 @@ unsafeUtxoFromTxOutsInRecentEra
 unsafeUtxoFromTxOutsInRecentEra =
     either (error . show) id . utxoFromTxOutsInRecentEra
 
+forceUTxOToEra
+    :: forall era1 era2. (IsRecentEra era1, IsRecentEra era2)
+    => UTxO era1
+    -> Either ErrInvalidTxOutInEra (UTxO era2)
+forceUTxOToEra (UTxO utxo) =
+    utxoFromTxOutsInRecentEra
+    $ map (second wrapTxOutInRecentEra)
+    $ Map.toList utxo
+
 --------------------------------------------------------------------------------
 -- Tx
 --------------------------------------------------------------------------------

lib/network-layer/cardano-wallet-network-layer.cabal

@@ -54,6 +54,7 @@ library
     Cardano.Wallet.Network.LocalStateQuery.PParams
     Cardano.Wallet.Network.LocalStateQuery.RewardAccount
     Cardano.Wallet.Network.LocalStateQuery.StakeDistribution
+    Cardano.Wallet.Network.LocalStateQuery.UTxO
     Cardano.Wallet.Network.Logging
     Cardano.Wallet.Network.Logging.Aggregation
     Cardano.Wallet.Network.RestorationMode

lib/network-layer/src/Cardano/Wallet/Network.hs

@@ -153,6 +153,9 @@ data NetworkLayer m block = NetworkLayer
     , stakeDistribution
         :: Coin -- Stake to consider for rewards
         -> m StakePoolsSummary
+    , getUTxOByTxIn
+        :: Set Write.TxIn
+        -> m (MaybeInRecentEra Write.UTxO)
     , getCachedRewardAccountBalance
         :: RewardAccount
         -- Either reward account from key hash or script hash

lib/network-layer/src/Cardano/Wallet/Network/Implementation.hs

@@ -211,6 +211,7 @@ import Data.Function
     )
 import Data.Functor
     ( ($>)
+    , (<&>)
     )
 import Data.Functor.Contravariant
     ( Contravariant (..)
@@ -517,6 +518,8 @@ withNodeNetworkLayerBase
                     _postTx txSubmissionQ readCurrentNodeEra
                 , stakeDistribution =
                     _stakeDistribution queryRewardQ
+                , getUTxOByTxIn =
+                    _getUTxOByTxIn queryRewardQ readCurrentNodeEra
                 , getCachedRewardAccountBalance =
                     _getCachedRewardAccountBalance rewardsObserver
                 , fetchRewardAccountBalances =
@@ -659,6 +662,19 @@ withNodeNetworkLayerBase
                     return res
                 Nothing -> pure $ StakePoolsSummary 0 mempty mempty
 
+        _getUTxOByTxIn queue readCachedEra ins
+            | ins == mempty = readCachedEra <&> \case
+                AnyCardanoEra ByronEra -> InNonRecentEraByron
+                AnyCardanoEra ShelleyEra -> InNonRecentEraShelley
+                AnyCardanoEra AllegraEra -> InNonRecentEraAllegra
+                AnyCardanoEra MaryEra -> InNonRecentEraMary
+                AnyCardanoEra AlonzoEra -> InNonRecentEraAlonzo
+                AnyCardanoEra BabbageEra -> InRecentEraBabbage mempty
+                AnyCardanoEra ConwayEra -> InRecentEraConway mempty
+            | otherwise
+                = bracketQuery "getUTxOByTxIn" tr
+                $ queue `send` SomeLSQ (LSQ.getUTxOByTxIn ins)
+
         _watchNodeTip readTip callback = do
             observeForever readTip $ \tip -> do
                 let tip' = fromOuroborosTip tip

lib/network-layer/src/Cardano/Wallet/Network/LocalStateQuery.hs

@@ -9,6 +9,7 @@ module Cardano.Wallet.Network.LocalStateQuery
     , module Cardano.Wallet.Network.LocalStateQuery.PParams
     , module Cardano.Wallet.Network.LocalStateQuery.RewardAccount
     , module Cardano.Wallet.Network.LocalStateQuery.StakeDistribution
+    , module Cardano.Wallet.Network.LocalStateQuery.UTxO
     ) where
 
 import Cardano.Wallet.Network.LocalStateQuery.Extra
@@ -25,3 +26,6 @@ import Cardano.Wallet.Network.LocalStateQuery.RewardAccount
 import Cardano.Wallet.Network.LocalStateQuery.StakeDistribution
     ( stakeDistribution
     )
+import Cardano.Wallet.Network.LocalStateQuery.UTxO
+    ( getUTxOByTxIn
+    )

lib/network-layer/src/Cardano/Wallet/Network/LocalStateQuery/UTxO.hs

@@ -0,0 +1,57 @@
+{-# LANGUAGE GADTs #-}
+-- |
+-- Copyright: © 2024 Cardano Foundation
+-- License: Apache-2.0
+--
+-- A local state query that looks up UTxOs based on TxIns.
+--
+module Cardano.Wallet.Network.LocalStateQuery.UTxO
+    ( getUTxOByTxIn
+    ) where
+
+import Prelude
+
+import Cardano.Ledger.TxIn
+    ( TxIn
+    )
+import Cardano.Ledger.UTxO
+    ( UTxO
+    )
+import Cardano.Wallet.Network.Implementation.Ouroboros
+    ( LSQ (..)
+    )
+import Cardano.Wallet.Network.LocalStateQuery.Extra
+    ( onAnyEra
+    )
+import Data.Set
+    ( Set
+    )
+import Internal.Cardano.Write.Tx
+    ( MaybeInRecentEra (..)
+    )
+import Ouroboros.Consensus.Cardano
+    ( CardanoBlock
+    )
+import Ouroboros.Consensus.Shelley.Eras
+    ( StandardCrypto
+    )
+
+import qualified Ouroboros.Consensus.Shelley.Ledger as Shelley
+
+{-----------------------------------------------------------------------------
+    Local State Query for GetUTxOByTxIn
+------------------------------------------------------------------------------}
+--
+type LSQ' m = LSQ (CardanoBlock StandardCrypto) m
+
+getUTxOByTxIn
+    :: Set (TxIn StandardCrypto) -> LSQ' m (MaybeInRecentEra UTxO)
+getUTxOByTxIn ins =
+    onAnyEra
+        (pure InNonRecentEraByron)
+        (pure InNonRecentEraShelley)
+        (pure InNonRecentEraAllegra)
+        (pure InNonRecentEraMary)
+        (pure InNonRecentEraAlonzo)
+        (InRecentEraBabbage <$> LSQry (Shelley.GetUTxOByTxIn ins))
+        (InRecentEraConway <$> LSQry (Shelley.GetUTxOByTxIn ins))

lib/unit/test-common/Cardano/Wallet/DummyTarget/Primitive/Types.hs

@@ -222,6 +222,7 @@ dummyNetworkLayer = NetworkLayer
     , currentProtocolParametersInRecentEras
         = err "currentProtocolParametersInRecentEras"
     , currentSlottingParameters = err "currentSlottingParameters"
+    , getUTxOByTxIn = err "getUTxOByTxIn"
     , postTx = err "postTx"
     , stakeDistribution = err "stakeDistribution"
     , getCachedRewardAccountBalance = err "getRewardCachedAccountBalance"

lib/wallet/src/Cardano/Wallet.hs

@@ -272,7 +272,8 @@ import Cardano.Crypto.Wallet
     ( toXPub
     )
 import Cardano.Ledger.Api
-    ( bodyTxL
+    ( EraTxBody (allInputsTxBodyF)
+    , bodyTxL
     , feeTxBodyL
     )
 import Cardano.Mnemonic
@@ -720,6 +721,7 @@ import Data.Functor.Contravariant
     )
 import Data.Generics.Internal.VL.Lens
     ( Lens'
+    , over
     , view
     , (.~)
     , (^.)
@@ -795,7 +797,8 @@ import GHC.TypeNats
     ( Nat
     )
 import Internal.Cardano.Write.Tx
-    ( recentEra
+    ( MaybeInRecentEra (..)
+    , recentEra
     , toRecentEraGADT
     )
 import Internal.Cardano.Write.Tx.Balance
@@ -867,15 +870,18 @@ import qualified Data.Vector as V
 import qualified Internal.Cardano.Write.Tx as Write
     ( AnyRecentEra
     , CardanoApiEra
+    , ErrInvalidTxOutInEra
     , FeePerByte
-    , IsRecentEra
+    , IsRecentEra (..)
+    , MaybeInRecentEra (..)
     , PParams
     , PParamsInAnyRecentEra (PParamsInAnyRecentEra)
     , RecentEra (..)
     , Tx
     , UTxO (UTxO)
     , cardanoEraFromRecentEra
     , feeOfBytes
+    , forceUTxOToEra
     , fromCardanoApiTx
     , getFeePerByte
     , stakeKeyDeposit
@@ -2155,6 +2161,9 @@ readNodeTipStateForTxWrite netLayer = do
 -- workflow with Shelley- and Shared- wallet flavors.
 --
 -- Changes to the change state are not written to the DB.
+--
+-- Inputs of the partial transaction are looked up using a local state query to
+-- the node.
 balanceTx
     :: forall s era.
         ( GenChange s
@@ -2175,6 +2184,13 @@ balanceTx wrk pp timeTranslation partialTx = do
             Write.constructUTxOIndex $
             Write.fromWalletUTxO utxo
 
+    -- Resolve inputs using LSQ. Useful for foreign reference inputs supplied by
+    -- the user when calling transactions-construct, or in transactions-balance.
+    let netLayer = wrk ^. networkLayer
+    let inputsToLookup = partialTx ^. #tx . bodyTxL . allInputsTxBodyF
+    lookedUpUTxO <- liftIO $
+        forceUTxOToEra =<< getUTxOByTxIn netLayer inputsToLookup
+
     let utxoAssumptions = case walletFlavor @s of
             ShelleyWallet -> AllKeyPaymentCredentials
             SharedWallet -> AllScriptPaymentCredentialsFrom
@@ -2190,10 +2206,29 @@ balanceTx wrk pp timeTranslation partialTx = do
         utxoIndex
         (defaultChangeAddressGen argGenChange)
         changeState
-        partialTx
+        -- In case of conflicts, the UTxO looked up from the node will win.
+        (over #extraUTxO (lookedUpUTxO <>) partialTx)
 
     return tx
   where
+    -- Assumes the 'utxo' was queried from the node /after/ the 'era'. As
+    -- rolling back to a previous era should be impossible, we know 'IsRecentEra
+    -- era => IsRecentEra eraOfUTxO'.
+    forceUTxOToEra
+        :: Write.MaybeInRecentEra Write.UTxO
+        -> IO (Write.UTxO era)
+    forceUTxOToEra = \case
+        InRecentEraConway utxo -> hoist $ Write.forceUTxOToEra utxo
+        InRecentEraBabbage utxo -> hoist $ Write.forceUTxOToEra utxo
+        InNonRecentEraAlonzo -> impossibleRollback
+        InNonRecentEraMary -> impossibleRollback
+        InNonRecentEraAllegra -> impossibleRollback
+        InNonRecentEraShelley -> impossibleRollback
+        InNonRecentEraByron -> impossibleRollback
+      where
+        impossibleRollback = error "forceUTxOToEra: era should not roll back"
+        hoist = either (throwIO . ExceptionInvalidTxOutInEra) pure
+
     argGenChange :: ArgGenChange s
     argGenChange = case walletFlavor @s of
         ShelleyWallet -> delegationAddressS @(NetworkOf s)
@@ -3850,6 +3885,7 @@ data WalletException
     | ExceptionSignPayment ErrSignPayment
     | forall era. Write.IsRecentEra era => ExceptionBalanceTx (ErrBalanceTx era)
     | ExceptionWriteTxEra ErrWriteTxEra
+    | ExceptionInvalidTxOutInEra Write.ErrInvalidTxOutInEra
     | ExceptionSubmitTransaction ErrSubmitTransaction
     | ExceptionConstructTx ErrConstructTx
     | ExceptionGetPolicyId ErrGetPolicyId