count_assets method for total ways to obtain assets

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

Author: Harshil-Jani

src/descriptor/mod.rs

@@ -569,6 +569,51 @@ impl Descriptor<DefiniteDescriptorKey> {
 }
 
 impl Descriptor<DescriptorPublicKey> {
+    /// 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<Vec<Assets>, Error> {
         match self {
@@ -659,6 +704,14 @@ impl Descriptor<DescriptorPublicKey> {
         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<P, Q> TranslatePk<P, Q> for Descriptor<P>

src/descriptor/tr.rs

@@ -161,6 +161,18 @@ impl TapTree<DescriptorPublicKey> {
             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<Pk: MiniscriptKey> fmt::Display for TapTree<Pk> {

src/miniscript/astelem.rs

@@ -597,6 +597,93 @@ impl<Pk: MiniscriptKey, Ctx: ScriptContext> Terminal<Pk, Ctx> {
 }
 
 impl<Ctx: ScriptContext> Terminal<DescriptorPublicKey, Ctx> {
+    /// 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<Assets> {
         match self {
@@ -824,4 +911,38 @@ 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/miniscript/mod.rs

@@ -410,6 +410,11 @@ impl<Ctx: ScriptContext> Miniscript<DescriptorPublicKey, Ctx> {
     pub fn get_all_assets(&self) -> Vec<Assets> {
         self.node.get_all_assets()
     }
+
+    /// Get the total number of assets possible
+    pub fn count_assets(&self) -> u64 {
+        self.node.count_assets()
+    }
 }
 
 impl<Pk: MiniscriptKey, Ctx: ScriptContext> ForEachKey<Pk> for Miniscript<Pk, Ctx> {