Wrap all_assets for desc and add count_assets method

This commit wraps the all_assets function for all possible descriptor
types. It also adds `count_assets` function which can count the total
possible ways to obtain an asset for a given descriptor.

This API will help developer to count the possible ways. In-fact this
will also help in protecting that if some descriptor is very large and
has huge number of possible spend paths then it could simply be
dropped off and we can protect the network from traffic.

Signed-off-by: Harshil Jani <[email protected]>

Author: Harshil-Jani

src/descriptor/mod.rs

@@ -23,8 +23,9 @@ use sync::Arc;
 use self::checksum::verify_checksum;
 use crate::miniscript::decode::Terminal;
 use crate::miniscript::{satisfy, Legacy, Miniscript, Segwitv0};
-use crate::plan::{AssetProvider, Plan};
+use crate::plan::{AssetProvider, Assets, Plan};
 use crate::prelude::*;
+use crate::util::get_asset_combination;
 use crate::{
     expression, hash256, BareCtx, Error, ForEachKey, MiniscriptKey, Satisfier, ToPublicKey,
     TranslateErr, TranslatePk, Translator,
@@ -546,6 +547,124 @@ impl Descriptor<DefiniteDescriptorKey> {
     }
 }
 
+impl Descriptor<DescriptorPublicKey> {
+    /// Count total possible assets for a given descriptor.
+    pub fn count_assets(&self) -> u64 {
+        match self {
+            Descriptor::Bare(k) => k.as_inner().count_assets(),
+            Descriptor::Pkh(_) => 1,
+            Descriptor::Wpkh(_) => 1,
+            Descriptor::Sh(k) => match k.as_inner() {
+                ShInner::Wsh(k) => match k.as_inner() {
+                    WshInner::SortedMulti(k) => {
+                        let n = k.pks.len() as u64;
+                        let k = k.k as u64;
+                        Self::k_of_n(k, n)
+                    }
+                    WshInner::Ms(k) => k.count_assets(),
+                },
+                ShInner::Wpkh(_) => 1,
+                ShInner::SortedMulti(k) => {
+                    let n = k.clone().pks.len() as u64;
+                    let k = k.clone().k as u64;
+                    Self::k_of_n(k, n)
+                }
+                ShInner::Ms(k) => k.count_assets(),
+            },
+            Descriptor::Wsh(k) => match k.as_inner() {
+                WshInner::SortedMulti(k) => {
+                    let n = k.clone().pks.len() as u64;
+                    let k = k.clone().k as u64;
+                    Self::k_of_n(k, n)
+                }
+                WshInner::Ms(k) => k.count_assets(),
+            },
+            Descriptor::Tr(k) => {
+                let s = k.tap_tree().clone().unwrap();
+                match s {
+                    TapTree::Tree { left, right, height: _ } => {
+                        let a = left.count_assets();
+                        let b = right.count_assets();
+                        a + b
+                    }
+                    TapTree::Leaf(k) => k.count_assets(),
+                }
+            }
+        }
+    }
+
+    /// Get all possible assets for a given descriptor
+    pub fn all_assets(&self) -> Result<Vec<Assets>, Error> {
+        match self {
+            Descriptor::Bare(k) => Ok(k.as_inner().all_assets()),
+            Descriptor::Pkh(k) => {
+                let mut asset = Assets::new();
+                asset = asset.add(k.as_inner().clone());
+                Ok(vec![asset])
+            }
+            Descriptor::Wpkh(k) => {
+                let mut asset = Assets::new();
+                asset = asset.add(k.as_inner().clone());
+                Ok(vec![asset])
+            }
+            Descriptor::Sh(k) => match k.as_inner() {
+                ShInner::Wsh(k) => match k.as_inner() {
+                    WshInner::SortedMulti(k) => {
+                        let dpk_v = k.clone().pks;
+                        let k = k.clone().k;
+                        Ok(get_asset_combination(k, &dpk_v))
+                    }
+                    WshInner::Ms(k) => Ok(k.all_assets()),
+                },
+                ShInner::Wpkh(k) => {
+                    let mut asset = Assets::new();
+                    asset = asset.add(k.as_inner().clone());
+                    Ok(vec![asset])
+                }
+                ShInner::SortedMulti(k) => {
+                    let dpk_v = k.clone().pks;
+                    let k = k.clone().k;
+                    Ok(get_asset_combination(k, &dpk_v))
+                }
+                ShInner::Ms(k) => Ok(k.all_assets()),
+            },
+            Descriptor::Wsh(k) => match k.as_inner() {
+                WshInner::SortedMulti(k) => {
+                    let dpk_v = k.clone().pks;
+                    let k = k.clone().k;
+                    Ok(get_asset_combination(k, &dpk_v))
+                }
+                WshInner::Ms(k) => {
+                    println!("{}", k);
+                    let a = k.all_assets();
+                    println!("{:#?}", a);
+                    Ok(a)
+                }
+            },
+            Descriptor::Tr(k) => {
+                let s = k.tap_tree().clone().unwrap();
+                match s {
+                    TapTree::Tree { left, right, height: _ } => {
+                        let mut a = left.all_assets();
+                        let b = right.all_assets();
+                        a.extend(b);
+                        Ok(a)
+                    }
+                    TapTree::Leaf(k) => Ok(k.all_assets()),
+                }
+            }
+        }
+    }
+
+    // ways to select k things out of n
+    fn k_of_n(k: u64, n: u64) -> u64 {
+        if k == 0 || k == n {
+            return 1;
+        }
+        Self::k_of_n(k - 1, n - 1) + Self::k_of_n(k - 1, n)
+    }
+}
+
 impl<P, Q> TranslatePk<P, Q> for Descriptor<P>
 where
     P: MiniscriptKey,
