21
21
#[ cfg( feature="std" ) ] use std:: collections:: HashMap ;
22
22
#[ cfg( all( feature="alloc" , not( feature="std" ) ) ) ] use alloc:: collections:: BTreeMap ;
23
23
24
+ #[ cfg( feature = "alloc" ) ] use distributions:: WeightedError ;
25
+
24
26
use super :: Rng ;
25
27
#[ cfg( feature="alloc" ) ] use distributions:: uniform:: { SampleUniform , SampleBorrow } ;
26
28
@@ -109,7 +111,7 @@ pub trait SliceRandom {
109
111
/// ```
110
112
/// [`choose`]: trait.SliceRandom.html#method.choose
111
113
#[ 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 >
113
115
where R : Rng + ?Sized ,
114
116
F : Fn ( & Self :: Item ) -> B ,
115
117
B : SampleBorrow < X > ,
@@ -129,7 +131,7 @@ pub trait SliceRandom {
129
131
/// [`choose_mut`]: trait.SliceRandom.html#method.choose_mut
130
132
/// [`choose_weighted`]: trait.SliceRandom.html#method.choose_weighted
131
133
#[ 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 >
133
135
where R : Rng + ?Sized ,
134
136
F : Fn ( & Self :: Item ) -> B ,
135
137
B : SampleBorrow < X > ,
@@ -327,7 +329,7 @@ impl<T> SliceRandom for [T] {
327
329
}
328
330
329
331
#[ 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 >
331
333
where R : Rng + ?Sized ,
332
334
F : Fn ( & Self :: Item ) -> B ,
333
335
B : SampleBorrow < X > ,
@@ -337,12 +339,12 @@ impl<T> SliceRandom for [T] {
337
339
Clone +
338
340
Default {
339
341
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) ] )
342
344
}
343
345
344
346
#[ 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 >
346
348
where R : Rng + ?Sized ,
347
349
F : Fn ( & Self :: Item ) -> B ,
348
350
B : SampleBorrow < X > ,
@@ -352,9 +354,8 @@ impl<T> SliceRandom for [T] {
352
354
Clone +
353
355
Default {
354
356
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) ] )
358
359
}
359
360
360
361
fn shuffle < R > ( & mut self , rng : & mut R ) where R : Rng + ?Sized
@@ -868,8 +869,10 @@ mod test {
868
869
869
870
// Check error cases
870
871
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 ) ) ;
874
877
}
875
878
}
0 commit comments