WIP

cardano-foundation · May 3, 2024 · e76ff81 · e76ff81
1 parent 13ccf9b
commit e76ff81
Show file tree

Hide file tree

Showing 3 changed files with 92 additions and 19 deletions.
diff --git a/lib/balance-tx/cardano-balance-tx.cabal b/lib/balance-tx/cardano-balance-tx.cabal
@@ -65,6 +65,7 @@ library internal
     , cardano-slotting
     , cardano-strict-containers
     , cardano-wallet-primitive
+    , cardano-wallet-test-utils
     , cborg
     , containers
     , deepseq
@@ -96,6 +97,7 @@ 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
 

diff --git a/lib/balance-tx/lib/internal/Internal/Cardano/Write/Tx/TxWithUTxO.hs b/lib/balance-tx/lib/internal/Internal/Cardano/Write/Tx/TxWithUTxO.hs
@@ -1,7 +1,4 @@
-{-# LANGUAGE DuplicateRecordFields #-}
 {-# LANGUAGE ExplicitNamespaces #-}
-{-# LANGUAGE NamedFieldPuns #-}
-{-# LANGUAGE NoFieldSelectors #-}
 {-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE StandaloneDeriving #-}
@@ -67,10 +64,7 @@ import qualified Data.Set.NonEmpty as NESet
 -- The UTxO set may also contain additional UTxOs that are not referenced by
 -- the transaction.
 --
-data TxWithUTxO era = UnsafeTxWithUTxO
-    { tx :: !(Tx era)
-    , utxo :: !(UTxO era)
-    }
+data TxWithUTxO era = UnsafeTxWithUTxO !(Tx era) !(UTxO era)
 
 deriving instance IsRecentEra era => Eq (TxWithUTxO era)
 
@@ -79,7 +73,7 @@ instance IsRecentEra era => Show (TxWithUTxO era) where
 
 {-# COMPLETE TxWithUTxO #-}
 pattern TxWithUTxO :: IsRecentEra era => Tx era -> UTxO era -> TxWithUTxO era
-pattern TxWithUTxO {tx, utxo} <- UnsafeTxWithUTxO {tx, utxo}
+pattern TxWithUTxO tx utxo <- UnsafeTxWithUTxO tx utxo
 
 -- | Constructs a 'TxWithUTxO' object, failing if there are unresolvable inputs.
 --
@@ -97,7 +91,7 @@ constructEither
 constructEither tx utxo =
     maybeToLeft txWithUTxO (unresolvableInputs txWithUTxO)
   where
-    txWithUTxO = UnsafeTxWithUTxO {tx, utxo}
+    txWithUTxO = UnsafeTxWithUTxO tx utxo
 
 -- | Constructs a 'TxWithUTxO' object, filtering out any unresolvable inputs.
 --
@@ -109,10 +103,10 @@ constructFiltered
     => Tx era
     -> UTxO era
     -> TxWithUTxO era
-constructFiltered tx0 utxo@(UTxO utxoMap) = UnsafeTxWithUTxO {tx, utxo}
+constructFiltered tx utxo@(UTxO utxoMap) = UnsafeTxWithUTxO txFiltered utxo
   where
-    tx :: Tx era
-    tx = over bodyTxL removeUnresolvableInputs tx0
+    txFiltered :: Tx era
+    txFiltered = over bodyTxL removeUnresolvableInputs tx
 
     removeUnresolvableInputs :: TxBody era -> TxBody era
     removeUnresolvableInputs
@@ -141,8 +135,8 @@ unresolvableInputs
     :: forall era. IsRecentEra era
     => TxWithUTxO era
     -> Maybe (NESet TxIn)
-unresolvableInputs TxWithUTxO {tx, utxo = UTxO utxoMap}
+unresolvableInputs (TxWithUTxO tx (UTxO utxo))
     = NESet.nonEmptySet
-    . Set.filter (`Map.notMember` utxoMap)
+    . Set.filter (`Map.notMember` utxo)
     . view (bodyTxL . allInputsTxBodyF)
     $ tx
diff --git a/lib/balance-tx/lib/internal/Internal/Cardano/Write/Tx/TxWithUTxO/Gen.hs b/lib/balance-tx/lib/internal/Internal/Cardano/Write/Tx/TxWithUTxO/Gen.hs
@@ -1,4 +1,7 @@
 {-# LANGUAGE ExplicitNamespaces #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 
 -- |
 -- Copyright: © 2024 Cardano Foundation
@@ -7,20 +10,94 @@
 -- Provides generators and shrinkers for the 'TxWithUTxO' data type.
 --
 module Internal.Cardano.Write.Tx.TxWithUTxO.Gen
-    ( genTxWithUTxO
+    ( generate
+    , generateWithMinimalUTxO
+    , generateWithSurplusUTxO
+    , shrinkTxOuts
     )
     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
-    ( type TxWithUTxO
+    ( pattern TxWithUTxO
+    , type TxWithUTxO
     )
 import Test.QuickCheck
     ( Gen
+    , oneof
+    , sublistOf
+    )
+import Test.QuickCheck.Extra
+    ( genMapFromKeysWith
+    , genMapWith
     )
 
---import qualified Internal.Cardano.Write.Tx.TxWithUTxO as TxWithUTxO
+import Data.Map.Strict
+    ( Map
+    )
+import Data.Set
+    ( Set
+    )
+import qualified Internal.Cardano.Write.Tx.TxWithUTxO as TxWithUTxO
+
+generate
+    :: IsRecentEra era
+    => Gen (Tx era)
+    -> Gen (TxIn)
+    -> Gen (TxOut era)
+    -> Gen (TxWithUTxO era)
+generate genTx genTxIn genTxOut =
+    oneof
+        [ generateWithMinimalUTxO genTx genTxIn genTxOut
+        , generateWithSurplusUTxO genTx genTxIn genTxOut
+        ]
+
+generateWithMinimalUTxO
+    :: IsRecentEra era
+    => Gen (Tx era)
+    -> Gen (TxIn)
+    -> Gen (TxOut era)
+    -> Gen (TxWithUTxO era)
+generateWithMinimalUTxO genTx _genTxIn genTxOut = do
+    tx <- genTx
+    utxo <- UTxO <$> genMapFromKeysWith genTxOut (txInputs tx)
+    pure $ TxWithUTxO.constructFiltered tx utxo
+  where
+    txInputs = view (bodyTxL . allInputsTxBodyF)
+
+generateWithSurplusUTxO
+    :: IsRecentEra era
+    => Gen (Tx era)
+    -> Gen (TxIn)
+    -> Gen (TxOut era)
+    -> Gen (TxWithUTxO era)
+generateWithSurplusUTxO genTx genTxIn genTxOut = do
+    generateWithMinimalUTxO genTx genTxIn genTxOut >>= \case
+        TxWithUTxO tx utxo -> do
+            utxoSurplus <- UTxO <$> genMapWith genTxIn genTxOut
+            pure $ TxWithUTxO.constructFiltered tx (utxo <> utxoSurplus)
 
-genTxWithUTxO :: Gen (TxWithUTxO era)
-genTxWithUTxO = undefined
+shrink :: IsRecentEra era => TxWithUTxO era -> [TxWithUTxO era]
+shrink (TxWithUTxO tx utxo) = do
+    utxoSubset <- UTxO <$> submapOf (unUTxO utxo)
+    TxWithUTxO.constructFiltered tx utxoSubset
+  where
+    submapOf :: Map k v -> [Map k v]
+    submapOf m = shrinkMapBy Map.fromList Map.toList sublistof