2020#[ cfg( feature="std" ) ] use std:: collections:: HashMap ;
2121#[ cfg( all( feature="alloc" , not( feature="std" ) ) ) ] use alloc:: btree_map:: BTreeMap ;
2222
23+ #[ cfg( feature = "alloc" ) ] use distributions:: WeightedError ;
2324
2425use super :: Rng ;
2526#[ cfg( feature="alloc" ) ] use distributions:: uniform:: { SampleUniform , SampleBorrow } ;
@@ -109,7 +110,7 @@ pub trait SliceRandom {
109110 /// ```
110111 /// [`choose`]: trait.SliceRandom.html#method.choose
111112 #[ cfg( feature = "alloc" ) ]
112- fn choose_weighted < R , F , B , X > ( & self , rng : & mut R , weight : F ) -> Option < & Self :: Item >
113+ fn choose_weighted < R , F , B , X > ( & self , rng : & mut R , weight : F ) -> Result < & Self :: Item , WeightedError >
113114 where R : Rng + ?Sized ,
114115 F : Fn ( & Self :: Item ) -> B ,
115116 B : SampleBorrow < X > ,
@@ -129,7 +130,7 @@ pub trait SliceRandom {
129130 /// [`choose_mut`]: trait.SliceRandom.html#method.choose_mut
130131 /// [`choose_weighted`]: trait.SliceRandom.html#method.choose_weighted
131132 #[ cfg( feature = "alloc" ) ]
132- fn choose_weighted_mut < R , F , B , X > ( & mut self , rng : & mut R , weight : F ) -> Option < & mut Self :: Item >
133+ fn choose_weighted_mut < R , F , B , X > ( & mut self , rng : & mut R , weight : F ) -> Result < & mut Self :: Item , WeightedError >
133134 where R : Rng + ?Sized ,
134135 F : Fn ( & Self :: Item ) -> B ,
135136 B : SampleBorrow < X > ,
@@ -327,7 +328,7 @@ impl<T> SliceRandom for [T] {
327328 }
328329
329330 #[ cfg( feature = "alloc" ) ]
330- fn choose_weighted < R , F , B , X > ( & self , rng : & mut R , weight : F ) -> Option < & Self :: Item >
331+ fn choose_weighted < R , F , B , X > ( & self , rng : & mut R , weight : F ) -> Result < & Self :: Item , WeightedError >
331332 where R : Rng + ?Sized ,
332333 F : Fn ( & Self :: Item ) -> B ,
333334 B : SampleBorrow < X > ,
@@ -337,12 +338,12 @@ impl<T> SliceRandom for [T] {
337338 Clone +
338339 Default {
339340 use distributions:: { Distribution , WeightedIndex } ;
340- WeightedIndex :: new ( self . iter ( ) . map ( weight) ) . ok ( )
341- . map ( |distr| & self [ distr. sample ( rng) ] )
341+ let distr = WeightedIndex :: new ( self . iter ( ) . map ( weight) ) ? ;
342+ Ok ( & self [ distr. sample ( rng) ] )
342343 }
343344
344345 #[ cfg( feature = "alloc" ) ]
345- fn choose_weighted_mut < R , F , B , X > ( & mut self , rng : & mut R , weight : F ) -> Option < & mut Self :: Item >
346+ fn choose_weighted_mut < R , F , B , X > ( & mut self , rng : & mut R , weight : F ) -> Result < & mut Self :: Item , WeightedError >
346347 where R : Rng + ?Sized ,
347348 F : Fn ( & Self :: Item ) -> B ,
348349 B : SampleBorrow < X > ,
@@ -352,9 +353,8 @@ impl<T> SliceRandom for [T] {
352353 Clone +
353354 Default {
354355 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] )
356+ let distr = WeightedIndex :: new ( self . iter ( ) . map ( weight) ) ?;
357+ Ok ( & mut self [ distr. sample ( rng) ] )
358358 }
359359
360360 fn shuffle < R > ( & mut self , rng : & mut R ) where R : Rng + ?Sized
@@ -868,8 +868,10 @@ mod test {
868868
869869 // Check error cases
870870 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 ) ;
871+ assert_eq ! ( empty_slice. choose_weighted( & mut r, |_| 1 ) , Err ( WeightedError :: NoItem ) ) ;
872+ assert_eq ! ( empty_slice. choose_weighted_mut( & mut r, |_| 1 ) , Err ( WeightedError :: NoItem ) ) ;
873+ assert_eq ! ( [ 'x' ] . choose_weighted_mut( & mut r, |_| 0 ) , Err ( WeightedError :: AllWeightsZero ) ) ;
874+ assert_eq ! ( [ 0 , -1 ] . choose_weighted_mut( & mut r, |x| * x) , Err ( WeightedError :: NegativeWeight ) ) ;
875+ assert_eq ! ( [ -1 , 0 ] . choose_weighted_mut( & mut r, |x| * x) , Err ( WeightedError :: NegativeWeight ) ) ;
874876 }
875877}
0 commit comments