count_assets method for total ways to obtain assets

Signed-off-by: Harshil Jani <[email protected]>

Author: Harshil-Jani

src/descriptor/mod.rs

@@ -25,7 +25,7 @@ use crate::miniscript::decode::Terminal;
 use crate::miniscript::{satisfy, Legacy, Miniscript, Segwitv0};
 use crate::plan::{AssetProvider, Assets, Plan};
 use crate::prelude::*;
-use crate::util::get_asset_combination;
+use crate::util::{asset_combination, k_of_n};
 use crate::{
     expression, hash256, BareCtx, Error, ForEachKey, MiniscriptKey, Satisfier, ToPublicKey,
     TranslateErr, TranslatePk, Translator,
@@ -581,23 +581,23 @@ impl Descriptor<DescriptorPublicKey> {
                     WshInner::SortedMulti(k) => {
                         let n = k.pks.len() as u64;
                         let k = k.k as u64;
-                        Self::k_of_n(k, n)
+                        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)
+                    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)
+                    k_of_n(k, n)
                 }
                 WshInner::Ms(k) => k.count_assets(),
             },
@@ -634,7 +634,7 @@ impl Descriptor<DescriptorPublicKey> {
                     WshInner::SortedMulti(k) => {
                         let dpk_v = k.clone().pks;
                         let k = k.clone().k;
-                        Ok(get_asset_combination(k, &dpk_v))
+                        Ok(asset_combination(k, &dpk_v))
                     }
                     WshInner::Ms(k) => Ok(k.all_assets()),
                 },
@@ -646,15 +646,15 @@ impl Descriptor<DescriptorPublicKey> {
                 ShInner::SortedMulti(k) => {
                     let dpk_v = k.clone().pks;
                     let k = k.clone().k;
-                    Ok(get_asset_combination(k, &dpk_v))
+                    Ok(asset_combination(k, &dpk_v))
                 }
                 ShInner::Ms(k) => Ok(k.all_assets()),
             },
             Descriptor::Wsh(k) => match k.as_inner() {
                 WshInner::SortedMulti(k) => {
                     let dpk_v = k.clone().pks;
                     let k = k.clone().k;
-                    Ok(get_asset_combination(k, &dpk_v))
+                    Ok(asset_combination(k, &dpk_v))
                 }
                 WshInner::Ms(k) => {
                     println!("{}", k);
@@ -677,14 +677,6 @@ impl Descriptor<DescriptorPublicKey> {
             }
         }
     }
-
-    // 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<P, Q> TranslatePk<P, Q> for Descriptor<P>

src/miniscript/astelem.rs

@@ -19,7 +19,7 @@ use crate::miniscript::types::{self, Property};
 use crate::miniscript::ScriptContext;
 use crate::plan::Assets;
 use crate::prelude::*;
