Skip to content

Commit b55f478

Browse files
committed
Consolidated repetitive functions into util.rs
There were helper functions to compute nCk (k combinations of n given things) and product of K different combinations. Those were duplicated in mod.rs of descriptor and astelem.rs of miniscript. This commit de-duplicates them and store them in util.rs Signed-off-by: Harshil Jani <[email protected]>
1 parent 48346ca commit b55f478

File tree

3 files changed

+52
-89
lines changed

3 files changed

+52
-89
lines changed

src/descriptor/mod.rs

Lines changed: 7 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ use crate::miniscript::decode::Terminal;
2525
use crate::miniscript::{satisfy, Legacy, Miniscript, Segwitv0};
2626
use crate::plan::{AssetProvider, Assets, Plan};
2727
use crate::prelude::*;
28-
use crate::util::get_asset_combination;
28+
use crate::util::{asset_combination, k_of_n};
2929
use crate::{
3030
expression, hash256, BareCtx, Error, ForEachKey, MiniscriptKey, Satisfier, ToPublicKey,
3131
TranslateErr, TranslatePk, Translator,
@@ -559,23 +559,23 @@ impl Descriptor<DescriptorPublicKey> {
559559
WshInner::SortedMulti(k) => {
560560
let n = k.pks.len() as u64;
561561
let k = k.k as u64;
562-
Self::k_of_n(k, n)
562+
k_of_n(k, n)
563563
}
564564
WshInner::Ms(k) => k.count_assets(),
565565
},
566566
ShInner::Wpkh(_) => 1,
567567
ShInner::SortedMulti(k) => {
568568
let n = k.clone().pks.len() as u64;
569569
let k = k.clone().k as u64;
570-
Self::k_of_n(k, n)
570+
k_of_n(k, n)
571571
}
572572
ShInner::Ms(k) => k.count_assets(),
573573
},
574574
Descriptor::Wsh(k) => match k.as_inner() {
575575
WshInner::SortedMulti(k) => {
576576
let n = k.clone().pks.len() as u64;
577577
let k = k.clone().k as u64;
578-
Self::k_of_n(k, n)
578+
k_of_n(k, n)
579579
}
580580
WshInner::Ms(k) => k.count_assets(),
581581
},
@@ -612,7 +612,7 @@ impl Descriptor<DescriptorPublicKey> {
612612
WshInner::SortedMulti(k) => {
613613
let dpk_v = k.clone().pks;
614614
let k = k.clone().k;
615-
Ok(get_asset_combination(k, &dpk_v))
615+
Ok(asset_combination(k, &dpk_v))
616616
}
617617
WshInner::Ms(k) => Ok(k.all_assets()),
618618
},
@@ -624,15 +624,15 @@ impl Descriptor<DescriptorPublicKey> {
624624
ShInner::SortedMulti(k) => {
625625
let dpk_v = k.clone().pks;
626626
let k = k.clone().k;
627-
Ok(get_asset_combination(k, &dpk_v))
627+
Ok(asset_combination(k, &dpk_v))
628628
}
629629
ShInner::Ms(k) => Ok(k.all_assets()),
630630
},
631631
Descriptor::Wsh(k) => match k.as_inner() {
632632
WshInner::SortedMulti(k) => {
633633
let dpk_v = k.clone().pks;
634634
let k = k.clone().k;
635-
Ok(get_asset_combination(k, &dpk_v))
635+
Ok(asset_combination(k, &dpk_v))
636636
}
637637
WshInner::Ms(k) => {
638638
println!("{}", k);
@@ -655,14 +655,6 @@ impl Descriptor<DescriptorPublicKey> {
655655
}
656656
}
657657
}
658-
659-
// ways to select k things out of n
660-
fn k_of_n(k: u64, n: u64) -> u64 {
661-
if k == 0 || k == n {
662-
return 1;
663-
}
664-
Self::k_of_n(k - 1, n - 1) + Self::k_of_n(k - 1, n)
665-
}
666658
}
667659

