Skip to content

Commit e0e9fb2

Browse files
committed
get_all_assets method for descriptors
Signed-off-by: Harshil Jani <harshiljani2002@gmail.com>
1 parent 4c0ef36 commit e0e9fb2

File tree

2 files changed

+183
-4
lines changed

2 files changed

+183
-4
lines changed

src/descriptor/mod.rs

Lines changed: 165 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ use sync::Arc;
2323
use self::checksum::verify_checksum;
2424
use crate::miniscript::decode::Terminal;
2525
use crate::miniscript::{satisfy, Legacy, Miniscript, Segwitv0};
26-
use crate::plan::{AssetProvider, Plan};
26+
use crate::plan::{AssetProvider, Assets, Plan};
2727
use crate::prelude::*;
2828
use crate::{
2929
expression, hash256, BareCtx, Error, ForEachKey, MiniscriptKey, Satisfier, ToPublicKey,
@@ -568,6 +568,99 @@ impl Descriptor<DefiniteDescriptorKey> {
568568
}
569569
}
570570

571+
impl Descriptor<DescriptorPublicKey> {
572+
/// Get all possible assets for a given descriptor
573+
pub fn get_all_assets(&self) -> Result<Vec<Assets>, Error> {
574+
match self {
575+
Descriptor::Bare(k) => Ok(k.as_inner().get_all_assets()),
576+
Descriptor::Pkh(k) => {
577+
let mut asset = Assets::new();
578+
asset = asset.add(k.as_inner().clone());
579+
Ok(vec![asset])
580+
}
581+
Descriptor::Wpkh(k) => {
582+
let mut asset = Assets::new();
583+
asset = asset.add(k.as_inner().clone());
584+
Ok(vec![asset])
585+
}
586+
Descriptor::Sh(k) => match k.as_inner() {
587+
ShInner::Wsh(k) => match k.as_inner() {
588+
WshInner::SortedMulti(k) => {
589+
let dpk_v = k.clone().pks;
590+
let k = k.clone().k;
591+
Ok(Self::get_asset_combination(k, &dpk_v))
592+
}
593+
WshInner::Ms(k) => Ok(k.get_all_assets()),
594+
},
595+
ShInner::Wpkh(k) => {
596+
let mut asset = Assets::new();
597+
asset = asset.add(k.as_inner().clone());
598+
Ok(vec![asset])
599+
}
600+
ShInner::SortedMulti(k) => {
601+
let dpk_v = k.clone().pks;
602+
let k = k.clone().k;
603+
Ok(Self::get_asset_combination(k, &dpk_v))
604+
}
605+
ShInner::Ms(k) => Ok(k.get_all_assets()),
606+
},
607+
Descriptor::Wsh(k) => match k.as_inner() {
608+
WshInner::SortedMulti(k) => {
609+
let dpk_v = k.clone().pks;
610+
let k = k.clone().k;
611+
Ok(Self::get_asset_combination(k, &dpk_v))
612+
}
613+
WshInner::Ms(k) => {
614+
println!("{}", k);
615+
let a = k.get_all_assets();
616+
println!("{:#?}", a);
617+
Ok(a)
618+
}
619+
},
620+
Descriptor::Tr(k) => {
621+
let s = k.taptree().clone().unwrap();
622+
match s {
623+
TapTree::Tree(ref left, ref right) => {
624+
let mut a = left.get_all_assets();
625+
let b = right.get_all_assets();
626+
a.extend(b);
627+
Ok(a)
628+
}
629+
TapTree::Leaf(k) => Ok(k.get_all_assets()),
630+
}
631+
}
632+
}
633+
}
634+
635+
fn get_asset_combination(k: usize, dpk_v: &Vec<DescriptorPublicKey>) -> Vec<Assets> {
636+
let mut all_assets: Vec<Assets> = Vec::new();
637+
let current_assets = Assets::new();
638+
Self::combine_assets(k, dpk_v, 0, current_assets, &mut all_assets);
639+
all_assets
640+
}
641+
642+
fn combine_assets(
643+
k: usize,
644+
dpk_v: &[DescriptorPublicKey],
645+
index: usize,
646+
current_assets: Assets,
647+
all_assets: &mut Vec<Assets>,
648+
) {
649+
if k == 0 {
650+
all_assets.push(current_assets);
651+
return;
652+
}
653+
if index >= dpk_v.len() {
654+
return;
655+
}
656+
Self::combine_assets(k, dpk_v, index + 1, current_assets.clone(), all_assets);
657+
let mut new_asset = current_assets;
658+
new_asset = new_asset.add(dpk_v[index].clone());
659+
println!("{:#?}", new_asset);
660+
Self::combine_assets(k - 1, dpk_v, index + 1, new_asset, all_assets)
661+
}
662+
}
663+
571664
impl<P, Q> TranslatePk<P, Q> for Descriptor<P>
572665
where
573666
P: MiniscriptKey,
@@ -2094,4 +2187,75 @@ pk(03f28773c2d975288bc7d1d205c3748651b075fbc6610e58cddeeddf8f19405aa8))";
20942187
Desc::from_str(&format!("tr({},pk({}))", x_only_key, uncomp_key)).unwrap_err();
20952188
Desc::from_str(&format!("tr({},pk({}))", x_only_key, x_only_key)).unwrap();
20962189
}
2190+
2191+
#[test]
2192+
fn test_get_all_assets_bare() {
2193+
let descriptor = Descriptor::<DescriptorPublicKey>::from_str(
2194+
"pk(0237b1c59ab055a8d3de40eaeb215c7b1922664b5ac049d849fb3346f81431e77f)",
2195+
)
2196+
.unwrap();
2197+
2198+
// Getting the assets from the get_all_assets method
2199+
let assets = descriptor.get_all_assets().unwrap();
2200+
2201+
let mut expected_asset = Assets::new();
2202+
expected_asset = expected_asset.add(
2203+
DescriptorPublicKey::from_str(
2204+
"0237b1c59ab055a8d3de40eaeb215c7b1922664b5ac049d849fb3346f81431e77f",
2205+
)
2206+
.unwrap(),
2207+
);
2208+
assert_eq!(assets, vec![expected_asset]);
2209+
}
2210+
2211+
#[test]
2212+
fn test_get_all_assets_sh_sortedmulti() {
2213+
let keys = vec![
2214+
"0360eabc52e04f70c89e944f379895cdff28fed60849afe7736815c091765afb3c",
2215+
"03a80a24196e66ccf6bca6e6f96633bb629ec702acd76b074de10922b0cf41cc81",
2216+
];
2217+
2218+
let descriptor = Descriptor::<DescriptorPublicKey>::from_str(&format!(
2219+
"sh(sortedmulti(1,{},{}))",
2220+
keys[0], keys[1]
2221+
))
2222+
.unwrap();
2223+
2224+
let assets = descriptor.get_all_assets().unwrap();
2225+
2226+
let mut expected_assets: Vec<Assets> = Vec::new();
2227+
2228+
let mut asset1 = Assets::new();
2229+
asset1 = asset1.add(DescriptorPublicKey::from_str(keys[0]).unwrap());
2230+
expected_assets.push(asset1);
2231+
2232+
let mut asset2 = Assets::new();
2233+
asset2 = asset2.add(DescriptorPublicKey::from_str(keys[1]).unwrap());
2234+
expected_assets.push(asset2);
2235+
2236+
for expected_asset in &expected_assets {
2237+
assert!(assets.contains(expected_asset));
2238+
}
2239+
}
2240+
2241+
#[test]
2242+
fn test_get_all_assets_taproot() {
2243+
let x_only_key = bitcoin::key::XOnlyPublicKey::from_str(
2244+
"015e4cb53458bf813db8c79968e76e10d13ed6426a23fa71c2f41ba021c2a7ab",
2245+
)
2246+
.unwrap();
2247+
let descriptor =
2248+
Descriptor::from_str(&format!("tr({},pk({}))", x_only_key, x_only_key)).unwrap();
2249+
let assets = descriptor.get_all_assets().unwrap();
2250+
let mut expected_assets: Vec<Assets> = Vec::new();
2251+
let mut asset_1 = Assets::new();
2252+
asset_1 = asset_1.add(
2253+
DescriptorPublicKey::from_str(
2254+
"015e4cb53458bf813db8c79968e76e10d13ed6426a23fa71c2f41ba021c2a7ab",
2255+
)
2256+
.unwrap(),
2257+
);
2258+
expected_assets.push(asset_1);
2259+
assert_eq!(assets, expected_assets);
2260+
}
20972261
}

