Skip to content

Commit 6327521

Browse files
Harshil-JaniHarshil-Jani
committed
get_all_assets method into miniscript
Signed-off-by: Harshil Jani <harshiljani2002@gmail.com>
1 parent 5166d85 commit 6327521

File tree

3 files changed

+380
-4
lines changed

3 files changed

+380
-4
lines changed

‎src/miniscript/astelem.rs

Lines changed: 233 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,12 @@ use sync::Arc;
1818
use crate::miniscript::context::SigType;
1919
use crate::miniscript::types::{self, Property};
2020
use crate::miniscript::ScriptContext;
21+
use crate::plan::Assets;
2122
use crate::prelude::*;
2223
use crate::util::MsKeyBuilder;
2324
use crate::{
24-
errstr, expression, AbsLockTime, Error, Miniscript, MiniscriptKey, Terminal, ToPublicKey,
25+
errstr, expression, AbsLockTime, DescriptorPublicKey, Error, Miniscript, MiniscriptKey,
26+
Terminal, ToPublicKey,
2527
};
2628

2729
impl<Pk: MiniscriptKey, Ctx: ScriptContext> Terminal<Pk, Ctx> {
@@ -594,3 +596,233 @@ impl<Pk: MiniscriptKey, Ctx: ScriptContext> Terminal<Pk, Ctx> {
594596
}
595597
}
596598
}
599+
600+
impl<Ctx: ScriptContext> Terminal<DescriptorPublicKey, Ctx> {
601+
/// Retrieve the assets associated with the type of miniscript element.
602+
pub fn get_all_assets(&self) -> Vec<Assets> {
603+
match self {
604+
Terminal::True => vec![Assets::new()],
605+
Terminal::False => Vec::new(),
606+
Terminal::PkK(k) => {
607+
let mut asset = Assets::new();
608+
asset = asset.add(k.clone());
609+
vec![asset]
610+
}
611+
Terminal::PkH(k) => {
612+
let mut asset = Assets::new();
613+
asset = asset.add(k.clone());
614+
vec![asset]
615+
}
616+
Terminal::RawPkH(k) => {
617+
let mut asset = Assets::new();
618+
asset = asset.add(k.clone());
619+
vec![asset]
620+
}
621+
Terminal::After(k) => {
622+
let mut asset = Assets::new();
623+
asset.absolute_timelock = Some(k.clone().into());
624+
vec![asset]
625+
}
626+
Terminal::Older(k) => {
627+
let mut asset = Assets::new();
628+
asset.relative_timelock = Some(k.clone());
629+
vec![asset]
630+
}
631+
Terminal::Sha256(k) => {
632+
let mut asset = Assets::new();
633+
asset = asset.add(k.clone());
634+
vec![asset]
635+
}
636+
Terminal::Hash256(k) => {
637+
let mut asset = Assets::new();
638+
asset = asset.add(k.clone());
639+
vec![asset]
640+
}
641+
Terminal::Ripemd160(k) => {
642+
let mut asset = Assets::new();
643+
asset = asset.add(k.clone());
644+
vec![asset]
645+
}
646+
Terminal::Hash160(k) => {
647+
let mut asset = Assets::new();
648+
asset = asset.add(k.clone());
649+
vec![asset]
650+
}
651+
Terminal::Alt(k) => k.get_all_assets(),
652+
Terminal::Swap(k) => k.get_all_assets(),
653+
Terminal::Check(k) => k.get_all_assets(),
654+
Terminal::DupIf(k) => k.get_all_assets(),
655+
Terminal::Verify(k) => k.get_all_assets(),
656+
Terminal::NonZero(k) => k.get_all_assets(),
657+
Terminal::ZeroNotEqual(k) => k.get_all_assets(),
658+
Terminal::AndV(left, right) => {
659+
let a = left.get_all_assets();
660+
let b = right.get_all_assets();
661+
let result: Vec<Assets> = a
662+
.into_iter()
663+
.flat_map(|x| {
664+
b.clone().into_iter().map(move |y| {
665+
let mut new_asset = Assets::new();
666+
new_asset = new_asset.add(x.clone());
667+
new_asset = new_asset.add(y.clone());
668+
new_asset
669+
})
670+
})
671+
.collect();
672+
result
673+
}
674+
Terminal::AndB(left, right) => {
675+
let a = left.get_all_assets();
676+
let b = right.get_all_assets();
677+
let result: Vec<Assets> = a
678+
.into_iter()
679+
.flat_map(|x| {
680+
b.clone().into_iter().map(move |y| {
681+
let mut new_asset = Assets::new();
682+
new_asset = new_asset.add(x.clone());
683+
new_asset = new_asset.add(y.clone());
684+
new_asset
685+
})
686+
})
687+
.collect();
688+
result
689+
}
690+
Terminal::AndOr(a, b, c) => {
691+
let a = a.get_all_assets();
692+
let b = b.get_all_assets();
693+
let mut c = c.get_all_assets();
694+
let and: Vec<Assets> = a
695+
.into_iter()
696+
.flat_map(|x| {
697+
b.clone().into_iter().map(move |y| {
698+
let mut new_asset = Assets::new();
699+
new_asset = new_asset.add(x.clone());
700+
new_asset = new_asset.add(y.clone());
701+
new_asset
702+
})
703+
})
704+
.collect();
705+
c.extend(and);
706+
c
707+
}
708+
Terminal::OrB(left, right) => {
709+
let mut a = left.get_all_assets();
710+
let b = right.get_all_assets();
711+
a.extend(b);
712+
a
713+
}
714+
Terminal::OrD(left, right) => {
715+
let mut a = left.get_all_assets();
716+
let b = right.get_all_assets();
717+
a.extend(b);
718+
a
719+
}
720+
Terminal::OrC(left, right) => {
721+
let mut a = left.get_all_assets();
722+
let b = right.get_all_assets();
723+
a.extend(b);
724+
a
725+
}
726+
Terminal::OrI(left, right) => {
727+
let mut a = left.get_all_assets();
728+
let b = right.get_all_assets();
729+
a.extend(b);
730+
a
731+
}
732+
Terminal::Thresh(k, ms) => {
733+
let ms_v = Self::get_ms_combination_thresh(*k, ms);
734+
// k = 2
735+
// ms = [ms(A),ms(B),ms(C)];
736+
// ms_v = [[ms(A),ms(B)],[ms(A),ms(C)],[ms(B),ms(C)]]
737+
// Do ms_v[0] OR ms_v[1] OR ms_v[2]
738+
// Also Do ms_v[0][0] AND ms_v[0][1] and so on in the inner for loop
739+
740+
let mut result = Vec::new();
741+
for ms in ms_v {
742+
let mut and: Vec<Assets> = Vec::new();
743+
if let Some(first_assets) = ms.first() {
744+
and = first_assets.get_all_assets().clone();
745+
}
746+
for i in ms.iter().skip(1) {
747+
let i_assets = i.get_all_assets();
748+
and = and
749+
.iter()
750+
.flat_map(|x| {
751+
i_assets.iter().map(move |y| {
752+
let mut new_asset = x.clone();
753+
new_asset = new_asset.add(y.clone());
754+
new_asset
755+
})
756+
})
757+
.collect();
758+
}
759+
// OR of all combinations.
760+
result.extend(and.clone());
761+
}
762+
result
763+
}
764+
Terminal::Multi(k, dpk_v) => Self::get_asset_combination(*k, dpk_v),
765+
Terminal::MultiA(k, dpk_v) => Self::get_asset_combination(*k, dpk_v),
766+
}
767+
}
768+
769+
// Helper to get all possible pairs of K of N assets
770+
fn get_asset_combination(k: usize, dpk_v: &Vec<DescriptorPublicKey>) -> Vec<Assets> {
771+
let mut all_assets: Vec<Assets> = Vec::new();
772+
let current_assets = Assets::new();
773+
Self::combine_assets(k, dpk_v, 0, current_assets, &mut all_assets);
774+
all_assets
775+
}
776+
777+
// Combine K of N assets
778+
fn combine_assets(
779+
k: usize,
780+
dpk_v: &[DescriptorPublicKey],
781+
index: usize,
782+
current_assets: Assets,
783+
all_assets: &mut Vec<Assets>,
784+
) {
785+
if k == 0 {
786+
all_assets.push(current_assets);
787+
return;
788+
}
789+
if index >= dpk_v.len() {
790+
return;
791+
}
792+
Self::combine_assets(k, dpk_v, index + 1, current_assets.clone(), all_assets);
793+
let mut new_asset = current_assets;
794+
new_asset = new_asset.add(dpk_v[index].clone());
795+
println!("{:#?}", new_asset);
796+
Self::combine_assets(k - 1, dpk_v, index + 1, new_asset, all_assets)
797+
}
798+
799+
// Helper to get all combinations of K policies of N for thresh
800+
fn get_ms_combination_thresh(
801+
k: usize,
802+
ms: &Vec<Arc<Miniscript<DescriptorPublicKey, Ctx>>>,
803+
) -> Vec<Vec<Arc<Miniscript<DescriptorPublicKey, Ctx>>>> {
804+
let mut result = Vec::new();
805+
let mut current_combination = Vec::new();
806+
Self::combine_ms(0, &mut current_combination, &mut result, ms, k);
807+
result
808+
}
809+
810+
// combine K policies of N for thresh
811+
fn combine_ms(
812+
start: usize,
813+
current_combination: &mut Vec<Arc<Miniscript<DescriptorPublicKey, Ctx>>>,
814+
result: &mut Vec<Vec<Arc<Miniscript<DescriptorPublicKey, Ctx>>>>,
815+
ms: &Vec<Arc<Miniscript<DescriptorPublicKey, Ctx>>>,
816+
k: usize,
817+
) {
818+
if current_combination.len() == k {
819+
result.push(current_combination.clone());
820+
return;
821+
}
822+
for i in start..ms.len() {
823+
current_combination.push(ms[i].clone());
824+
Self::combine_ms(i + 1, current_combination, result, ms, k);
825+
current_combination.truncate(current_combination.len() - 1);
826+
}
827+
}
828+
}

0 commit comments

Comments
 (0)