Resolve PartialTx inputs over LSQ

Includes squashed commits:
- Avoid querying node for `getUTxOByTxIn mempty`
- Move logic to `Write.forceUTxOInEra`

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/src/Cardano/Wallet/Network/Implementation.hs

@@ -211,6 +211,7 @@ import Data.Function
     )
 import Data.Functor
     ( ($>)
+    , (<&>)
     )
 import Data.Functor.Contravariant
     ( Contravariant (..)
@@ -518,7 +519,7 @@ withNodeNetworkLayerBase
                 , stakeDistribution =
                     _stakeDistribution queryRewardQ
                 , getUTxOByTxIn =
-                    _getUTxOByTxIn queryRewardQ
+                    _getUTxOByTxIn queryRewardQ readCurrentNodeEra
                 , getCachedRewardAccountBalance =
                     _getCachedRewardAccountBalance rewardsObserver
                 , fetchRewardAccountBalances =
@@ -661,8 +662,17 @@ withNodeNetworkLayerBase
                     return res
                 Nothing -> pure $ StakePoolsSummary 0 mempty mempty
 
-        _getUTxOByTxIn queue ins =
-            bracketQuery "getUTxOByTxIn" tr
+        _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

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,28 @@ balanceTx wrk pp timeTranslation partialTx = do
         utxoIndex
         (defaultChangeAddressGen argGenChange)
         changeState
-        partialTx
+        (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 +3884,7 @@ data WalletException
     | ExceptionSignPayment ErrSignPayment
     | forall era. Write.IsRecentEra era => ExceptionBalanceTx (ErrBalanceTx era)
     | ExceptionWriteTxEra ErrWriteTxEra
+    | ExceptionInvalidTxOutInEra Write.ErrInvalidTxOutInEra
     | ExceptionSubmitTransaction ErrSubmitTransaction
     | ExceptionConstructTx ErrConstructTx
     | ExceptionGetPolicyId ErrGetPolicyId