get_all_assets method for descriptors

Signed-off-by: Harshil Jani <harshiljani2002@gmail.com>

Author: Harshil-Jani

src/descriptor/mod.rs

@@ -23,7 +23,7 @@ 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::{
     expression, hash256, BareCtx, Error, ForEachKey, MiniscriptKey, Satisfier, ToPublicKey,
@@ -568,6 +568,99 @@ impl Descriptor<DefiniteDescriptorKey> {
     }
 }
 
+impl Descriptor<DescriptorPublicKey> {
+    /// Get all possible assets for a given descriptor
+    pub fn get_all_assets(&self) -> Result<Vec<Assets>, Error> {
+        match self {
+            Descriptor::Bare(k) => Ok(k.as_inner().get_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(Self::get_asset_combination(k, &dpk_v))
+                    }
+                    WshInner::Ms(k) => Ok(k.get_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(Self::get_asset_combination(k, &dpk_v))
+                }
+                ShInner::Ms(k) => Ok(k.get_all_assets()),
+            },
+            Descriptor::Wsh(k) => match k.as_inner() {
+                WshInner::SortedMulti(k) => {
+                    let dpk_v = k.clone().pks;
+                    let k = k.clone().k;
+                    Ok(Self::get_asset_combination(k, &dpk_v))
+                }
+                WshInner::Ms(k) => {
+                    println!("{}", k);
+                    let a = k.get_all_assets();
+                    println!("{:#?}", a);
+                    Ok(a)
+                }
+            },
+            Descriptor::Tr(k) => {
+                let s = k.taptree().clone().unwrap();
+                match s {
+                    TapTree::Tree(ref left, ref right) => {
+                        let mut a = left.get_all_assets();
+                        let b = right.get_all_assets();
+                        a.extend(b);
+                        Ok(a)
+                    }
+                    TapTree::Leaf(k) => Ok(k.get_all_assets()),
+                }
+            }
+        }
+    }
+
+    fn get_asset_combination(k: usize, dpk_v: &Vec<DescriptorPublicKey>) -> Vec<Assets> {
+        let mut all_assets: Vec<Assets> = Vec::new();
+        let current_assets = Assets::new();
+        Self::combine_assets(k, dpk_v, 0, current_assets, &mut all_assets);
+        all_assets
+    }
+
+    fn combine_assets(
+        k: usize,
+        dpk_v: &[DescriptorPublicKey],
+        index: usize,
+        current_assets: Assets,
+        all_assets: &mut Vec<Assets>,
+    ) {
+        if k == 0 {
+            all_assets.push(current_assets);
+            return;
+        }
+        if index >= dpk_v.len() {
+            return;
+        }
+        Self::combine_assets(k, dpk_v, index + 1, current_assets.clone(), all_assets);
+        let mut new_asset = current_assets;
+        new_asset = new_asset.add(dpk_v[index].clone());
+        println!("{:#?}", new_asset);
+        Self::combine_assets(k - 1, dpk_v, index + 1, new_asset, all_assets)
+    }
+}
+
 impl<P, Q> TranslatePk<P, Q> for Descriptor<P>
 where
     P: MiniscriptKey,
@@ -2094,4 +2187,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_get_all_assets_bare() {
+        let descriptor = Descriptor::<DescriptorPublicKey>::from_str(
+            "pk(0237b1c59ab055a8d3de40eaeb215c7b1922664b5ac049d849fb3346f81431e77f)",
+        )
+        .unwrap();
+
+        // Getting the assets from the get_all_assets method
+        let assets = descriptor.get_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_get_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.get_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_get_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.get_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

@@ -16,14 +16,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.
@@ -148,6 +148,21 @@ impl<Pk: MiniscriptKey> TapTree<Pk> {
     }
 }
 
+impl TapTree<DescriptorPublicKey> {
+    /// Get all possible assets for Taptree
+    pub fn get_all_assets(&self) -> Vec<Assets> {
+        match self {
+            TapTree::Tree(left, right) => {
+                let mut a = left.get_all_assets();
+                let b = right.get_all_assets();
+                a.extend(b);
+                a
+            }
+            TapTree::Leaf(k) => k.get_all_assets(),
+        }
+    }
+}
+
 impl<Pk: MiniscriptKey> fmt::Display for TapTree<Pk> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {