diff --git a/CHANGELOG.md b/CHANGELOG.md index aaccde2ce7..b431a9ef73 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,9 @@ A [separate changelog is kept for rand_core](rand_core/CHANGELOG.md). You may also find the [Upgrade Guide](https://rust-random.github.io/book/update.html) useful. +## [0.9.1] - unreleased +- Add the `Slice::num_choices` method to the Slice distribution (#1402) + ## [0.9.0-alpha.0] - 2024-02-18 This is a pre-release. To depend on this version, use `rand = "=0.9.0-alpha.0"` to prevent automatic updates (which can be expected to include breaking changes). diff --git a/src/distributions/slice.rs b/src/distributions/slice.rs index 5fc08751f6..1c96680aa6 100644 --- a/src/distributions/slice.rs +++ b/src/distributions/slice.rs @@ -6,6 +6,8 @@ // option. This file may not be copied, modified, or distributed // except according to those terms. +use core::num::NonZeroUsize; + use crate::distributions::{Distribution, Uniform}; #[cfg(feature = "alloc")] use alloc::string::String; @@ -67,19 +69,25 @@ use alloc::string::String; pub struct Slice<'a, T> { slice: &'a [T], range: Uniform, + num_choices: NonZeroUsize, } impl<'a, T> Slice<'a, T> { /// Create a new `Slice` instance which samples uniformly from the slice. /// Returns `Err` if the slice is empty. pub fn new(slice: &'a [T]) -> Result { - match slice.len() { - 0 => Err(EmptySlice), - len => Ok(Self { - slice, - range: Uniform::new(0, len).unwrap(), - }), - } + let num_choices = NonZeroUsize::new(slice.len()).ok_or(EmptySlice)?; + + Ok(Self { + slice, + range: Uniform::new(0, num_choices.get()).unwrap(), + num_choices, + }) + } + + /// Returns the count of choices in this distribution + pub fn num_choices(&self) -> NonZeroUsize { + self.num_choices } }