File tree Expand file tree Collapse file tree 1 file changed +9
-1
lines changed Expand file tree Collapse file tree 1 file changed +9
-1
lines changed Original file line number Diff line number Diff line change @@ -117,7 +117,11 @@ impl<X: SampleUniform + PartialOrd> WeightedIndex<X> {
117117 return Err ( WeightedError :: InvalidWeight ) ;
118118 }
119119 weights. push ( total_weight. clone ( ) ) ;
120+
120121 total_weight += w. borrow ( ) ;
122+ if total_weight < * w. borrow ( ) {
123+ return Err ( WeightedError :: Overflow ) ;
124+ }
121125 }
122126
123127 if total_weight == zero {
@@ -423,7 +427,7 @@ mod test {
423427
424428 #[ test]
425429 fn overflow ( ) {
426- assert ! ( WeightedIndex :: new( [ 2 , usize :: MAX ] ) . is_err ( ) ) ;
430+ assert_eq ! ( WeightedIndex :: new( [ 2 , usize :: MAX ] ) , Err ( WeightedError :: Overflow ) ) ;
427431 }
428432}
429433
@@ -443,6 +447,9 @@ pub enum WeightedError {
443447
444448 /// Too many weights are provided (length greater than `u32::MAX`)
445449 TooMany ,
450+
451+ /// The sum of weights overflows
452+ Overflow ,
446453}
447454
448455#[ cfg( feature = "std" ) ]
@@ -455,6 +462,7 @@ impl fmt::Display for WeightedError {
455462 WeightedError :: InvalidWeight => "A weight is invalid in distribution" ,
456463 WeightedError :: AllWeightsZero => "All weights are zero in distribution" ,
457464 WeightedError :: TooMany => "Too many weights (hit u32::MAX) in distribution" ,
465+ WeightedError :: Overflow => "The sum of weights overflowed" ,
458466 } )
459467 }
460468}
You can’t perform that action at this time.
0 commit comments