From f4b42a50c8082e3f81970d13e3015ddf770b9427 Mon Sep 17 00:00:00 2001
From: tadeubas <tadeubas@gmail.com>
Date: Wed, 16 Apr 2025 23:35:11 -0300
Subject: [PATCH] fix Descriptor.num_branches for miniscript

---
 src/embit/descriptor/descriptor.py |  3 +++
 tests/tests/test_descriptor.py     | 35 +++++++++++++++++++++++++-----
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/src/embit/descriptor/descriptor.py b/src/embit/descriptor/descriptor.py
index 9f585dc..6b1a22f 100644
--- a/src/embit/descriptor/descriptor.py
+++ b/src/embit/descriptor/descriptor.py
@@ -61,6 +61,9 @@ def script_len(self):
 
     @property
     def num_branches(self):
+        if self.miniscript is not None:
+            return max({k.num_branches for k in self.miniscript.keys})
+
         return max([k.num_branches for k in self.keys])
 
     def branch(self, branch_index=None):
diff --git a/tests/tests/test_descriptor.py b/tests/tests/test_descriptor.py
index de3e4a6..bdba715 100644
--- a/tests/tests/test_descriptor.py
+++ b/tests/tests/test_descriptor.py
@@ -34,6 +34,7 @@ def test_descriptors(self):
                 "c:pk_k(0250863ad64a87ae8a2fe83c1af1a8403cb53f53e486d8511dad8a04887e5b2352)"
                 "))",
                 "21020e0338c96a8870479f2396c373cc7696ba124e8635d41b0ea581112b67817261ac7364210250863ad64a87ae8a2fe83c1af1a8403cb53f53e486d8511dad8a04887e5b2352ac68",
+                1,
             ),
             # # pkh - 8e5d7457d33a978d1c3c1e440f92a195e00cc7d8
             # ("wsh(v:pk_h(03e7d285b4817f83f724cd29394da75dfc84fe639ed147a944e7e6064703b14130))", None),
@@ -41,36 +42,44 @@ def test_descriptors(self):
                 "sh(wsh(and_v(or_c(pk(%s),or_c(pk(%s),v:older(1000))),pk(%s))))"
                 % tuple(keys[-3:]),
                 "2103e7d285b4817f83f724cd29394da75dfc84fe639ed147a944e7e6064703b14130ac642103b8fa5d5959fa4027ccbf0736a86ccde4242e3051ea363437b4ff0d52598d7cecac6402e803b26968682103e7d285b4817f83f724cd29394da75dfc84fe639ed147a944e7e6064703b14130ac",
+                2,
             ),
             (
                 "sh(or_b(pk(%s),s:pk(%s)))" % tuple(keys[:2]),
                 "2103801b3a4e3ca0d61d469445621561c47f6c1424d0fd353a44c2c3ebb84ae78f59ac7c2103e7d285b4817f83f724cd29394da75dfc84fe639ed147a944e7e6064703b14130ac9b",
+                2,
             ),
             (
                 "wsh(or_d(pk(%s),pkh(%s)))" % tuple(keys[-2:]),
                 "2103b8fa5d5959fa4027ccbf0736a86ccde4242e3051ea363437b4ff0d52598d7cecac736476a9148e5d7457d33a978d1c3c1e440f92a195e00cc7d888ac68",
+                2,
             ),
             (
                 "wsh(and_v(v:pk(%s),or_d(pk(%s),older(12960))))" % tuple(keys[:2]),
                 "2103801b3a4e3ca0d61d469445621561c47f6c1424d0fd353a44c2c3ebb84ae78f59ad2103e7d285b4817f83f724cd29394da75dfc84fe639ed147a944e7e6064703b14130ac736402a032b268",
+                2,
             ),
             (
                 "wsh(andor(pk(%s),older(1008),pk(%s)))" % tuple(keys[:2]),
                 "2103801b3a4e3ca0d61d469445621561c47f6c1424d0fd353a44c2c3ebb84ae78f59ac642103e7d285b4817f83f724cd29394da75dfc84fe639ed147a944e7e6064703b14130ac6702f003b268",
+                2,
             ),
             (
                 "wsh(t:or_c(pk(%s),and_v(v:pk(%s),or_c(pk(%s),v:hash160(e7d285b4817f83f724cd29394da75dfc84fe639e)))))"
                 % tuple(keys[:3]),
                 "2103801b3a4e3ca0d61d469445621561c47f6c1424d0fd353a44c2c3ebb84ae78f59ac642103e7d285b4817f83f724cd29394da75dfc84fe639ed147a944e7e6064703b14130ad2103b8fa5d5959fa4027ccbf0736a86ccde4242e3051ea363437b4ff0d52598d7cecac6482012088a914e7d285b4817f83f724cd29394da75dfc84fe639e88686851",
+                2,
             ),
             (
                 "wsh(andor(pk(%s),or_i(and_v(v:pkh(%s),hash160(e7d285b4817f83f724cd29394da75dfc84fe639e)),older(1008)),pk(%s)))"
                 % tuple(keys[:3]),
                 "2103801b3a4e3ca0d61d469445621561c47f6c1424d0fd353a44c2c3ebb84ae78f59ac642103b8fa5d5959fa4027ccbf0736a86ccde4242e3051ea363437b4ff0d52598d7cecac676376a9148e5d7457d33a978d1c3c1e440f92a195e00cc7d888ad82012088a914e7d285b4817f83f724cd29394da75dfc84fe639e876702f003b26868",
+                2,
             ),
             (
                 "wsh(multi(2,%s,%s,%s))" % tuple(keys[:3]),
                 "522103801b3a4e3ca0d61d469445621561c47f6c1424d0fd353a44c2c3ebb84ae78f592103e7d285b4817f83f724cd29394da75dfc84fe639ed147a944e7e6064703b141302103b8fa5d5959fa4027ccbf0736a86ccde4242e3051ea363437b4ff0d52598d7cec53ae",
+                2,
             ),
             # TODO: invalid miniscript for segwit, but valid for taproot
             # ("wsh(thresh(3,pk(%s),s:pk(%s),s:pk(%s),sdv:older(12960)))" % tuple(keys[:3]),
@@ -88,6 +97,7 @@ def test_descriptors(self):
                 "038c8f919f70062c084376223fd8b4f0c08958e70499df496411dde83a1bb64b0d,"
                 "02d0ea7084e344b56625277b074d15a15301b9d96b0b2dd9fc905e01fc3de408e1))",
                 "5a210373b665b6fe153c5872de1344339ee60588491257d2c34567aa026af237143a6c2102916ee61974fc4892afb2d3cad4c13472138b5521411de24a78910afb97b95f22210244efc096ea3b7df99071b1cfa1630144e20d8ccd1540e726034a051aa1802d3b2102d9c51dc3f4088d5ce0b83f188fb14901b98c1c9e8cf771c49b7b441e56272b8a2103094990a34af21ef3ed766c8e0cb1e44f5e0d80412bbe00a2ade82a024ca91d232102722a386ad0f6d7f1261808a3e70fab143303bd2264283486411c3183ea3ed1c321036070b1f2995d8ffda8478ef55affd39795689a3982d54b12180397b1ad1f5f7521026515fa7603c10c44f6d316ae7592b5899d46d87ac1e574ec53de8b59f95efad621038c8f919f70062c084376223fd8b4f0c08958e70499df496411dde83a1bb64b0d2102d0ea7084e344b56625277b074d15a15301b9d96b0b2dd9fc905e01fc3de408e15aae",