src/descriptor/tr.rs

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,14 @@ use crate::descriptor::DefiniteDescriptorKey;
1616
use crate::expression::{self, FromTree};
1717
use crate::miniscript::satisfy::{Placeholder, Satisfaction, SchnorrSigType, Witness};
1818
use crate::miniscript::Miniscript;
19-
use crate::plan::AssetProvider;
19+
use crate::plan::{AssetProvider, Assets};
2020
use crate::policy::semantic::Policy;
2121
use crate::policy::Liftable;
2222
use crate::prelude::*;
2323
use crate::util::{varint_len, witness_size};
2424
use crate::{
25-
errstr, Error, ForEachKey, MiniscriptKey, Satisfier, ScriptContext, Tap, ToPublicKey,
26-
TranslateErr, TranslatePk, Translator,
25+
errstr, DescriptorPublicKey, Error, ForEachKey, MiniscriptKey, Satisfier, ScriptContext, Tap,
26+
ToPublicKey, TranslateErr, TranslatePk, Translator,
2727
};
2828

2929
/// A Taproot Tree representation.
@@ -148,6 +148,21 @@ impl<Pk: MiniscriptKey> TapTree<Pk> {
148148
}
149149
}
150150

151+
impl TapTree<DescriptorPublicKey> {
152+
/// Get all possible assets for Taptree
153+
pub fn get_all_assets(&self) -> Vec<Assets> {
154+
match self {
155+
TapTree::Tree(left, right) => {
156+
let mut a = left.get_all_assets();
157+
let b = right.get_all_assets();
158+
a.extend(b);
159+
a
160+
}
161+
TapTree::Leaf(k) => k.get_all_assets(),
162+
}
163+
}
164+
}
165+
151166
impl<Pk: MiniscriptKey> fmt::Display for TapTree<Pk> {
152167
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
153168
match self {

0 commit comments

Comments
 (0)