From ec17ec356d1d1e1ddf3ec6b903c3958ffd512131 Mon Sep 17 00:00:00 2001 From: Jeremy Rubin Date: Thu, 24 Mar 2022 14:54:58 -0700 Subject: [PATCH] Move with_huffman_tree logic to TaprootBuilder --- src/util/taproot.rs | 93 ++++++++++++++++++++++++++------------------- 1 file changed, 53 insertions(+), 40 deletions(-) diff --git a/src/util/taproot.rs b/src/util/taproot.rs index cb9950db..7da73a70 100644 --- a/src/util/taproot.rs +++ b/src/util/taproot.rs @@ -191,24 +191,9 @@ pub struct TaprootSpendInfo { } impl TaprootSpendInfo { - /// Create a new [`TaprootSpendInfo`] from a list of script(with default script version) and - /// weights of satisfaction for that script. The weights represent the probability of - /// each branch being taken. If probabilities/weights for each condition are known, - /// constructing the tree as a Huffman tree is the optimal way to minimize average - /// case satisfaction cost. This function takes input an iterator of tuple(u64, &Script) - /// where usize represents the satisfaction weights of the branch. - /// For example, [(3, S1), (2, S2), (5, S3)] would construct a TapTree that has optimal - /// satisfaction weight when probability for S1 is 30%, S2 is 20% and S3 is 50%. + /// Create a new [`TaprootSpendInfo`] from a list of script(with default script version). /// - /// # Errors: - /// - /// - When the optimal huffman tree has a depth more than 128 - /// - If the provided list of script weights is empty - /// - /// # Edge Cases: - /// - If the script weight calculations overflow, a sub-optimal tree may be generated. This - /// should not happen unless you are dealing with billions of branches with weights close to - /// 2^32. + /// See [`TaprootBuilder::with_huffman_tree`] for more detailed documentation pub fn with_huffman_tree( secp: &Secp256k1, internal_key: UntweakedPublicKey, @@ -218,29 +203,7 @@ impl TaprootSpendInfo { I: IntoIterator, C: secp256k1::Verification, { - let mut node_weights = BinaryHeap::<(Reverse, NodeInfo)>::new(); - for (p, leaf) in script_weights { - node_weights.push((Reverse(p as u64), NodeInfo::new_leaf_with_ver(leaf, LeafVersion::TapScript))); - } - if node_weights.is_empty() { - return Err(TaprootBuilderError::IncompleteTree); - } - while node_weights.len() > 1 { - // Combine the last two elements and insert a new node - let (p1, s1) = node_weights.pop().expect("len must be at least two"); - let (p2, s2) = node_weights.pop().expect("len must be at least two"); - // Insert the sum of first two in the tree as a new node - // N.B.: p1 + p2 can not practically saturate as you would need to have 2**32 max u32s - // from the input to overflow. However, saturating is a reasonable behavior here as - // huffman tree construction would treat all such elements as "very likely". - let p = Reverse(p1.0.saturating_add(p2.0)); - node_weights.push((p, NodeInfo::combine(s1, s2)?)); - } - // Every iteration of the loop reduces the node_weights.len() by exactly 1 - // Therefore, the loop will eventually terminate with exactly 1 element - debug_assert!(node_weights.len() == 1); - let node = node_weights.pop().expect("huffman tree algorithm is broken").1; - return Ok(Self::from_node_info(secp, internal_key, node)); + TaprootBuilder::with_huffman_tree(script_weights)?.finalize(secp, internal_key) } /// Create a new key spend with internal key and proided merkle root. @@ -396,6 +359,56 @@ impl TaprootBuilder { pub fn new() -> Self { TaprootBuilder { branch: vec![] } } + + /// Create a new [`TaprootBuilder`] from a list of script(with default script version) and + /// weights of satisfaction for that script. The weights represent the probability of + /// each branch being taken. If probabilities/weights for each condition are known, + /// constructing the tree as a Huffman tree is the optimal way to minimize average + /// case satisfaction cost. This function takes input an iterator of tuple(u64, &Script) + /// where usize represents the satisfaction weights of the branch. + /// For example, [(3, S1), (2, S2), (5, S3)] would construct a TapTree that has optimal + /// satisfaction weight when probability for S1 is 30%, S2 is 20% and S3 is 50%. + /// + /// # Errors: + /// + /// - When the optimal huffman tree has a depth more than 128 + /// - If the provided list of script weights is empty + /// + /// # Edge Cases: + /// - If the script weight calculations overflow, a sub-optimal tree may be generated. This + /// should not happen unless you are dealing with billions of branches with weights close to + /// 2^32. + pub fn with_huffman_tree( + script_weights: I, + ) -> Result + where + I: IntoIterator, + { + let mut node_weights = BinaryHeap::<(Reverse, NodeInfo)>::new(); + for (p, leaf) in script_weights { + node_weights.push((Reverse(p as u64), NodeInfo::new_leaf_with_ver(leaf, LeafVersion::TapScript))); + } + if node_weights.is_empty() { + return Err(TaprootBuilderError::IncompleteTree); + } + while node_weights.len() > 1 { + // Combine the last two elements and insert a new node + let (p1, s1) = node_weights.pop().expect("len must be at least two"); + let (p2, s2) = node_weights.pop().expect("len must be at least two"); + // Insert the sum of first two in the tree as a new node + // N.B.: p1 + p2 can not practically saturate as you would need to have 2**32 max u32s + // from the input to overflow. However, saturating is a reasonable behavior here as + // huffman tree construction would treat all such elements as "very likely". + let p = Reverse(p1.0.saturating_add(p2.0)); + node_weights.push((p, NodeInfo::combine(s1, s2)?)); + } + // Every iteration of the loop reduces the node_weights.len() by exactly 1 + // Therefore, the loop will eventually terminate with exactly 1 element + debug_assert!(node_weights.len() == 1); + let node = node_weights.pop().expect("huffman tree algorithm is broken").1; + Ok(TaprootBuilder{branch: vec![Some(node)]}) + } + /// Just like [`TaprootBuilder::add_leaf`] but allows to specify script version pub fn add_leaf_with_ver( self,