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}