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}