@@ -2041,4 +2160,75 @@ pk(03f28773c2d975288bc7d1d205c3748651b075fbc6610e58cddeeddf8f19405aa8))";
         Desc::from_str(&format!("tr({},pk({}))", x_only_key, uncomp_key)).unwrap_err();
         Desc::from_str(&format!("tr({},pk({}))", x_only_key, x_only_key)).unwrap();
     }
+
+    #[test]
+    fn test_all_assets_bare() {
+        let descriptor = Descriptor::<DescriptorPublicKey>::from_str(
+            "pk(0237b1c59ab055a8d3de40eaeb215c7b1922664b5ac049d849fb3346f81431e77f)",
+        )
+        .unwrap();
+
+        // Getting the assets from the all_assets method
+        let assets = descriptor.all_assets().unwrap();
+
+        let mut expected_asset = Assets::new();
+        expected_asset = expected_asset.add(
+            DescriptorPublicKey::from_str(
+                "0237b1c59ab055a8d3de40eaeb215c7b1922664b5ac049d849fb3346f81431e77f",
+            )
+            .unwrap(),
+        );
+        assert_eq!(assets, vec![expected_asset]);
+    }
+
+    #[test]
+    fn test_all_assets_sh_sortedmulti() {
+        let keys = vec![
+            "0360eabc52e04f70c89e944f379895cdff28fed60849afe7736815c091765afb3c",
+            "03a80a24196e66ccf6bca6e6f96633bb629ec702acd76b074de10922b0cf41cc81",
+        ];
+
+        let descriptor = Descriptor::<DescriptorPublicKey>::from_str(&format!(
+            "sh(sortedmulti(1,{},{}))",
+            keys[0], keys[1]
+        ))
+        .unwrap();
+
+        let assets = descriptor.all_assets().unwrap();
+
+        let mut expected_assets: Vec<Assets> = Vec::new();
+
+        let mut asset1 = Assets::new();
+        asset1 = asset1.add(DescriptorPublicKey::from_str(keys[0]).unwrap());
+        expected_assets.push(asset1);
+
+        let mut asset2 = Assets::new();
+        asset2 = asset2.add(DescriptorPublicKey::from_str(keys[1]).unwrap());
+        expected_assets.push(asset2);
+
+        for expected_asset in &expected_assets {
+            assert!(assets.contains(expected_asset));
+        }
+    }
+
+    #[test]
+    fn test_all_assets_taproot() {
+        let x_only_key = bitcoin::key::XOnlyPublicKey::from_str(
+            "015e4cb53458bf813db8c79968e76e10d13ed6426a23fa71c2f41ba021c2a7ab",
+        )
+        .unwrap();
+        let descriptor =
+            Descriptor::from_str(&format!("tr({},pk({}))", x_only_key, x_only_key)).unwrap();
+        let assets = descriptor.all_assets().unwrap();
+        let mut expected_assets: Vec<Assets> = Vec::new();
+        let mut asset_1 = Assets::new();
+        asset_1 = asset_1.add(
+            DescriptorPublicKey::from_str(
+                "015e4cb53458bf813db8c79968e76e10d13ed6426a23fa71c2f41ba021c2a7ab",
+            )
+            .unwrap(),
+        );
+        expected_assets.push(asset_1);
+        assert_eq!(assets, expected_assets);
+    }
 }

