WIP

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

@@ -45,11 +45,10 @@ import Test.QuickCheck
     )
 import Test.QuickCheck.Extra
     ( genMapFromKeysWith
-    , genMapWith
+    , genNonEmptyDisjointMap
     , shrinkMapToSubmaps
     )
 
-import qualified Data.Map.Strict as Map
 import qualified Internal.Cardano.Write.Tx.TxWithUTxO as TxWithUTxO
 
 -- | Generates a 'TxWithUTxO' object.
@@ -86,7 +85,7 @@ generateWithMinimalUTxO genTx _genTxIn genTxOut = do
   where
     txInputs = view (bodyTxL . allInputsTxBodyF)
 
--- | Generates a 'TxWithUTxO' object that has a non-minimal UTxO set.
+-- | 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.
 --
@@ -99,14 +98,9 @@ generateWithSurplusUTxO
     -> Gen (TxWithUTxO era)
 generateWithSurplusUTxO genTx genTxIn genTxOut =
     generateWithMinimalUTxO genTx genTxIn genTxOut >>= \case
-        TxWithUTxO tx minimalUTxO -> do
-            surplusUTxO <- genNonEmptyUTxO
-            pure $ TxWithUTxO.constructFiltered tx (minimalUTxO <> surplusUTxO)
-  where
-    genNonEmptyUTxO :: Gen (UTxO era)
-    genNonEmptyUTxO
-        = fmap UTxO
-        $ Map.insert <$> genTxIn <*> genTxOut <*> genMapWith genTxIn genTxOut
+        TxWithUTxO tx (UTxO utxo) -> do
+            utxoSurplus <- genNonEmptyDisjointMap genTxIn genTxOut utxo
+            pure $ TxWithUTxO.constructFiltered tx $ UTxO (utxo <> utxoSurplus)
 
 -- TODO: Parameterise this by `Tx era -> [Tx era]` shrinker.
 shrink :: IsRecentEra era => TxWithUTxO era -> [TxWithUTxO era]