+                1,
             ),
             (
                 "wsh(andor("
@@ -113,17 +123,24 @@ def test_descriptors(self):
                 "a:pkh(379ed952eb4740386acc59c2d28d9aa62e63968d),"
                 "a:pkh(c30d2795e70b1ee6f8af0b33d9460d60cfcf10b3))))",
                 "5421036070b1f2995d8ffda8478ef55affd39795689a3982d54b12180397b1ad1f5f7521026515fa7603c10c44f6d316ae7592b5899d46d87ac1e574ec53de8b59f95efad621038c8f919f70062c084376223fd8b4f0c08958e70499df496411dde83a1bb64b0d2102d0ea7084e344b56625277b074d15a15301b9d96b0b2dd9fc905e01fc3de408e154ae6476a9141ad3ca2d247b8e8888e41f89ac8bef217d83f33f88ac6b76a914f94f2eadc9c1bc3a8b8c2c6364af2c070fd4120688ac6c936b76a9143c306c2c97e4ba62ac0d7fb3965aba66b28e895988ac6c936b76a914ba7b9e846eb6b16420976c6bead54d9bb2b08d3588ac6c936b76a914379ed952eb4740386acc59c2d28d9aa62e63968d88ac6c936b76a914c30d2795e70b1ee6f8af0b33d9460d60cfcf10b388ac6c93558767562103856d447f1b890cc6e0e0114cd5bac58662c37ce7f458c458b72bd396597edfc72103e080e99896384aa8a07da837b2042a4c0d824eeaa8d51e6c9cff20682be75d4f2102c6d258e728005d4d00e55ac4b87786df507921b3ba3efec244a47f4a2e61b4b02102edfc1d6088f9b6470ed4550d8bf2326ebebc0464a7f78581fa7283fc54edecf02102f3630d1f51b2ebaaf1c7ebae9c24318279d4cff5ad16cb290b6d26edf96dca9c210353ecc8e7b1cc90d405cd6fc9d9f24d44b6b5649abc2773f28a6ca4fa7a4cd62956af029000b268",
+                1,
             ),
             (
                 "wsh(sortedmulti(2,%s,%s,%s))" % tuple(keys[:3]),
                 "522103801b3a4e3ca0d61d469445621561c47f6c1424d0fd353a44c2c3ebb84ae78f592103b8fa5d5959fa4027ccbf0736a86ccde4242e3051ea363437b4ff0d52598d7cec2103e7d285b4817f83f724cd29394da75dfc84fe639ed147a944e7e6064703b1413053ae",
+                2,
             ),
-            ("wpkh(%s)" % keys[0], "0014f8f93df2160de8fd3ca716e2f905c74da3f9839f"),
-            ("sh(wpkh(%s))" % keys[0], "0014f8f93df2160de8fd3ca716e2f905c74da3f9839f"),
-            ("pkh(%s)" % keys[0], "76a914f8f93df2160de8fd3ca716e2f905c74da3f9839f88ac"),
+            (
+                "wsh(sortedmulti(2,%s,%s,%s))" % tuple(key.replace("<0;1>", "0") for key in keys[:3]),
+                "522103801b3a4e3ca0d61d469445621561c47f6c1424d0fd353a44c2c3ebb84ae78f592103b8fa5d5959fa4027ccbf0736a86ccde4242e3051ea363437b4ff0d52598d7cec2103e7d285b4817f83f724cd29394da75dfc84fe639ed147a944e7e6064703b1413053ae",
+                1,
+            ),
+            ("wpkh(%s)" % keys[0], "0014f8f93df2160de8fd3ca716e2f905c74da3f9839f", 2,),
+            ("sh(wpkh(%s))" % keys[0], "0014f8f93df2160de8fd3ca716e2f905c74da3f9839f", 2,),
+            ("pkh(%s)" % keys[0], "76a914f8f93df2160de8fd3ca716e2f905c74da3f9839f88ac", 2,),
         ]
 
-        for i, (d, a) in enumerate(dd):
+        for d, a, n_branches in dd:
             sc = Descriptor.from_string(d)
             self.assertEqual(str(sc), d)
             # get top level script
@@ -132,6 +149,7 @@ def test_descriptors(self):
             schex = hexlify(scc.data).decode()
             self.assertEqual(schex, a)
             self.assertEqual(str(sc), d)
+            assert sc.num_branches == n_branches
 
     def test_keys(self):
         keys = [
@@ -222,32 +240,39 @@ def test_multisig(self):
                 "wsh(c:andor(multi(1,%s,%s),pk_k(%s),pk_k(%s)))" % keys,
                 False,
                 False,
+                1
             ),
             (
                 "wsh(multi(2,%s,%s,%s,%s))" % keys,
                 True,
                 False,
+                1
             ),
             (
                 "wsh(sortedmulti(2,%s,%s,%s,%s))" % keys,
                 True,
                 True,
+                1
             ),
             (
                 "tr(%s,multi_a(2,%s,%s,%s))" % keys,
                 False,
                 False,
+                1
             ),
             (
                 "tr(%s,sortedmulti_a(2,%s,%s,%s))" % keys,
                 False,
                 False,
+                1
             ),
         ]
-        for dstr, is_basic, is_sorted in descriptors:
+        for dstr, is_basic, is_sorted, n_branches in descriptors:
             d = Descriptor.from_string(dstr)
             self.assertEqual(d.is_basic_multisig, is_basic)
             self.assertEqual(d.is_sorted, is_sorted)
+            
+            assert n_branches == d.num_branches
 
 
 # test that: