Fix Weighting Addition to never error on overflow + prevent overflows from ever happening with wider integers

This commit is contained in:
Jeremy Rubin 2021-11-14 17:55:30 -08:00
parent abc242dfe1
commit cccd75d004
1 changed files with 6 additions and 8 deletions

View File

@ -208,9 +208,9 @@ impl TaprootSpendInfo {
I: IntoIterator<Item = (u64, Script)>, I: IntoIterator<Item = (u64, Script)>,
C: secp256k1::Verification, C: secp256k1::Verification,
{ {
let mut node_weights = BinaryHeap::<(u64, NodeInfo)>::new(); let mut node_weights = BinaryHeap::<(u128, NodeInfo)>::new();
for (p, leaf) in script_weights { for (p, leaf) in script_weights {
node_weights.push((p, NodeInfo::new_leaf_with_ver(leaf, LeafVersion::default()))); node_weights.push((p as u128, NodeInfo::new_leaf_with_ver(leaf, LeafVersion::default())));
} }
if node_weights.is_empty() { if node_weights.is_empty() {
return Err(TaprootBuilderError::IncompleteTree); return Err(TaprootBuilderError::IncompleteTree);
@ -220,7 +220,10 @@ impl TaprootSpendInfo {
let (p1, s1) = node_weights.pop().expect("len must be at least two"); 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"); 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 // Insert the sum of first two in the tree as a new node
let p = p1.checked_add(p2).ok_or(TaprootBuilderError::ScriptWeightOverflow)?; // N.B.: p1 + p2 can never actually saturate as you would need to have 2**64 max u64s
// 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 = p1.saturating_add(p2);
node_weights.push((p, NodeInfo::combine(s1, s2)?)); node_weights.push((p, NodeInfo::combine(s1, s2)?));
} }
// Every iteration of the loop reduces the node_weights.len() by exactly 1 // Every iteration of the loop reduces the node_weights.len() by exactly 1
@ -801,8 +804,6 @@ pub enum TaprootBuilderError {
IncompleteTree, IncompleteTree,
/// Called finalize on a empty tree /// Called finalize on a empty tree
EmptyTree, EmptyTree,
/// Script weight overflow
ScriptWeightOverflow,
} }
impl fmt::Display for TaprootBuilderError { impl fmt::Display for TaprootBuilderError {
@ -830,9 +831,6 @@ impl fmt::Display for TaprootBuilderError {
TaprootBuilderError::EmptyTree => { TaprootBuilderError::EmptyTree => {
write!(f, "Called finalize on an empty tree") write!(f, "Called finalize on an empty tree")
} }
TaprootBuilderError::ScriptWeightOverflow => {
write!(f, "Script weight overflow in Huffman tree construction")
},
} }
} }
} }