From 8834b7f992979143e0e6e27dea9a5b79d2e34f03 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BF=97=E5=AE=87?= <hello@evanlinjin.me>
Date: Thu, 6 Oct 2022 13:42:06 +0800
Subject: [PATCH] Fix `max_satisfaction_weight` calculations for taproot

* Key spend path should take into consideration the scriptSigLen,
  stackLen and itemLen (for the one and only stack item).
* Script spend path should consider `control_block` and `script` to be
  items on the stack.

I also restructured the function to improve readability.
---
 examples/taproot.rs  |  8 +++----
 src/descriptor/tr.rs | 52 ++++++++++++++++++++++++++------------------
 2 files changed, 35 insertions(+), 25 deletions(-)

diff --git a/examples/taproot.rs b/examples/taproot.rs
index 8af46f0c2..d9068cfea 100644
--- a/examples/taproot.rs
+++ b/examples/taproot.rs
@@ -116,11 +116,11 @@ fn main() {
 
     // Max Satisfaction Weight for compilation, corresponding to the script-path spend
     // `multi_a(2,PUBKEY_1,PUBKEY_2) at taptree depth 1, having
-    // Max Witness Size = scriptSig len + control_block size + varint(script_size) + script_size +
-    //                     varint(max satisfaction elements) + max satisfaction size
-    //                  = 4 + 65 + 1 + 70 + 1 + 132
+    // Max Witness Size = scriptSig len + witnessStack len + varint(control_block_size) +
+    //                    control_block size + varint(script_size) + script_size + max_satisfaction_size
+    //                  = 4 + 1 + 1 + 65 + 1 + 70 + 132 = 274
     let max_sat_wt = real_desc.max_satisfaction_weight().unwrap();
-    assert_eq!(max_sat_wt, 273);
+    assert_eq!(max_sat_wt, 274);
 
     // Compute the bitcoin address and check if it matches
     let network = Network::Bitcoin;
diff --git a/src/descriptor/tr.rs b/src/descriptor/tr.rs
index f1ba69d01..ecfa2a9af 100644
--- a/src/descriptor/tr.rs
+++ b/src/descriptor/tr.rs
@@ -264,27 +264,37 @@ impl<Pk: MiniscriptKey> Tr<Pk> {
     /// # Errors
     /// When the descriptor is impossible to safisfy (ex: sh(OP_FALSE)).
     pub fn max_satisfaction_weight(&self) -> Result<usize, Error> {
-        let mut max_wieght = Some(65);
-        for (depth, ms) in self.iter_scripts() {
-            let script_size = ms.script_size();
-            let max_sat_elems = match ms.max_satisfaction_witness_elements() {
-                Ok(elem) => elem,
-                Err(..) => continue,
-            };
-            let max_sat_size = match ms.max_satisfaction_size() {
-                Ok(sz) => sz,
-                Err(..) => continue,
-            };
-            let control_block_sz = control_block_len(depth);
-            let wit_size = 4 + // scriptSig len byte
-            control_block_sz + // first element control block
-            varint_len(script_size) +
-            script_size + // second element script len with prefix
-            varint_len(max_sat_elems) +
-            max_sat_size; // witness
-            max_wieght = cmp::max(max_wieght, Some(wit_size));
-        }
-        max_wieght.ok_or(Error::ImpossibleSatisfaction)
+        let tree = match self.taptree() {
+            // key spend path:
+            // scriptSigLen(4) + stackLen(1) + stack[Sig]Len(1) + stack[Sig](65)
+            None => return Ok(4 + 1 + 1 + 65),
+            // script path spend..
+            Some(tree) => tree,
+        };
+
+        tree.iter()
+            .filter_map(|(depth, ms)| {
+                let script_size = ms.script_size();
+                let max_sat_elems = ms.max_satisfaction_witness_elements().ok()?;
+                let max_sat_size = ms.max_satisfaction_size().ok()?;
+                let control_block_size = control_block_len(depth);
+                Some(
+                    // scriptSig len byte
+                    4 +
+                    // witness field stack len (+2 for control block & script)
+                    varint_len(max_sat_elems + 2) +
+                    // size of elements to satisfy script
+                    max_sat_size +
+                    // second to last element: script
+                    varint_len(script_size) +
+                    script_size +
+                    // last element: control block
+                    varint_len(control_block_size) +
+                    control_block_size,
+                )
+            })
+            .max()
+            .ok_or(Error::ImpossibleSatisfaction)
     }
 }