1// Copyright 2018 Developers of the Rand project.
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8
9//! Weighted (index) sampling
10//!
11//! Primarily, this module houses the [`WeightedIndex`] distribution.
12//! See also [`rand_distr::weighted`] for alternative implementations supporting
13//! potentially-faster sampling or a more easily modifiable tree structure.
14//!
15//! [`rand_distr::weighted`]: https://docs.rs/rand_distr/latest/rand_distr/weighted/index.html
16
17use core::fmt;
18mod weighted_index;
19
20pub use weighted_index::WeightedIndex;
21
22/// Bounds on a weight
23///
24/// See usage in [`WeightedIndex`].
25pub trait Weight: Clone {
26 /// Representation of 0
27 const ZERO: Self;
28
29 /// Checked addition
30 ///
31 /// - `Result::Ok`: On success, `v` is added to `self`
32 /// - `Result::Err`: Returns an error when `Self` cannot represent the
33 /// result of `self + v` (i.e. overflow). The value of `self` should be
34 /// discarded.
35 #[allow(clippy::result_unit_err)]
36 fn checked_add_assign(&mut self, v: &Self) -> Result<(), ()>;
37}
38
39macro_rules! impl_weight_int {
40 ($t:ty) => {
41 impl Weight for $t {
42 const ZERO: Self = 0;
43 fn checked_add_assign(&mut self, v: &Self) -> Result<(), ()> {
44 match self.checked_add(*v) {
45 Some(sum) => {
46 *self = sum;
47 Ok(())
48 }
49 None => Err(()),
50 }
51 }
52 }
53 };
54 ($t:ty, $($tt:ty),*) => {
55 impl_weight_int!($t);
56 impl_weight_int!($($tt),*);
57 }
58}
59impl_weight_int!(i8, i16, i32, i64, i128, isize);
60impl_weight_int!(u8, u16, u32, u64, u128, usize);
61
62macro_rules! impl_weight_float {
63 ($t:ty) => {
64 impl Weight for $t {
65 const ZERO: Self = 0.0;
66
67 fn checked_add_assign(&mut self, v: &Self) -> Result<(), ()> {
68 // Floats have an explicit representation for overflow
69 *self += *v;
70 Ok(())
71 }
72 }
73 };
74}
75impl_weight_float!(f32);
76impl_weight_float!(f64);
77
78/// Invalid weight errors
79///
80/// This type represents errors from [`WeightedIndex::new`],
81/// [`WeightedIndex::update_weights`] and other weighted distributions.
82#[derive(Debug, Clone, Copy, PartialEq, Eq)]
83// Marked non_exhaustive to allow a new error code in the solution to #1476.
84#[non_exhaustive]
85pub enum Error {
86 /// The input weight sequence is empty, too long, or wrongly ordered
87 InvalidInput,
88
89 /// A weight is negative, too large for the distribution, or not a valid number
90 InvalidWeight,
91
92 /// Not enough non-zero weights are available to sample values
93 ///
94 /// When attempting to sample a single value this implies that all weights
95 /// are zero. When attempting to sample `amount` values this implies that
96 /// less than `amount` weights are greater than zero.
97 InsufficientNonZero,
98
99 /// Overflow when calculating the sum of weights
100 Overflow,
101}
102
103#[cfg(feature = "std")]
104impl std::error::Error for Error {}
105
106impl fmt::Display for Error {
107 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
108 f.write_str(match *self {
109 Error::InvalidInput => "Weights sequence is empty/too long/unordered",
110 Error::InvalidWeight => "A weight is negative, too large or not a valid number",
111 Error::InsufficientNonZero => "Not enough weights > zero",
112 Error::Overflow => "Overflow when summing weights",
113 })
114 }
115}
116