668660
impl<P, Q> TranslatePk<P, Q> for Descriptor<P>

src/miniscript/astelem.rs

Lines changed: 10 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ use crate::miniscript::types::{self, Property};
1919
use crate::miniscript::ScriptContext;
2020
use crate::plan::Assets;
2121
use crate::prelude::*;
22-
use crate::util::{get_asset_combination, MsKeyBuilder};
22+
use crate::util::{asset_combination, get_combinations_product, k_of_n, MsKeyBuilder};
2323
use crate::{
2424
errstr, expression, AbsLockTime, DescriptorPublicKey, Error, Miniscript, MiniscriptKey,
2525
Terminal, ToPublicKey,
@@ -605,7 +605,6 @@ impl<Ctx: ScriptContext> Terminal<DescriptorPublicKey, Ctx> {
605605
Terminal::PkK(_) => 1,
606606
Terminal::PkH(_) => 1,
607607
Terminal::RawPkH(_) => 1,
608-
// What happens to timelocks ? for both the assets and the count.
609608
Terminal::After(_) => 0,
610609
Terminal::Older(_) => 0,
611610
Terminal::Sha256(_) => 1,
@@ -619,12 +618,7 @@ impl<Ctx: ScriptContext> Terminal<DescriptorPublicKey, Ctx> {
619618
Terminal::Verify(k) => k.count_assets(),
620619
Terminal::NonZero(k) => k.count_assets(),
621620
Terminal::ZeroNotEqual(k) => k.count_assets(),
622-
Terminal::AndV(left, right) => {
623-
let left_count = left.count_assets();
624-
let right_count = right.count_assets();
625-
left_count * right_count
626-
}
627-
Terminal::AndB(left, right) => {
621+
Terminal::AndV(left, right) | Terminal::AndB(left, right) => {
628622
let left_count = left.count_assets();
629623
let right_count = right.count_assets();
630624
left_count * right_count
@@ -635,51 +629,30 @@ impl<Ctx: ScriptContext> Terminal<DescriptorPublicKey, Ctx> {
635629
let c = c.count_assets();
636630
(a * b) + c
637631
}
638-
Terminal::OrB(left, right) => {
639-
let left_count = left.count_assets();
640-
let right_count = right.count_assets();
641-
left_count + right_count
642-
}
643-
Terminal::OrD(left, right) => {
644-
let left_count = left.count_assets();
645-
let right_count = right.count_assets();
646-
left_count + right_count
647-
}
648-
Terminal::OrC(left, right) => {
649-
let left_count = left.count_assets();
650-
let right_count = right.count_assets();
651-
left_count + right_count
652-
}
653-
Terminal::OrI(left, right) => {
632+
Terminal::OrB(left, right)
633+
| Terminal::OrC(left, right)
634+
| Terminal::OrD(left, right)
635+
| Terminal::OrI(left, right) => {
654636
let left_count = left.count_assets();
655637
let right_count = right.count_assets();
656638
left_count + right_count
657639
}
658640
Terminal::Thresh(k, ms_v) => {
659-
// k = 2, n = ms_v.len()
660-
// ms_v = [ms(A),ms(B),ms(C)];
661-
// Assume count array as [5,7,8] and k=2
662-
// get_combinations_product gives [5*7,5*8,7*8] = [35,40,56]
663641
let mut count_array = Vec::new();
664642
for ms in ms_v {
665643
count_array.push(ms.count_assets());
666644
}
667-
let products = Self::get_combinations_product(&count_array, *k as u64);
645+
let products = get_combinations_product(&count_array, *k as u64);
668646
let mut total_count: u64 = 0;
669647
for product in products {
670648
total_count += product;
671649
}
672650
total_count
673651
}
674-
Terminal::Multi(k, dpk) => {
652+
Terminal::Multi(k, dpk) | Terminal::MultiA(k, dpk) => {
675653
let k: u64 = *k as u64;
676654
let n: u64 = dpk.len() as u64;
677-
Self::k_of_n(k, n)
678-
}
679-
Terminal::MultiA(k, dpk) => {
680-
let k: u64 = *k as u64;
681-
let n: u64 = dpk.len() as u64;
682-
Self::k_of_n(k, n)
655+
k_of_n(k, n)
683656
}
684657
}
685658
}
@@ -823,9 +796,7 @@ impl<Ctx: ScriptContext> Terminal<DescriptorPublicKey, Ctx> {
823796
}
824797
result
825798
}
826-
Terminal::Multi(k, dpk_v) | Terminal::MultiA(k, dpk_v) => {
827-
get_asset_combination(*k, dpk_v)
828-
}
799+
Terminal::Multi(k, dpk_v) | Terminal::MultiA(k, dpk_v) => asset_combination(*k, dpk_v),
829800
}
830801
}
831802

@@ -858,38 +829,4 @@ impl<Ctx: ScriptContext> Terminal<DescriptorPublicKey, Ctx> {
858829
current_combination.truncate(current_combination.len() - 1);
859830
}
860831
}
861-
862-
// Do product of K combinations
863-
fn get_combinations_product(values: &[u64], k: u64) -> Vec<u64> {
864-
let mut products = Vec::new();
865-
let n = values.len();
866-
867-
if k == 0 {
868-
return vec![1]; // Empty combination has a product of 1
869-
}
870-
871-
// Using bitwise operations to generate combinations
872-
let max_combinations = 1u32 << n;
873-
for combination_bits in 1..max_combinations {
874-
if combination_bits.count_ones() as usize == k as usize {
875-
let mut product = 1;
876-
for i in 0..n {
877-
if combination_bits & (1u32 << i) != 0 {
878-
product *= values[i];
879-
}
880-
}
881-
products.push(product);
882-
}
883-
}
884-
885-
products
886-
}
887-
888-
// ways to select k things out of n
889-
fn k_of_n(k: u64, n: u64) -> u64 {
890-
if k == 0 || k == n {
891-
return 1;
892-
}
893-
Self::k_of_n(k - 1, n - 1) + Self::k_of_n(k, n - 1)
894-
}
895832
}

