Skip to content

Commit

Permalink
WeightedIndex::new: trap overflow in release builds only
Browse files Browse the repository at this point in the history
  • Loading branch information
dhardy committed Dec 4, 2023
1 parent 910883c commit e5e5b1f
Showing 1 changed file with 9 additions and 1 deletion.
10 changes: 9 additions & 1 deletion src/distributions/weighted_index.rs
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,11 @@ impl<X: SampleUniform + PartialOrd> WeightedIndex<X> {
return Err(WeightedError::InvalidWeight);
}
weights.push(total_weight.clone());

total_weight += w.borrow();
if total_weight < *w.borrow() {
return Err(WeightedError::Overflow);
}
}

if total_weight == zero {
Expand Down Expand Up @@ -423,7 +427,7 @@ mod test {

#[test]
fn overflow() {
assert!(WeightedIndex::new([2, usize::MAX]).is_err());
assert_eq!(WeightedIndex::new([2, usize::MAX]), Err(WeightedError::Overflow));
}
}

Expand All @@ -443,6 +447,9 @@ pub enum WeightedError {

/// Too many weights are provided (length greater than `u32::MAX`)
TooMany,

/// The sum of weights overflows
Overflow,
}

#[cfg(feature = "std")]
Expand All @@ -455,6 +462,7 @@ impl fmt::Display for WeightedError {
WeightedError::InvalidWeight => "A weight is invalid in distribution",
WeightedError::AllWeightsZero => "All weights are zero in distribution",
WeightedError::TooMany => "Too many weights (hit u32::MAX) in distribution",
WeightedError::Overflow => "The sum of weights overflowed",
})
}
}

0 comments on commit e5e5b1f

Please sign in to comment.