From 806a384bd17451e5863b71adc3bc398727426980 Mon Sep 17 00:00:00 2001
From: Pablo Lamela <pablo.lamela@iohk.io>
Date: Tue, 18 Feb 2025 14:17:37 +0100
Subject: [PATCH] Add test for estimating key witness count with simple script

---
 cardano-api/gen/Test/Gen/Cardano/Api/Typed.hs | 11 ++++--
 .../Test/Cardano/Api/TxBody.hs                | 35 +++++++++++++++++++
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/cardano-api/gen/Test/Gen/Cardano/Api/Typed.hs b/cardano-api/gen/Test/Gen/Cardano/Api/Typed.hs
index 1117d29e99..8e95784bd0 100644
--- a/cardano-api/gen/Test/Gen/Cardano/Api/Typed.hs
+++ b/cardano-api/gen/Test/Gen/Cardano/Api/Typed.hs
@@ -132,6 +132,7 @@ module Test.Gen.Cardano.Api.Typed
   , genProposals
   , genProposal
   , genVotingProcedures
+  , genSimpleScriptWithoutEmptyAnys
   )
 where
 
@@ -233,8 +234,14 @@ genScript SimpleScriptLanguage =
 genScript (PlutusScriptLanguage lang) =
   PlutusScript lang <$> genPlutusScript lang
 
+genSimpleScriptWithoutEmptyAnys :: Gen SimpleScript
+genSimpleScriptWithoutEmptyAnys = genRandomSimpleScript False
+
 genSimpleScript :: Gen SimpleScript
-genSimpleScript =
+genSimpleScript = genRandomSimpleScript True
+
+genRandomSimpleScript :: Bool -> Gen SimpleScript
+genRandomSimpleScript hasEmptyAnys=
   genTerm
  where
   genTerm = Gen.recursive Gen.choice nonRecursive recursive
@@ -249,7 +256,7 @@ genSimpleScript =
   -- Recursive generators
   recursive =
     [ RequireAllOf <$> Gen.list (Range.linear 0 10) genTerm
-    , RequireAnyOf <$> Gen.list (Range.linear 0 10) genTerm
+    , RequireAnyOf <$> Gen.list (Range.linear (if hasEmptyAnys then 0 else 1) 10) genTerm
     , do
         ts <- Gen.list (Range.linear 0 10) genTerm
         m <- Gen.integral (Range.constant 0 (length ts))
diff --git a/cardano-api/test/cardano-api-test/Test/Cardano/Api/TxBody.hs b/cardano-api/test/cardano-api-test/Test/Cardano/Api/TxBody.hs
index 5dc668c025..ad86fb447a 100644
--- a/cardano-api/test/cardano-api-test/Test/Cardano/Api/TxBody.hs
+++ b/cardano-api/test/cardano-api-test/Test/Cardano/Api/TxBody.hs
@@ -4,6 +4,7 @@
 {-# LANGUAGE OverloadedLists #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TupleSections #-}
 
 {- HLINT ignore "Use camelCase" -}
 
@@ -18,6 +19,8 @@ import qualified Cardano.Api.Ledger as L
 import           Cardano.Api.Shelley (ShelleyLedgerEra)
 
 import           Data.Maybe (isJust)
+import           Data.Set (Set)
+import qualified Data.Set as Set
 import           Data.Type.Equality (TestEquality (testEquality))
 import           GHC.Exts (IsList (..))
 
@@ -27,6 +30,7 @@ import           Test.Cardano.Api.Orphans ()
 
 import           Hedgehog (MonadTest, Property, (===))
 import qualified Hedgehog as H
+import           Hedgehog.Gen (shuffle)
 import           Test.Tasty (TestTree, testGroup)
 import           Test.Tasty.Hedgehog (testProperty)
 
@@ -108,6 +112,34 @@ prop_roundtrip_txbodycontent_conway_fields = H.property $ do
   getProposalProcedures TxProposalProceduresNone = Nothing
   getProposalProcedures (TxProposalProcedures pp) = Just $ fst <$> toList pp
 
+prop_simple_script_witness_count :: Property
+prop_simple_script_witness_count = H.property $ do
+  let sbe = ShelleyBasedEraConway
+  (_, contentWithoutScript) <- H.forAll $ genValidTxBody sbe
+  script <- H.forAll genSimpleScriptWithoutEmptyAnys
+  newTxIn <-
+    H.forAll $
+      (,BuildTxWith
+          ( ScriptWitness
+              ScriptWitnessForSpending
+              (SimpleScriptWitness SimpleScriptInConway (SScript script))
+          ))
+        <$> genTxIn
+  witList <- H.forAll $ satisfyScript script
+  let witCount = fromIntegral $ Set.size witList
+  H.diff
+    (estimateTransactionKeyWitnessCount contentWithoutScript + witCount)
+    (<=)
+    (estimateTransactionKeyWitnessCount (addTxIn newTxIn contentWithoutScript))
+ where
+  satisfyScript :: SimpleScript -> H.Gen (Set (Hash PaymentKey))
+  satisfyScript (RequireSignature paymentKeyHash) = return $ Set.singleton paymentKeyHash
+  satisfyScript (RequireTimeBefore _) = return mempty
+  satisfyScript (RequireTimeAfter _) = return mempty
+  satisfyScript (RequireAllOf simpleScripts) = Set.unions <$> traverse satisfyScript simpleScripts
+  satisfyScript (RequireMOf n simpleScripts) = shuffle simpleScripts >>= satisfyScript . RequireAllOf . take n
+  satisfyScript (RequireAnyOf simpleScripts) = satisfyScript (RequireMOf 1 simpleScripts)
+
 tests :: TestTree
 tests =
   testGroup
@@ -119,4 +151,7 @@ tests =
     , testProperty
         "roundtrip txbodycontent new conway fields"
         prop_roundtrip_txbodycontent_conway_fields
+    , testProperty
+        "simple script witness count"
+        prop_simple_script_witness_count
     ]