[ADP-3272] Add the TxWithUTxO data type. (#4570)

## Summary

This PR:

- adds the `TxWithUTxO` data type, along with smart constructors,
generators, and shrinkers.
- uses the `TxWithUTxO` data type to simplify the `Arbitrary` instance
for `PartialTx`.

## Details

An object of type `TxWithUTxO` combines a transaction `t` with an
associated UTxO set `u`, such that:
- every input in transaction `t` is **_guaranteed_** to resolve to a
UTxO within UTxO set `u`.
- UTXO set `u` **_may_** also contain **_additional_** UTxOs that are
not referenced by `t`.

In order to uphold the above guarantee:
- the `TxWithUTxO` data type is defined within its own module;
- the `TxWithUTxO` data constructor is not exported;
- a pattern synonym `TxWithUTxO` is provided to facilitate safe pattern
matching.

## Motivation

This data type will be used for testing `balanceTx` in situations where
we want to guarantee that `balanceTx` will not fail with
`ErrBalanceTxUnresolvedInputs`.

## Issue

ADP-3272

lib/balance-tx/cardano-balance-tx.cabal

@@ -65,9 +65,11 @@ library internal
     , cardano-slotting
     , cardano-strict-containers
     , cardano-wallet-primitive
+    , cardano-wallet-test-utils
     , cborg
     , containers
     , deepseq
+    , either
     , fmt
     , generic-lens
     , groups
@@ -95,6 +97,8 @@ library internal
     Internal.Cardano.Write.Tx.Sign
     Internal.Cardano.Write.Tx.SizeEstimation
     Internal.Cardano.Write.Tx.TimeTranslation
+    Internal.Cardano.Write.Tx.TxWithUTxO
+    Internal.Cardano.Write.Tx.TxWithUTxO.Gen
     Internal.Cardano.Write.UTxOAssumptions
 
 test-suite test

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

@@ -0,0 +1,143 @@
+{-# LANGUAGE ExplicitNamespaces #-}
+{-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE StandaloneDeriving #-}
+
+-- |
+-- Copyright: © 2024 Cardano Foundation
+-- License: Apache-2.0
+--
+-- Provides the 'TxWithUTxO' data type.
+--
+module Internal.Cardano.Write.Tx.TxWithUTxO
+    ( type TxWithUTxO
+    , pattern TxWithUTxO
+    , construct
+    , constructFiltered
+    , isValid
+    )
+    where
+
+import Prelude
+
+import Cardano.Ledger.Api
+    ( AlonzoEraTxBody (collateralInputsTxBodyL)
+    , BabbageEraTxBody (referenceInputsTxBodyL)
+    , EraTx (bodyTxL)
+    , EraTxBody (TxBody, inputsTxBodyL)
+    )
+import Cardano.Ledger.Api.Tx.Body
+    ( allInputsTxBodyF
+    )
+import Control.Lens
+    ( over
+    , view
+    )
+import Data.Either.Combinators
+    ( maybeToLeft
+    )
+import Data.Maybe
+    ( fromMaybe
+    )
+import Data.Semigroup.Cancellative
+    ( LeftReductive (stripPrefix)
+    )
+import Data.Set.NonEmpty
+    ( NESet
+    )
+import Internal.Cardano.Write.Tx
+    ( IsRecentEra
+    , Tx
+    , TxIn
+    , UTxO (UTxO)
+    )
+
+import qualified Data.Map.Strict as Map
+import qualified Data.Set as Set
+import qualified Data.Set.NonEmpty as NESet
+
+-- | A transaction with an associated UTxO set.
+--
+-- Every input in the transaction is guaranteed to resolve to a UTxO within the
+-- associated UTxO set.
+--
+-- The UTxO set may also contain additional UTxOs that are not referenced by
+-- the transaction.
+--
+data TxWithUTxO era = UnsafeTxWithUTxO !(Tx era) !(UTxO era)
+
+deriving instance IsRecentEra era => Eq (TxWithUTxO era)
+
+instance IsRecentEra era => Show (TxWithUTxO era) where
+    show = fromMaybe "TxWithUTxO" . stripPrefix "Unsafe" . show
+
+{-# COMPLETE TxWithUTxO #-}
+pattern TxWithUTxO :: IsRecentEra era => Tx era -> UTxO era -> TxWithUTxO era
+pattern TxWithUTxO tx utxo <- UnsafeTxWithUTxO tx utxo
+
+-- | Constructs a 'TxWithUTxO' object from an existing transaction and UTxO set.
+--
+-- Construction succeeds if (and only if) every single input within the given
+-- transaction resolves to a UTxO within the accompanying UTxO set.
+--
+-- Otherwise, if the transaction has any unresolvable inputs, this function
+-- returns the non-empty set of those inputs.
+--
+construct
+    :: IsRecentEra era
+    => Tx era
+    -> UTxO era
+    -> Either (NESet TxIn) (TxWithUTxO era)
+construct tx utxo =
+    maybeToLeft txWithUTxO (unresolvableInputs txWithUTxO)
+  where
+    txWithUTxO = UnsafeTxWithUTxO tx utxo
+
+-- | Constructs a 'TxWithUTxO' object from an existing transaction and UTxO set,
+--   automatically filtering out any unresolvable inputs from the transaction.
+--
+-- A transaction input is unresolvable if (and only if) it does not resolve to
+-- a UTxO within the given UTxO set.
+--
+constructFiltered
+    :: forall era. IsRecentEra era
+    => Tx era
+    -> UTxO era
+    -> TxWithUTxO era
+constructFiltered tx utxo@(UTxO utxoMap) = UnsafeTxWithUTxO txFiltered utxo
+  where
+    txFiltered :: Tx era
+    txFiltered = over bodyTxL removeUnresolvableInputs tx
+
+    removeUnresolvableInputs :: TxBody era -> TxBody era
+    removeUnresolvableInputs
+        = over           inputsTxBodyL f
+        . over collateralInputsTxBodyL f
+        . over  referenceInputsTxBodyL f
+      where
+        f = Set.filter (`Map.member` utxoMap)
+
+-- | Indicates whether or not a given 'TxWithUTxO' object is valid.
+--
+-- A 'TxWithUTxO' object is valid if (and only if) all inputs within the
+-- transaction resolve to a UTxO within the associated UTxO set.
+--
+isValid :: IsRecentEra era => TxWithUTxO era -> Bool
+isValid = null . unresolvableInputs
+
+-- | Finds the complete set of unresolvable transaction inputs.
+--
+-- A transaction input is unresolvable if (and only if) it does not resolve
+-- to a UTxO within the associated UTxO set.
+--
+-- For a valid 'TxWithUTxO' object, this function will return 'Nothing'.
+--
+unresolvableInputs
+    :: forall era. IsRecentEra era
+    => TxWithUTxO era
+    -> Maybe (NESet TxIn)
+unresolvableInputs (TxWithUTxO tx (UTxO utxo))
+    = NESet.nonEmptySet
+    . Set.filter (`Map.notMember` utxo)
+    . view (bodyTxL . allInputsTxBodyF)
+    $ tx

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

@@ -0,0 +1,129 @@
+{-# LANGUAGE ExplicitNamespaces #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+
+-- |
+-- Copyright: © 2024 Cardano Foundation
+-- License: Apache-2.0
+--
+-- Provides generators and shrinkers for the 'TxWithUTxO' data type.
+--
+module Internal.Cardano.Write.Tx.TxWithUTxO.Gen
+    ( generate
+    , generateWithMinimalUTxO
+    , generateWithSurplusUTxO
+    , shrinkWith
+    , shrinkTxWith
+    , shrinkUTxOWith
+    )
+    where
+
+import Prelude
+
+import Cardano.Ledger.Api
+    ( EraTx (bodyTxL)
+    )
+import Cardano.Ledger.Api.Tx.Body
+    ( allInputsTxBodyF
+    )
+import Control.Lens
+    ( view
+    )
+import Internal.Cardano.Write.Tx
+    ( IsRecentEra
+    , Tx
+    , TxIn
+    , TxOut
+    , UTxO (UTxO)
+    )
+import Internal.Cardano.Write.Tx.TxWithUTxO
+    ( pattern TxWithUTxO
+    , type TxWithUTxO
+    )
+import Test.QuickCheck
+    ( Gen
+    , frequency
+    )
+import Test.QuickCheck.Extra
+    ( genMapFromKeysWith
+    , genNonEmptyDisjointMap
+    , interleaveRoundRobin
+    )
+
+import qualified Internal.Cardano.Write.Tx.TxWithUTxO as TxWithUTxO
+
+-- | Generates a 'TxWithUTxO' object.
+--
+-- The domain of the UTxO map is a superset of the transaction input set, but
+-- it may or may not be a strict superset.
+--
+generate
+    :: IsRecentEra era
+    => Gen (Tx era)
+    -> Gen (TxIn)
+    -> Gen (TxOut era)
+    -> Gen (TxWithUTxO era)
+generate genTx genTxIn genTxOut =
+    frequency
+        [ (9, generateWithMinimalUTxO genTx         genTxOut)
+        , (1, generateWithSurplusUTxO genTx genTxIn genTxOut)
+        ]
+
+-- | Generates a 'TxWithUTxO' object that has a minimal UTxO set.
+--
+-- The domain of the UTxO map is exactly equal to the transaction input set.
+--
+generateWithMinimalUTxO
+    :: IsRecentEra era
+    => Gen (Tx era)
+    -> Gen (TxOut era)
+    -> Gen (TxWithUTxO era)
+generateWithMinimalUTxO genTx genTxOut = do
+    tx <- genTx
+    utxo <- UTxO <$> genMapFromKeysWith genTxOut (txAllInputs tx)
+    pure $ TxWithUTxO.constructFiltered tx utxo
+  where
+    txAllInputs = view (bodyTxL . allInputsTxBodyF)
+
+-- | Generates a 'TxWithUTxO' object that has a surplus UTxO set.
+--
+-- The domain of the UTxO map is a strict superset of the transaction input set.
+--
+generateWithSurplusUTxO
+    :: forall era. ()
+    => IsRecentEra era
+    => Gen (Tx era)
+    -> Gen (TxIn)
+    -> Gen (TxOut era)
+    -> Gen (TxWithUTxO era)
+generateWithSurplusUTxO genTx genTxIn genTxOut =
+    generateWithMinimalUTxO genTx genTxOut >>= \case
+        TxWithUTxO tx (UTxO utxo) -> do
+            utxoSurplus <- genNonEmptyDisjointMap genTxIn genTxOut utxo
+            pure $ TxWithUTxO.constructFiltered tx $ UTxO (utxo <> utxoSurplus)
+
+shrinkWith
+    :: IsRecentEra era
+    => (Tx         era -> [Tx         era])
+    -> (UTxO       era -> [UTxO       era])
+    -> (TxWithUTxO era -> [TxWithUTxO era])
+shrinkWith shrinkTx shrinkUTxO txWithUTxO =
+    interleaveRoundRobin
+        [ shrinkTxWith   shrinkTx   txWithUTxO
+        , shrinkUTxOWith shrinkUTxO txWithUTxO
+        ]
+
+shrinkTxWith
+    :: IsRecentEra era
+    => (Tx         era -> [Tx         era])
+    -> (TxWithUTxO era -> [TxWithUTxO era])
+shrinkTxWith shrinkTx (TxWithUTxO tx utxo) =
+    [ TxWithUTxO.constructFiltered tx' utxo | tx' <- shrinkTx tx ]
+
+shrinkUTxOWith
+    :: IsRecentEra era
+    => (UTxO       era -> [UTxO       era])
+    -> (TxWithUTxO era -> [TxWithUTxO era])
+shrinkUTxOWith shrinkUTxO (TxWithUTxO tx utxo) =
+    [ TxWithUTxO.constructFiltered tx utxo' | utxo' <- shrinkUTxO utxo ]