commonware_utils/
probability.rs

1//! A platform-independent probability value and sampler.
2
3use core::fmt;
4use rand::Rng;
5
6// Number of possible `u64` samples and denominator of the threshold grid.
7const SCALE: u128 = 1u128 << u64::BITS;
8
9// Biased `f64` exponent where scaling by 2^64 leaves the 53-bit significand unshifted.
10const SCALED_EXPONENT: u64 = 1023 + 52 - 64;
11
12/// Error returned when an `f64` cannot be represented as a probability.
13#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)]
14#[error(
15    "probability must be finite, within [0, 1], and exactly representable as a 64-bit threshold"
16)]
17pub struct InvalidProbability;
18
19/// A probability represented as a threshold over all possible `u64` samples.
20///
21/// Ratios are rounded down to the nearest multiple of 2^-64. Sampling consumes one `u64` for
22/// probabilities strictly between zero and one, and consumes no randomness for either endpoint.
23/// Given the same sequence of `u64` samples, decisions are identical on every platform.
24#[derive(Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd)]
25#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
26pub struct Probability(u64);
27
28impl Probability {
29    /// Creates a probability from `numerator / denominator`.
30    ///
31    /// Returns [`None`] if the denominator is zero or the numerator exceeds the denominator.
32    pub const fn new(numerator: u64, denominator: u64) -> Option<Self> {
33        if denominator == 0 || numerator > denominator {
34            return None;
35        }
36
37        if numerator == denominator {
38            return Some(Self(u64::MAX));
39        }
40
41        let threshold = ((numerator as u128) << u64::BITS) / denominator as u128;
42
43        // A proper fraction with a `u64` denominator is at least 2^-64 below one, so its rounded
44        // threshold cannot collide with the sentinel reserved for probability one.
45        assert!(threshold < u64::MAX as u128);
46        Some(Self(threshold as u64))
47    }
48
49    /// Creates a probability from an `f64` that maps exactly to a 64-bit threshold.
50    ///
51    /// Returns [`None`] if `value` is not finite, is outside `[0, 1]`, or would require rounding.
52    /// The exact IEEE-754 value is preserved rather than interpreting its source spelling as a
53    /// decimal ratio. Use [`TryFrom`] when const evaluation is not required.
54    pub const fn from_f64(value: f64) -> Option<Self> {
55        let bits = value.to_bits();
56        let magnitude = bits & (u64::MAX >> 1);
57
58        if magnitude == 0 {
59            return Some(Self(0));
60        }
61        if bits != magnitude || magnitude > 1.0f64.to_bits() {
62            return None;
63        }
64
65        let exponent = magnitude >> 52;
66        if exponent == 0 {
67            return None;
68        }
69        let significand = (1u64 << 52) | (magnitude & ((1u64 << 52) - 1));
70        if exponent < SCALED_EXPONENT {
71            let shift = (SCALED_EXPONENT - exponent) as u32;
72            if significand.trailing_zeros() < shift {
73                return None;
74            }
75            return Some(Self(significand >> shift));
76        }
77
78        let threshold = (significand as u128) << (exponent - SCALED_EXPONENT);
79        if threshold == SCALE {
80            return Some(Self(u64::MAX));
81        }
82
83        // Exact one is handled above, so the narrowed threshold cannot use its reserved sentinel.
84        assert!(threshold < u64::MAX as u128);
85        Some(Self(threshold as u64))
86    }
87
88    /// Returns whether this probability never occurs.
89    pub const fn is_zero(self) -> bool {
90        self.0 == 0
91    }
92
93    /// Returns whether this probability always occurs.
94    pub const fn is_one(self) -> bool {
95        self.0 == u64::MAX
96    }
97
98    /// Converts this probability to an `f64` in the inclusive range `[0, 1]`.
99    ///
100    /// This conversion is intended for APIs that require floating-point probabilities. Interior
101    /// probabilities remain strictly below one even when rounding to `f64`.
102    pub fn as_f64(self) -> f64 {
103        if self.is_one() {
104            return 1.0;
105        }
106
107        let value = self.0 as f64 / SCALE as f64;
108        if value == 1.0 {
109            f64::from_bits(1.0f64.to_bits() - 1)
110        } else {
111            value
112        }
113    }
114
115    /// Samples this probability using the next `u64` from `rng`.
116    pub fn sample<R: Rng + ?Sized>(self, rng: &mut R) -> bool {
117        match self.0 {
118            0 => false,
119            u64::MAX => true,
120            threshold => rng.next_u64() < threshold,
121        }
122    }
123}
124
125impl TryFrom<f64> for Probability {
126    type Error = InvalidProbability;
127
128    fn try_from(value: f64) -> Result<Self, Self::Error> {
129        Self::from_f64(value).ok_or(InvalidProbability)
130    }
131}
132
133impl fmt::Debug for Probability {
134    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
135        fmt::Debug::fmt(&self.as_f64(), f)
136    }
137}
138
139/// Creates a [`Probability`] from an integer ratio or an exactly representable `f64`.
140///
141/// The two-argument form preserves the exact ratio. The one-argument form preserves the exact
142/// IEEE-754 value and accepts only a literal that maps exactly to a 64-bit threshold. Ratio
143/// literals are validated at compile time; ratio expressions are validated at runtime.
144///
145/// # Panics
146///
147/// The ratio expression form panics if its denominator is zero or its numerator exceeds the
148/// denominator. Use [`Probability::new`] or [`Probability::try_from`] to validate untrusted values
149/// without panicking.
150///
151/// # Examples
152///
153/// ```
154/// use commonware_utils::{Probability, probability};
155///
156/// const HALF: Probability = probability!(1, 2);
157/// const NINETY_EIGHT_PERCENT: Probability = probability!(0.98);
158/// assert_eq!(HALF.as_f64(), 0.5);
159/// assert_eq!(NINETY_EIGHT_PERCENT.as_f64(), 0.98);
160/// ```
161///
162/// ```compile_fail
163/// use commonware_utils::{Probability, probability};
164///
165/// const INVALID: Probability = probability!(2, 1);
166/// ```
167///
168/// ```compile_fail
169/// use commonware_utils::{Probability, probability};
170///
171/// const REQUIRES_ROUNDING: Probability = probability!(1e-20);
172/// ```
173#[cfg(not(any(
174    commonware_stability_GAMMA,
175    commonware_stability_DELTA,
176    commonware_stability_EPSILON,
177    commonware_stability_RESERVED
178)))] // BETA
179#[macro_export]
180macro_rules! probability {
181    ($value:literal) => {
182        const {
183            $crate::Probability::from_f64($value).expect(
184                "probability requires a value in [0, 1] exactly representable as a 64-bit threshold",
185            )
186        }
187    };
188    ($numerator:literal, $denominator:literal) => {
189        const {
190            $crate::Probability::new($numerator, $denominator)
191                .expect("probability requires a non-zero denominator and numerator <= denominator")
192        }
193    };
194    ($numerator:expr, $denominator:expr) => {
195        $crate::Probability::new($numerator, $denominator)
196            .expect("probability requires a non-zero denominator and numerator <= denominator")
197    };
198}
199
200#[cfg(test)]
201mod tests {
202    use super::*;
203    use core::convert::Infallible;
204    use rand::TryRng;
205
206    struct CountingRng {
207        value: u64,
208        calls: usize,
209    }
210
211    impl TryRng for CountingRng {
212        type Error = Infallible;
213
214        fn try_next_u32(&mut self) -> Result<u32, Self::Error> {
215            self.calls += 1;
216            Ok(self.value as u32)
217        }
218
219        fn try_next_u64(&mut self) -> Result<u64, Self::Error> {
220            self.calls += 1;
221            Ok(self.value)
222        }
223
224        fn try_fill_bytes(&mut self, dst: &mut [u8]) -> Result<(), Self::Error> {
225            self.calls += 1;
226            dst.fill(0);
227            Ok(())
228        }
229    }
230
231    #[test]
232    fn construction() {
233        assert_eq!(Probability::new(0, u64::MAX), Some(probability!(0.0)));
234        assert_eq!(
235            Probability::new(u64::MAX, u64::MAX),
236            Some(probability!(1.0))
237        );
238        assert_eq!(probability!(1, 2), probability!(2, 4));
239        assert_eq!(probability!(1, 2).as_f64(), 0.5);
240        assert!(Probability::new(1, 0).is_none());
241        assert!(Probability::new(2, 1).is_none());
242    }
243
244    #[test]
245    fn f64_construction_preserves_clean_binary_value() {
246        const FROM_LITERAL: Probability = probability!(0.98);
247        const MINIMUM_INTERIOR: f64 = f64::from_bits(959u64 << 52);
248
249        assert_eq!(FROM_LITERAL.0, 18_077_809_192_235_360_256);
250        assert_eq!(FROM_LITERAL.as_f64(), 0.98);
251        assert_ne!(FROM_LITERAL, probability!(49, 50));
252        assert_eq!(probability!(0.5), probability!(1, 2));
253        assert_eq!(Probability::try_from(0.5), Ok(probability!(0.5)));
254        assert_eq!(Probability::from_f64(0.0), Some(probability!(0.0)));
255        assert_eq!(Probability::from_f64(-0.0), Some(probability!(0.0)));
256        assert_eq!(Probability::from_f64(1.0), Some(probability!(1.0)));
257        assert_eq!(
258            Probability::from_f64(MINIMUM_INTERIOR),
259            Some(Probability(1))
260        );
261
262        let below_one = f64::from_bits(1.0f64.to_bits() - 1);
263        assert_eq!(
264            Probability::from_f64(below_one).unwrap().as_f64(),
265            below_one
266        );
267    }
268
269    #[test]
270    fn f64_construction_rejects_lossy_or_invalid_values() {
271        const BELOW_MINIMUM: f64 = f64::from_bits(958u64 << 52);
272        const MINIMUM_INTERIOR_BITS: u64 = 959u64 << 52;
273
274        for value in [
275            BELOW_MINIMUM,
276            f64::from_bits(MINIMUM_INTERIOR_BITS + 1),
277            f64::from_bits(1),
278            -0.1,
279            1.1,
280            f64::from_bits(1.0f64.to_bits() + 1),
281            f64::NAN,
282            f64::from_bits(f64::NAN.to_bits() | (1u64 << 63)),
283            f64::INFINITY,
284            f64::NEG_INFINITY,
285        ] {
286            assert!(Probability::from_f64(value).is_none());
287            assert_eq!(Probability::try_from(value), Err(InvalidProbability));
288        }
289    }
290
291    #[test]
292    fn representation_matches_a_raw_rate() {
293        assert_eq!(core::mem::size_of::<Probability>(), size_of::<u64>());
294    }
295
296    #[test]
297    fn ratios_use_platform_independent_thresholds() {
298        assert_eq!(probability!(1, 3).0 as u128, SCALE / 3);
299        assert_eq!(probability!(2, 3).0 as u128, (2 * SCALE) / 3);
300
301        let below_one = Probability::new(u64::MAX - 1, u64::MAX).unwrap();
302        assert_eq!(below_one.0, u64::MAX - 1);
303        assert!(below_one.as_f64() < 1.0);
304    }
305
306    #[test]
307    fn sampling_uses_threshold_and_skips_endpoints() {
308        let mut rng = CountingRng { value: 0, calls: 0 };
309        assert!(!probability!(0.0).sample(&mut rng));
310        assert!(probability!(1.0).sample(&mut rng));
311        assert_eq!(rng.calls, 0);
312
313        rng.value = (1u64 << 63) - 1;
314        assert!(probability!(1, 2).sample(&mut rng));
315        assert_eq!(rng.calls, 1);
316
317        rng.value = 1u64 << 63;
318        assert!(!probability!(1, 2).sample(&mut rng));
319        assert_eq!(rng.calls, 2);
320    }
321
322    #[test]
323    #[should_panic(
324        expected = "probability requires a non-zero denominator and numerator <= denominator"
325    )]
326    fn expression_macro_rejects_invalid_probability() {
327        let numerator = 2;
328        let denominator = 1;
329        let _ = probability!(numerator, denominator);
330    }
331}
commonware_utils/probability.rs

commonware_utils/
probability.rs
