From 261c8d8ae659d98fd3f63228f16d8f2eb4010b7e Mon Sep 17 00:00:00 2001
From: yancy <yancy@yancy.lol>
Date: Mon, 25 Nov 2024 11:07:24 -0600
Subject: [PATCH] Add total_weight to Input Weight Prediction

---
 bitcoin/src/blockdata/transaction.rs | 33 ++++++++++++++--------------
 1 file changed, 17 insertions(+), 16 deletions(-)

diff --git a/bitcoin/src/blockdata/transaction.rs b/bitcoin/src/blockdata/transaction.rs
index b8526d072..0ed381ef4 100644
--- a/bitcoin/src/blockdata/transaction.rs
+++ b/bitcoin/src/blockdata/transaction.rs
@@ -851,19 +851,13 @@ where
     I: IntoIterator<Item = InputWeightPrediction>,
     O: IntoIterator<Item = usize>,
 {
-    // This fold() does three things:
-    // 1) Counts the inputs and returns the sum as `input_count`.
-    // 2) Sums all of the input weights and returns the sum as `partial_input_weight`
-    //    For every input: script_size * 4 + witness_size
-    //    Since script_size is non-witness data, it gets a 4x multiplier.
-    // 3) Counts the number of inputs that have a witness data and returns the count as
-    //    `num_inputs_with_witnesses`.
-    let (input_count, partial_input_weight, inputs_with_witnesses) = inputs.into_iter().fold(
+    // sum input_weights, input_count and count of inputs with witness data
+    let (input_count, input_weight, inputs_with_witnesses) = inputs.into_iter().fold(
         (0, 0, 0),
-        |(count, partial_input_weight, inputs_with_witnesses), prediction| {
+        |(count, input_weight, inputs_with_witnesses), prediction| {
             (
                 count + 1,
-                partial_input_weight + prediction.witness_weight().to_wu() as usize,
+                input_weight + prediction.total_weight().to_wu() as usize,
                 inputs_with_witnesses + (prediction.witness_size > 0) as usize,
             )
         },
@@ -882,7 +876,7 @@ where
     );
     predict_weight_internal(
         input_count,
-        partial_input_weight,
+        input_weight,
         inputs_with_witnesses,
         output_count,
         output_scripts_size,
@@ -891,15 +885,11 @@ where
 
 const fn predict_weight_internal(
     input_count: usize,
-    partial_input_weight: usize,
+    input_weight: usize,
     inputs_with_witnesses: usize,
     output_count: usize,
     output_scripts_size: usize,
 ) -> Weight {
-    // Lengths of txid, index and sequence: (32, 4, 4).
-    // Multiply the lengths by 4 since the fields are all non-witness fields.
-    let input_weight = partial_input_weight + input_count * 4 * (32 + 4 + 4);
-
     // The value field of a TxOut is 8 bytes.
     let output_size = 8 * output_count + output_scripts_size;
     let non_input_size =
@@ -1147,6 +1137,17 @@ impl InputWeightPrediction {
         Self::witness_weight(self)
     }
 
+    /// Computes the signature, prevout (txid, index), and sequence weights of this weight
+    /// prediction.
+    ///
+    /// See also [`InputWeightPrediction::witness_weight`]
+    pub const fn total_weight(&self) -> Weight {
+        // `impl const Trait` is currently unavailable: rust/issues/67792
+        // Convert to u64s because we can't use `Add` in const context.
+        let weight = TX_IN_BASE_WEIGHT.to_wu() + Self::witness_weight(self).to_wu();
+        Weight::from_wu(weight)
+    }
+
     /// Computes the **signature weight** added to a transaction by an input with this weight prediction,
     /// not counting the prevout (txid, index), sequence, potential witness flag bytes or the witness count varint.
     ///