src/descriptor/tr.rs

@@ -15,14 +15,14 @@ use crate::descriptor::DefiniteDescriptorKey;
 use crate::expression::{self, FromTree};
 use crate::miniscript::satisfy::{Placeholder, Satisfaction, SchnorrSigType, Witness};
 use crate::miniscript::Miniscript;
-use crate::plan::AssetProvider;
+use crate::plan::{AssetProvider, Assets};
 use crate::policy::semantic::Policy;
 use crate::policy::Liftable;
 use crate::prelude::*;
 use crate::util::{varint_len, witness_size};
 use crate::{
-    errstr, Error, ForEachKey, MiniscriptKey, Satisfier, ScriptContext, Tap, ToPublicKey,
-    TranslateErr, TranslatePk, Translator,
+    errstr, DescriptorPublicKey, Error, ForEachKey, MiniscriptKey, Satisfier, ScriptContext, Tap,
+    ToPublicKey, TranslateErr, TranslatePk, Translator,
 };
 
 /// A Taproot Tree representation.
@@ -153,6 +153,33 @@ impl<Pk: MiniscriptKey> TapTree<Pk> {
     }
 }
 
+impl TapTree<DescriptorPublicKey> {
+    /// Get all possible assets for Taptree
+    pub fn all_assets(&self) -> Vec<Assets> {
+        match self {
+            TapTree::Tree { left, right, height: _ } => {
+                let mut a = left.all_assets();
+                let b = right.all_assets();
+                a.extend(b);
+                a
+            }
+            TapTree::Leaf(k) => k.all_assets(),
+        }
+    }
+
+    /// Get total possible assets for TapTree
+    pub fn count_assets(&self) -> u64 {
+        match self {
+            TapTree::Tree { left, right, height: _ } => {
+                let a = left.count_assets();
+                let b = right.count_assets();
+                a + b
+            }
+            TapTree::Leaf(k) => k.count_assets(),
+        }
+    }
+}
+
 impl<Pk: MiniscriptKey> fmt::Display for TapTree<Pk> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {

src/lib.rs

@@ -501,6 +501,8 @@ pub enum Error {
     /// At least two BIP389 key expressions in the descriptor contain tuples of
     /// derivation indexes of different lengths.
     MultipathDescLenMismatch,
+    /// Cannot get assets for this large descriptor. Exceeds 1000 assets.
+    MaxAssetThresholdExceeded,
 }
 
 // https://github.com/sipa/miniscript/pull/5 for discussion on this number
@@ -577,6 +579,7 @@ impl fmt::Display for Error {
             Error::TrNoScriptCode => write!(f, "No script code for Tr descriptors"),
             Error::TrNoExplicitScript => write!(f, "No script code for Tr descriptors"),
             Error::MultipathDescLenMismatch => write!(f, "At least two BIP389 key expressions in the descriptor contain tuples of derivation indexes of different lengths"),
+            Error::MaxAssetThresholdExceeded => write!(f,"Cannot plan descriptors having more than 1000 possible spend paths."),
         }
     }
 }
@@ -619,6 +622,7 @@ impl error::Error for Error {
             | TrNoScriptCode
             | TrNoExplicitScript
             | MultipathDescLenMismatch => None,
+            MaxAssetThresholdExceeded => None,
             Script(e) => Some(e),
             AddrError(e) => Some(e),
             BadPubkey(e) => Some(e),