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.

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),