Skip to content

Commit 4c0ef36

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

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
@@ -17,10 +17,12 @@ use sync::Arc;
1717
use crate::miniscript::context::SigType;
1818
use crate::miniscript::types::{self, Property};
1919
use crate::miniscript::ScriptContext;
20+
use crate::plan::Assets;
2021
use crate::prelude::*;
2122
use crate::util::MsKeyBuilder;
2223
use crate::{
23-
errstr, expression, AbsLockTime, Error, Miniscript, MiniscriptKey, Terminal, ToPublicKey,
24+
errstr, expression, AbsLockTime, DescriptorPublicKey, Error, Miniscript, MiniscriptKey,
25+
Terminal, ToPublicKey,
2426
};
2527

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

0 commit comments

Comments
 (0)