Skip to main content

amari_flynn/
prob.rs

1//! Probabilistic value wrapper
2//!
3//! The core `Prob<T>` type wraps values with runtime probability tracking.
4
5use rand::Rng;
6
7/// A value with associated probability
8///
9/// Wraps any value `T` with runtime metadata tracking its probability.
10/// This is the foundation for probabilistic contracts.
11///
12/// # Examples
13///
14/// ```
15/// use amari_flynn::prob::Prob;
16///
17/// let coin = Prob::with_probability(0.5, true);
18/// assert_eq!(coin.probability(), 0.5);
19/// assert_eq!(coin.into_inner(), true);
20/// ```
21#[derive(Clone, Copy, Debug)]
22pub struct Prob<T> {
23    value: T,
24    probability: f64,
25}
26
27impl<T> Prob<T> {
28    /// Create a probabilistic value with associated probability
29    ///
30    /// # Examples
31    ///
32    /// ```
33    /// use amari_flynn::prob::Prob;
34    ///
35    /// let certain = Prob::new(42);
36    /// assert_eq!(certain.probability(), 1.0);
37    /// ```
38    #[inline]
39    pub fn new(value: T) -> Self {
40        Self {
41            value,
42            probability: 1.0,
43        }
44    }
45
46    /// Create a probabilistic value with specified probability
47    ///
48    /// # Panics
49    ///
50    /// Panics if probability is not in [0, 1]
51    ///
52    /// # Examples
53    ///
54    /// ```
55    /// use amari_flynn::prob::Prob;
56    ///
57    /// let rare = Prob::with_probability(0.01, "miracle");
58    /// assert_eq!(rare.probability(), 0.01);
59    /// ```
60    #[inline]
61    pub fn with_probability(probability: f64, value: T) -> Self {
62        assert!(
63            (0.0..=1.0).contains(&probability),
64            "Probability must be in [0, 1]"
65        );
66        Self { value, probability }
67    }
68
69    /// Get the probability associated with this value
70    #[inline]
71    pub fn probability(&self) -> f64 {
72        self.probability
73    }
74
75    /// Extract the inner value
76    #[inline]
77    pub fn into_inner(self) -> T {
78        self.value
79    }
80
81    /// Get a reference to the inner value
82    #[inline]
83    pub fn inner(&self) -> &T {
84        &self.value
85    }
86
87    /// Map over the inner value while preserving probability
88    ///
89    /// # Examples
90    ///
91    /// ```
92    /// use amari_flynn::prob::Prob;
93    ///
94    /// let x = Prob::with_probability(0.5, 10);
95    /// let y = x.map(|v| v * 2);
96    /// assert_eq!(y.into_inner(), 20);
97    /// assert_eq!(y.probability(), 0.5);
98    /// ```
99    #[inline]
100    pub fn map<U, F>(self, f: F) -> Prob<U>
101    where
102        F: FnOnce(T) -> U,
103    {
104        Prob {
105            value: f(self.value),
106            probability: self.probability,
107        }
108    }
109
110    /// Monadic bind for probabilistic values
111    ///
112    /// Combines probabilities multiplicatively (assumes independence)
113    #[inline]
114    pub fn and_then<U, F>(self, f: F) -> Prob<U>
115    where
116        F: FnOnce(T) -> Prob<U>,
117    {
118        let result = f(self.value);
119        Prob {
120            value: result.value,
121            probability: self.probability * result.probability,
122        }
123    }
124
125    /// Sample this probabilistic value
126    ///
127    /// Returns Some(value) with probability p, None otherwise
128    ///
129    /// # Examples
130    ///
131    /// ```
132    /// use amari_flynn::prob::Prob;
133    ///
134    /// let coin = Prob::with_probability(0.5, "heads");
135    /// let result = coin.sample(&mut rand::thread_rng());
136    /// // result is either Some("heads") or None
137    /// ```
138    #[inline]
139    pub fn sample<R: Rng>(self, rng: &mut R) -> Option<T> {
140        if rng.gen::<f64>() < self.probability {
141            Some(self.value)
142        } else {
143            None
144        }
145    }
146}
147
148/// Phantom type for probability density (future verification use)
149///
150/// This ghost type exists for future formal verification integration
151/// where we can prove properties about probability distributions.
152#[derive(Clone, Copy, Debug)]
153pub struct ProbabilityDensity<T> {
154    _phantom: core::marker::PhantomData<T>,
155}
156
157impl<T> ProbabilityDensity<T> {
158    /// Create a phantom probability density
159    pub fn new() -> Self {
160        Self {
161            _phantom: core::marker::PhantomData,
162        }
163    }
164}
165
166impl<T> Default for ProbabilityDensity<T> {
167    fn default() -> Self {
168        Self::new()
169    }
170}
171
172#[cfg(test)]
173mod tests {
174    use super::*;
175
176    #[test]
177    fn test_new_has_probability_one() {
178        let p = Prob::new(42);
179        assert_eq!(p.probability(), 1.0);
180        assert_eq!(p.into_inner(), 42);
181    }
182
183    #[test]
184    fn test_with_probability() {
185        let p = Prob::with_probability(0.5, "test");
186        assert_eq!(p.probability(), 0.5);
187        assert_eq!(p.into_inner(), "test");
188    }
189
190    #[test]
191    #[should_panic]
192    fn test_invalid_probability_panics() {
193        let _ = Prob::with_probability(1.5, ());
194    }
195
196    #[test]
197    fn test_map_preserves_probability() {
198        let p = Prob::with_probability(0.3, 10);
199        let q = p.map(|x| x * 2);
200        assert_eq!(q.into_inner(), 20);
201        assert_eq!(q.probability(), 0.3);
202    }
203
204    #[test]
205    fn test_and_then_multiplies_probabilities() {
206        let p = Prob::with_probability(0.5, 10);
207        let q = p.and_then(|x| Prob::with_probability(0.4, x + 5));
208        assert_eq!(q.into_inner(), 15);
209        assert!((q.probability() - 0.2).abs() < 1e-10);
210    }
211}