-use crate::util::{MsKeyBuilder, get_asset_combination};
+use crate::util::{asset_combination, get_combinations_product, k_of_n, MsKeyBuilder};
 use crate::{
     errstr, expression, AbsLockTime, DescriptorPublicKey, Error, Miniscript, MiniscriptKey,
     Terminal, ToPublicKey,
@@ -605,7 +605,6 @@ impl<Ctx: ScriptContext> Terminal<DescriptorPublicKey, Ctx> {
             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,
@@ -619,12 +618,7 @@ impl<Ctx: ScriptContext> Terminal<DescriptorPublicKey, Ctx> {
             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) => {
+            Terminal::AndV(left, right) | Terminal::AndB(left, right) => {
                 let left_count = left.count_assets();
                 let right_count = right.count_assets();
                 left_count * right_count
@@ -635,51 +629,30 @@ impl<Ctx: ScriptContext> Terminal<DescriptorPublicKey, Ctx> {
                 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) => {
+            Terminal::OrB(left, right)
+            | Terminal::OrC(left, right)
+            | Terminal::OrD(left, right)
+            | 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 products = 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) => {
+            Terminal::Multi(k, dpk) | Terminal::MultiA(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)
+                k_of_n(k, n)
             }
         }
     }
@@ -741,9 +714,7 @@ impl<Ctx: ScriptContext> Terminal<DescriptorPublicKey, Ctx> {
             Terminal::Verify(k) => k.all_assets(),
             Terminal::NonZero(k) => k.all_assets(),
             Terminal::ZeroNotEqual(k) => k.all_assets(),
-            Terminal::AndB(left, right) |
-            Terminal::AndV(left, right)
-            => {
+            Terminal::AndB(left, right) | Terminal::AndV(left, right) => {
                 let a = left.all_assets();
                 let b = right.all_assets();
                 let result: Vec<Assets> = a
@@ -777,11 +748,10 @@ impl<Ctx: ScriptContext> Terminal<DescriptorPublicKey, Ctx> {
                 c.extend(and);
                 c
             }
-            Terminal::OrB(left, right) |
-            Terminal::OrC(left, right) |
-            Terminal::OrD(left, right) |
-            Terminal::OrI(left, right)
-            => {
+            Terminal::OrB(left, right)
+            | Terminal::OrC(left, right)
+            | Terminal::OrD(left, right)
+            | Terminal::OrI(left, right) => {
                 let mut a = left.all_assets();
                 let b = right.all_assets();
                 a.extend(b);
@@ -792,14 +762,14 @@ impl<Ctx: ScriptContext> Terminal<DescriptorPublicKey, Ctx> {
                 // Eg : thresh(2,ms(A),ms(B),ms(C)) Here ms(A),ms(B) and ms(C) are miniscript policies
                 // k = 2
                 // ms = [ms(A),ms(B),ms(C)];
-                // We would consider the possible combinations of k policies into the ms_v 
+                // We would consider the possible combinations of k policies into the ms_v
                 // here k=2 so all possible combinations of 2.
                 // ms_v = [[ms(A),ms(B)],[ms(A),ms(C)],[ms(B),ms(C)]]
-                // Between each set of combination we would need to do an OR 
+                // Between each set of combination we would need to do an OR
                 // (i.e ms_v[0] OR ms_v[1] OR ms_v[3])
                 // Now inside of each policy combination we need to have AND
                 // Eg : ms_v[0] = [ms(A),ms(B)] so here -> ms(A) AND ms(B)
-                
+
                 let ms_v = Self::get_ms_combination_thresh(*k, ms);
                 let mut result = Vec::new();
                 for ms in ms_v {
@@ -826,8 +796,7 @@ impl<Ctx: ScriptContext> Terminal<DescriptorPublicKey, Ctx> {
                 }
                 result
             }
-            Terminal::Multi(k, dpk_v) |
-            Terminal::MultiA(k,dpk_v ) => get_asset_combination(*k, dpk_v),
+            Terminal::Multi(k, dpk_v) | Terminal::MultiA(k, dpk_v) => asset_combination(*k, dpk_v),
         }
     }
 
@@ -860,38 +829,4 @@ impl<Ctx: ScriptContext> Terminal<DescriptorPublicKey, Ctx> {
             current_combination.truncate(current_combination.len() - 1);
         }
     }
-
-    // Do product of K combinations
-    fn get_combinations_product(values: &[u64], k: u64) -> Vec<u64> {
-        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)
-    }
 }

src/util.rs

@@ -9,8 +9,8 @@ use bitcoin::PubkeyHash;
 use crate::miniscript::context;
 use crate::miniscript::satisfy::Placeholder;
 use crate::plan::Assets;
-use crate::{prelude::*, DescriptorPublicKey};
-use crate::{MiniscriptKey, ScriptContext, ToPublicKey};
+use crate::prelude::*;
+use crate::{DescriptorPublicKey, MiniscriptKey, ScriptContext, ToPublicKey};
 pub(crate) fn varint_len(n: usize) -> usize {
     bitcoin::VarInt(n as u64).len()
 }
@@ -108,7 +108,7 @@ impl MsKeyBuilder for script::Builder {
 }
 
 // Helper to get all possible pairs of K of N assets
-pub fn get_asset_combination(k: usize, dpk_v: &Vec<DescriptorPublicKey>) -> Vec<Assets> {
+pub fn asset_combination(k: usize, dpk_v: &Vec<DescriptorPublicKey>) -> Vec<Assets> {
     let mut all_assets: Vec<Assets> = Vec::new();
     let current_assets = Assets::new();
     combine_assets(k, dpk_v, 0, current_assets, &mut all_assets);
@@ -136,3 +136,37 @@ pub fn combine_assets(
     println!("{:#?}", new_asset);
     combine_assets(k - 1, dpk_v, index + 1, new_asset, all_assets)
 }
+
+// Do product of K combinations
+pub fn get_combinations_product(values: &[u64], k: u64) -> Vec<u64> {
+    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
+pub fn k_of_n(k: u64, n: u64) -> u64 {
+    if k == 0 || k == n {
+        return 1;
+    }
+    k_of_n(k - 1, n - 1) + k_of_n(k, n - 1)
+}