From 8e99191d85428bd6f819e55dc45fb46a0c9c31f2 Mon Sep 17 00:00:00 2001
From: rooooooooob <raftias@gmail.com>
Date: Tue, 10 Sep 2024 09:10:55 -0700
Subject: [PATCH] Ref script fee fixes

* Check everywhere that could use a ref script (not sure if everyone
  it's allowed)

* Duplicate size for different ref inptus with the same ref script as
  per Haskell code

* look in spent inputs too (as this is possible to ref a script from)

Possibly addresses #354
---
 chain/rust/src/builders/tx_builder.rs | 67 +++++++++++++++++++++++++--
 1 file changed, 62 insertions(+), 5 deletions(-)

diff --git a/chain/rust/src/builders/tx_builder.rs b/chain/rust/src/builders/tx_builder.rs
index 8c2cb01c..5bc8d649 100644
--- a/chain/rust/src/builders/tx_builder.rs
+++ b/chain/rust/src/builders/tx_builder.rs
@@ -240,24 +240,81 @@ fn min_fee_with_exunits(tx_builder: &TransactionBuilder) -> Result<Coin, TxBuild
         })
     }
 
-    let ref_script_orig_sizes: HashMap<ScriptHash, u64> =
+    // if the same script is included in multiple places its size is multiplied
+    let ref_script_orig_sizes_non_distinct: Vec<(ScriptHash, u64)> =
         if let Some(ref_inputs) = &tx_builder.reference_inputs {
-            ref_inputs
+            tx_builder
+                .inputs
                 .iter()
+                .chain(ref_inputs.iter())
                 .filter_map(ref_script_orig_size_builder)
                 .collect()
         } else {
-            HashMap::default()
+            tx_builder
+                .inputs
+                .iter()
+                .filter_map(ref_script_orig_size_builder)
+                .collect()
         };
+    let mut ref_script_orig_sizes: HashMap<ScriptHash, u64> = HashMap::default();
+    for (hash, size) in ref_script_orig_sizes_non_distinct {
+        *ref_script_orig_sizes.entry(hash).or_default() += size;
+    }
 
-    let mut total_ref_script_size = 0;
+    let mut used_ref_scripts: HashMap<ScriptHash, u64> = HashMap::default();
     for utxo in tx_builder.inputs.iter() {
         if let Some(Credential::Script { hash, .. }) = utxo.output.address().payment_cred() {
             if let Some(orig_size) = ref_script_orig_sizes.get(hash) {
-                total_ref_script_size += *orig_size;
+                used_ref_scripts.insert(*hash, *orig_size);
+            }
+        }
+    }
+    if let Some(certs) = &tx_builder.certs {
+        let mut cert_wits = RequiredWitnessSet::default();
+        for cert in certs.iter() {
+            cert_required_wits(cert, &mut cert_wits);
+        }
+        for hash in cert_wits.scripts {
+            if let Some(orig_size) = ref_script_orig_sizes.get(&hash) {
+                used_ref_scripts.insert(hash, *orig_size);
+            }
+        }
+    }
+    if let Some(proposals) = &tx_builder.proposals {
+        for proposal in proposals.iter() {
+            if let Some(hash) = proposal.gov_action.script_hash() {
+                if let Some(orig_size) = ref_script_orig_sizes.get(hash) {
+                    used_ref_scripts.insert(*hash, *orig_size);
+                }
+            }
+        }
+    }
+    if let Some(votes) = &tx_builder.votes {
+        for (voter, _) in votes.iter() {
+            if let Some(hash) = voter.script_hash() {
+                if let Some(orig_size) = ref_script_orig_sizes.get(hash) {
+                    used_ref_scripts.insert(*hash, *orig_size);
+                }
+            }
+        }
+    }
+    if let Some(withdrawals) = &tx_builder.withdrawals {
+        for (addr, _) in withdrawals.iter() {
+            if let Credential::Script { hash, .. } = &addr.payment {
+                if let Some(orig_size) = ref_script_orig_sizes.get(hash) {
+                    used_ref_scripts.insert(*hash, *orig_size);
+                }
+            }
+        }
+    }
+    if let Some(mint) = &tx_builder.mint {
+        for (policy, _) in mint.iter() {
+            if let Some(orig_size) = ref_script_orig_sizes.get(policy) {
+                used_ref_scripts.insert(*policy, *orig_size);
             }
         }
     }
+    let total_ref_script_size = used_ref_scripts.values().sum();
 
     crate::fees::min_fee(
         &full_tx,