Skip to main content

rand/distr/weighted/
mod.rs

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    /// Note that for floating-point formats with a representation of infinity,
32    /// overflow-to-infinity is not considered an error.
33    ///
34    /// -   `Result::Ok`: On success, `v` is added to `self`
35    /// -   `Result::Err`: Returns an error when `Self` cannot represent the
36    ///     result of `self + v` (i.e. overflow). The value of `self` should be
37    ///     discarded.
38    #[allow(clippy::result_unit_err)]
39    fn checked_add_assign(&mut self, v: &Self) -> Result<(), ()>;
40}
41
42macro_rules! impl_weight_int {
43    ($t:ty) => {
44        impl Weight for $t {
45            const ZERO: Self = 0;
46            fn checked_add_assign(&mut self, v: &Self) -> Result<(), ()> {
47                match self.checked_add(*v) {
48                    Some(sum) => {
49                        *self = sum;
50                        Ok(())
51                    }
52                    None => Err(()),
53                }
54            }
55        }
56    };
57    ($t:ty, $($tt:ty),*) => {
58        impl_weight_int!($t);
59        impl_weight_int!($($tt),*);
60    }
61}
62impl_weight_int!(i8, i16, i32, i64, i128, isize);
63impl_weight_int!(u8, u16, u32, u64, u128, usize);
64
65macro_rules! impl_weight_float {
66    ($t:ty) => {
67        impl Weight for $t {
68            const ZERO: Self = 0.0;
69
70            fn checked_add_assign(&mut self, v: &Self) -> Result<(), ()> {
71                // Floats have an explicit representation for overflow
72                *self += *v;
73                Ok(())
74            }
75        }
76    };
77}
78impl_weight_float!(f32);
79impl_weight_float!(f64);
80
81/// Invalid weight errors
82///
83/// This type represents errors from [`WeightedIndex::new`],
84/// [`WeightedIndex::update_weights`] and other weighted distributions.
85#[derive(Debug, Clone, Copy, PartialEq, Eq)]
86// Marked non_exhaustive to allow a new error code in the solution to #1476.
87#[non_exhaustive]
88pub enum Error {
89    /// The input weight sequence is empty, too long, or wrongly ordered
90    InvalidInput,
91
92    /// A weight is negative, too large for the distribution, or not a valid number
93    InvalidWeight,
94
95    /// Not enough non-zero weights are available to sample values
96    ///
97    /// When attempting to sample a single value this implies that all weights
98    /// are zero. When attempting to sample `amount` values this implies that
99    /// less than `amount` weights are greater than zero.
100    InsufficientNonZero,
101
102    /// Overflow when calculating the sum of weights
103    Overflow,
104}
105
106impl core::error::Error for Error {}
107
108impl fmt::Display for Error {
109    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
110        f.write_str(match *self {
111            Error::InvalidInput => "Weights sequence is empty/too long/unordered",
112            Error::InvalidWeight => "A weight is negative, too large or not a valid number",
113            Error::InsufficientNonZero => "Not enough weights > zero",
114            Error::Overflow => "Overflow when summing weights",
115        })
116    }
117}