src/util.rs

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ impl MsKeyBuilder for script::Builder {
104104
}
105105

106106
// Helper to get all possible pairs of K of N assets
107-
pub fn get_asset_combination(k: usize, dpk_v: &Vec<DescriptorPublicKey>) -> Vec<Assets> {
107+
pub fn asset_combination(k: usize, dpk_v: &Vec<DescriptorPublicKey>) -> Vec<Assets> {
108108
let mut all_assets: Vec<Assets> = Vec::new();
109109
let current_assets = Assets::new();
110110
combine_assets(k, dpk_v, 0, current_assets, &mut all_assets);
@@ -132,3 +132,37 @@ pub fn combine_assets(
132132
println!("{:#?}", new_asset);
133133
combine_assets(k - 1, dpk_v, index + 1, new_asset, all_assets)
134134
}
135+
136+
// Do product of K combinations
137+
pub fn get_combinations_product(values: &[u64], k: u64) -> Vec<u64> {
138+
let mut products = Vec::new();
139+
let n = values.len();
140+
141+
if k == 0 {
142+
return vec![1]; // Empty combination has a product of 1
143+
}
144+
145+
// Using bitwise operations to generate combinations
146+
let max_combinations = 1u32 << n;
147+
for combination_bits in 1..max_combinations {
148+
if combination_bits.count_ones() as usize == k as usize {
149+
let mut product = 1;
150+
for i in 0..n {
151+
if combination_bits & (1u32 << i) != 0 {
152+
product *= values[i];
153+
}
154+
}
155+
products.push(product);
156+
}
157+
}
158+
159+
products
160+
}
161+
162+
// ways to select k things out of n
163+
pub fn k_of_n(k: u64, n: u64) -> u64 {
164+
if k == 0 || k == n {
165+
return 1;
166+
}
167+
k_of_n(k - 1, n - 1) + k_of_n(k, n - 1)
168+
}

0 commit comments

Comments
 (0)