2121#[ cfg( feature="std" ) ] use std:: collections:: HashMap ;
2222#[ cfg( all( feature="alloc" , not( feature="std" ) ) ) ] use alloc:: collections:: BTreeMap ;
2323
24+ #[ cfg( feature = "alloc" ) ] use distributions:: WeightedError ;
25+
2426use super :: Rng ;
2527#[ cfg( feature="alloc" ) ] use distributions:: uniform:: { SampleUniform , SampleBorrow } ;
2628
@@ -109,7 +111,7 @@ pub trait SliceRandom {
109111 /// ```
110112 /// [`choose`]: trait.SliceRandom.html#method.choose
111113 #[ cfg( feature = "alloc" ) ]
112- fn choose_weighted < R , F , B , X > ( & self , rng : & mut R , weight : F ) -> Option < & Self :: Item >
114+ fn choose_weighted < R , F , B , X > ( & self , rng : & mut R , weight : F ) -> Result < & Self :: Item , WeightedError >
113115 where R : Rng + ?Sized ,
114116 F : Fn ( & Self :: Item ) -> B ,
115117 B : SampleBorrow < X > ,
@@ -129,7 +131,7 @@ pub trait SliceRandom {
129131 /// [`choose_mut`]: trait.SliceRandom.html#method.choose_mut
130132 /// [`choose_weighted`]: trait.SliceRandom.html#method.choose_weighted
131133 #[ cfg( feature = "alloc" ) ]
132- fn choose_weighted_mut < R , F , B , X > ( & mut self , rng : & mut R , weight : F ) -> Option < & mut Self :: Item >
134+ fn choose_weighted_mut < R , F , B , X > ( & mut self , rng : & mut R , weight : F ) -> Result < & mut Self :: Item , WeightedError >
133135 where R : Rng + ?Sized ,
134136 F : Fn ( & Self :: Item ) -> B ,
135137 B : SampleBorrow < X > ,
@@ -327,7 +329,7 @@ impl<T> SliceRandom for [T] {
327329 }
328330
329331 #[ cfg( feature = "alloc" ) ]
330- fn choose_weighted < R , F , B , X > ( & self , rng : & mut R , weight : F ) -> Option < & Self :: Item >
332+ fn choose_weighted < R , F , B , X > ( & self , rng : & mut R , weight : F ) -> Result < & Self :: Item , WeightedError >
331333 where R : Rng + ?Sized ,
332334 F : Fn ( & Self :: Item ) -> B ,
333335 B : SampleBorrow < X > ,
@@ -337,12 +339,12 @@ impl<T> SliceRandom for [T] {
337339 Clone +
338340 Default {
339341 use distributions:: { Distribution , WeightedIndex } ;
340- WeightedIndex :: new ( self . iter ( ) . map ( weight) ) . ok ( )
341- . map ( |distr| & self [ distr. sample ( rng) ] )
342+ let distr = WeightedIndex :: new ( self . iter ( ) . map ( weight) ) ? ;
343+ Ok ( & self [ distr. sample ( rng) ] )
342344 }
343345
344346 #[ cfg( feature = "alloc" ) ]
345- fn choose_weighted_mut < R , F , B , X > ( & mut self , rng : & mut R , weight : F ) -> Option < & mut Self :: Item >
347+ fn choose_weighted_mut < R , F , B , X > ( & mut self , rng : & mut R , weight : F ) -> Result < & mut Self :: Item , WeightedError >
346348 where R : Rng + ?Sized ,
347349 F : Fn ( & Self :: Item ) -> B ,
348350 B : SampleBorrow < X > ,
@@ -352,9 +354,8 @@ impl<T> SliceRandom for [T] {
352354 Clone +
353355 Default {
354356 use distributions:: { Distribution , WeightedIndex } ;
355- WeightedIndex :: new ( self . iter ( ) . map ( weight) ) . ok ( )
356- . map ( |distr| distr. sample ( rng) )
357- . map ( move |ix| & mut self [ ix] )
357+ let distr = WeightedIndex :: new ( self . iter ( ) . map ( weight) ) ?;
358+ Ok ( & mut self [ distr. sample ( rng) ] )
358359 }
359360
360361 fn shuffle < R > ( & mut self , rng : & mut R ) where R : Rng + ?Sized
@@ -868,8 +869,10 @@ mod test {
868869
869870 // Check error cases
870871 let empty_slice = & mut [ 10 ] [ 0 ..0 ] ;
871- assert_eq ! ( empty_slice. choose_weighted( & mut r, |_| 1 ) , None ) ;
872- assert_eq ! ( empty_slice. choose_weighted_mut( & mut r, |_| 1 ) , None ) ;
873- assert_eq ! ( [ 'x' ] . choose_weighted_mut( & mut r, |_| 0 ) , None ) ;
872+ assert_eq ! ( empty_slice. choose_weighted( & mut r, |_| 1 ) , Err ( WeightedError :: NoItem ) ) ;
873+ assert_eq ! ( empty_slice. choose_weighted_mut( & mut r, |_| 1 ) , Err ( WeightedError :: NoItem ) ) ;
874+ assert_eq ! ( [ 'x' ] . choose_weighted_mut( & mut r, |_| 0 ) , Err ( WeightedError :: AllWeightsZero ) ) ;
875+ assert_eq ! ( [ 0 , -1 ] . choose_weighted_mut( & mut r, |x| * x) , Err ( WeightedError :: NegativeWeight ) ) ;
876+ assert_eq ! ( [ -1 , 0 ] . choose_weighted_mut( & mut r, |x| * x) , Err ( WeightedError :: NegativeWeight ) ) ;
874877 }
875878}
0 commit comments