From 9a2d4f5a46578a6fa57b03726eb4e296a3ba3171 Mon Sep 17 00:00:00 2001 From: Harshil Jani Date: Tue, 1 Aug 2023 05:20:06 +0530 Subject: [PATCH] count_assets method for total ways to obtain assets Signed-off-by: Harshil Jani --- src/descriptor/mod.rs | 53 +++++++++++++++++ src/descriptor/tr.rs | 12 ++++ src/miniscript/astelem.rs | 121 ++++++++++++++++++++++++++++++++++++++ src/miniscript/mod.rs | 5 ++ 4 files changed, 191 insertions(+) diff --git a/src/descriptor/mod.rs b/src/descriptor/mod.rs index ec08d335e..35548ad68 100644 --- a/src/descriptor/mod.rs +++ b/src/descriptor/mod.rs @@ -569,6 +569,51 @@ impl Descriptor { } impl Descriptor { + /// Count total possible assets for a given descriptor. + pub fn count_assets(&self) -> u64 { + match self { + Descriptor::Bare(k) => k.as_inner().count_assets(), + Descriptor::Pkh(_) => 1, + Descriptor::Wpkh(_) => 1, + Descriptor::Sh(k) => match k.as_inner() { + ShInner::Wsh(k) => match k.as_inner() { + WshInner::SortedMulti(k) => { + let n = k.pks.len() as u64; + let k = k.k as u64; + Self::k_of_n(k, n) + } + WshInner::Ms(k) => k.count_assets(), + }, + ShInner::Wpkh(_) => 1, + ShInner::SortedMulti(k) => { + let n = k.clone().pks.len() as u64; + let k = k.clone().k as u64; + Self::k_of_n(k, n) + } + ShInner::Ms(k) => k.count_assets(), + }, + Descriptor::Wsh(k) => match k.as_inner() { + WshInner::SortedMulti(k) => { + let n = k.clone().pks.len() as u64; + let k = k.clone().k as u64; + Self::k_of_n(k, n) + } + WshInner::Ms(k) => k.count_assets(), + }, + Descriptor::Tr(k) => { + let s = k.taptree().clone().unwrap(); + match s { + TapTree::Tree(ref left, ref right) => { + let a = left.count_assets(); + let b = right.count_assets(); + a + b + } + TapTree::Leaf(k) => k.count_assets(), + } + } + } + } + /// Get all possible assets for a given descriptor pub fn get_all_assets(&self) -> Result, Error> { match self { @@ -659,6 +704,14 @@ impl Descriptor { println!("{:#?}", new_asset); Self::combine_assets(k - 1, dpk_v, index + 1, new_asset, all_assets) } + + // ways to select k things out of n + fn k_of_n(k: u64, n: u64) -> u64 { + if k == 0 || k == n { + return 1; + } + Self::k_of_n(k - 1, n - 1) + Self::k_of_n(k - 1, n) + } } impl TranslatePk for Descriptor

diff --git a/src/descriptor/tr.rs b/src/descriptor/tr.rs index 2954fe93f..7dbc2c72c 100644 --- a/src/descriptor/tr.rs +++ b/src/descriptor/tr.rs @@ -161,6 +161,18 @@ impl TapTree { TapTree::Leaf(k) => k.get_all_assets(), } } + + /// Get total possible assets for TapTree + pub fn count_assets(&self) -> u64 { + match self { + TapTree::Tree(left, right) => { + let a = left.count_assets(); + let b = right.count_assets(); + a + b + } + TapTree::Leaf(k) => k.count_assets(), + } + } } impl fmt::Display for TapTree { diff --git a/src/miniscript/astelem.rs b/src/miniscript/astelem.rs index 9a7b6e9d5..fbe6a653d 100644 --- a/src/miniscript/astelem.rs +++ b/src/miniscript/astelem.rs @@ -597,6 +597,93 @@ impl Terminal { } impl Terminal { + /// Count total possible assets + pub fn count_assets(&self) -> u64 { + match self { + Terminal::True => 0, + Terminal::False => 0, + Terminal::PkK(_) => 1, + Terminal::PkH(_) => 1, + Terminal::RawPkH(_) => 1, + // What happens to timelocks ? for both the assets and the count. + Terminal::After(_) => 0, + Terminal::Older(_) => 0, + Terminal::Sha256(_) => 1, + Terminal::Hash256(_) => 1, + Terminal::Ripemd160(_) => 1, + Terminal::Hash160(_) => 1, + Terminal::Alt(k) => k.count_assets(), + Terminal::Swap(k) => k.count_assets(), + Terminal::Check(k) => k.count_assets(), + Terminal::DupIf(k) => k.count_assets(), + Terminal::Verify(k) => k.count_assets(), + Terminal::NonZero(k) => k.count_assets(), + Terminal::ZeroNotEqual(k) => k.count_assets(), + Terminal::AndV(left, right) => { + let left_count = left.count_assets(); + let right_count = right.count_assets(); + left_count * right_count + } + Terminal::AndB(left, right) => { + let left_count = left.count_assets(); + let right_count = right.count_assets(); + left_count * right_count + } + Terminal::AndOr(a, b, c) => { + let a = a.count_assets(); + let b = b.count_assets(); + let c = c.count_assets(); + (a * b) + c + } + Terminal::OrB(left, right) => { + let left_count = left.count_assets(); + let right_count = right.count_assets(); + left_count + right_count + } + Terminal::OrD(left, right) => { + let left_count = left.count_assets(); + let right_count = right.count_assets(); + left_count + right_count + } + Terminal::OrC(left, right) => { + let left_count = left.count_assets(); + let right_count = right.count_assets(); + left_count + right_count + } + Terminal::OrI(left, right) => { + let left_count = left.count_assets(); + let right_count = right.count_assets(); + left_count + right_count + } + Terminal::Thresh(k, ms_v) => { + // k = 2, n = ms_v.len() + // ms_v = [ms(A),ms(B),ms(C)]; + // Assume count array as [5,7,8] and k=2 + // get_combinations_product gives [5*7,5*8,7*8] = [35,40,56] + let mut count_array = Vec::new(); + for ms in ms_v { + count_array.push(ms.count_assets()); + } + let products = Self::get_combinations_product(&count_array, *k as u64); + let mut total_count: u64 = 0; + for product in products { + total_count += product; + } + total_count + } + Terminal::Multi(k, dpk) => { + let k: u64 = *k as u64; + let n: u64 = dpk.len() as u64; + Self::k_of_n(k, n) + } + Terminal::MultiA(k, dpk) => { + let k: u64 = *k as u64; + let n: u64 = dpk.len() as u64; + Self::k_of_n(k, n) + } + } + } + /// Retrieve the assets associated with the type of miniscript element. pub fn get_all_assets(&self) -> Vec { match self { @@ -824,4 +911,38 @@ impl Terminal { current_combination.truncate(current_combination.len() - 1); } } + + // Do product of K combinations + fn get_combinations_product(values: &[u64], k: u64) -> Vec { + let mut products = Vec::new(); + let n = values.len(); + + if k == 0 { + return vec![1]; // Empty combination has a product of 1 + } + + // Using bitwise operations to generate combinations + let max_combinations = 1u32 << n; + for combination_bits in 1..max_combinations { + if combination_bits.count_ones() as usize == k as usize { + let mut product = 1; + for i in 0..n { + if combination_bits & (1u32 << i) != 0 { + product *= values[i]; + } + } + products.push(product); + } + } + + products + } + + // ways to select k things out of n + fn k_of_n(k: u64, n: u64) -> u64 { + if k == 0 || k == n { + return 1; + } + Self::k_of_n(k - 1, n - 1) + Self::k_of_n(k, n - 1) + } } diff --git a/src/miniscript/mod.rs b/src/miniscript/mod.rs index a5cb03ffd..f3547eab8 100644 --- a/src/miniscript/mod.rs +++ b/src/miniscript/mod.rs @@ -410,6 +410,11 @@ impl Miniscript { pub fn get_all_assets(&self) -> Vec { self.node.get_all_assets() } + + /// Get the total number of assets possible + pub fn count_assets(&self) -> u64 { + self.node.count_assets() + } } impl ForEachKey for Miniscript {