1use std::ops::{Add, Div, Mul, Rem};
2
3use get_size2::GetSize;
4use rand::{Rng, RngExt};
5
6use crate::{
7 Expression, Type, TypeError, Val,
8 grammar::{BooleanExpr, IntegerExpr},
9};
10
11pub type Natural = u64;
13
14#[derive(Debug, Clone, GetSize)]
16pub enum NaturalExpr<V>
17where
18 V: Clone,
19{
20 Const(Natural),
25 Var(V),
27 Rand(Box<(NaturalExpr<V>, NaturalExpr<V>)>),
32 Sum(Vec<NaturalExpr<V>>),
37 Product(Vec<NaturalExpr<V>>),
39 Rem(Box<(NaturalExpr<V>, NaturalExpr<V>)>),
41 Div(Box<(NaturalExpr<V>, NaturalExpr<V>)>),
43 Abs(Box<IntegerExpr<V>>),
45 Min(Box<(NaturalExpr<V>, NaturalExpr<V>)>),
50 Max(Box<(NaturalExpr<V>, NaturalExpr<V>)>),
52 Ite(Box<(BooleanExpr<V>, NaturalExpr<V>, NaturalExpr<V>)>),
59}
60
61impl<V> NaturalExpr<V>
62where
63 V: Copy,
64{
65 pub fn is_constant(&self) -> bool {
67 match self {
68 NaturalExpr::Const(_) => true,
69 NaturalExpr::Var(_) | NaturalExpr::Rand(_) => false,
70 NaturalExpr::Sum(natural_exprs) | NaturalExpr::Product(natural_exprs) => {
71 natural_exprs.iter().all(NaturalExpr::is_constant)
72 }
73 NaturalExpr::Rem(args)
74 | NaturalExpr::Div(args)
75 | NaturalExpr::Min(args)
76 | NaturalExpr::Max(args) => {
77 let (lhs, rhs) = args.as_ref();
78 lhs.is_constant() && rhs.is_constant()
79 }
80 NaturalExpr::Abs(integer_expr) => integer_expr.is_constant(),
81 NaturalExpr::Ite(args) => {
82 let (ite, lhs, rhs) = args.as_ref();
83 ite.is_constant() && lhs.is_constant() && rhs.is_constant()
84 }
85 }
86 }
87
88 pub fn eval<R: Rng>(&self, vars: &dyn Fn(V) -> Val, mut rng: Option<&mut R>) -> Natural {
97 match self {
98 NaturalExpr::Const(nat) => *nat,
99 NaturalExpr::Var(var) => {
100 if let Val::Natural(nat) = vars(*var) {
101 nat
102 } else {
103 panic!("type mismatch: expected natural variable")
104 }
105 }
106 NaturalExpr::Rand(bounds) => {
107 let (lower_bound_expr, upper_bound_expr) = bounds.as_ref();
108 let lower_bound = lower_bound_expr.eval(vars, rng.as_deref_mut());
109 let upper_bound = upper_bound_expr.eval(vars, rng.as_deref_mut());
110 rng.as_mut()
111 .expect("rng")
112 .random_range(lower_bound..upper_bound)
113 }
114 NaturalExpr::Sum(natural_exprs) => natural_exprs.iter().fold(0, |acc, expr| {
115 acc.strict_add(expr.eval(vars, rng.as_deref_mut()))
116 }),
117 NaturalExpr::Product(natural_exprs) => natural_exprs.iter().fold(1, |acc, expr| {
118 acc.strict_mul(expr.eval(vars, rng.as_deref_mut()))
119 }),
120 NaturalExpr::Rem(args) => {
121 let (lhs_expr, rhs_expr) = args.as_ref();
122 let lhs = lhs_expr.eval(vars, rng.as_deref_mut());
123 let rhs = rhs_expr.eval(vars, rng);
124 lhs.strict_rem_euclid(rhs)
125 }
126 NaturalExpr::Div(args) => {
127 let (lhs_expr, rhs_expr) = args.as_ref();
128 let lhs = lhs_expr.eval(vars, rng.as_deref_mut());
129 let rhs = rhs_expr.eval(vars, rng);
130 lhs / rhs
131 }
132 NaturalExpr::Abs(integer_expr) => integer_expr.eval(vars, rng).unsigned_abs(),
133 NaturalExpr::Ite(args) => {
134 let (ite, lhs, rhs) = args.as_ref();
135 if ite.eval(vars, rng.as_deref_mut()) {
136 lhs.eval(vars, rng)
137 } else {
138 rhs.eval(vars, rng)
139 }
140 }
141 NaturalExpr::Min(args) => {
142 let (lhs_expr, rhs_expr) = args.as_ref();
143 let lhs = lhs_expr.eval(vars, rng.as_deref_mut());
144 let rhs = rhs_expr.eval(vars, rng);
145 lhs.min(rhs)
146 }
147 NaturalExpr::Max(args) => {
148 let (lhs_expr, rhs_expr) = args.as_ref();
149 let lhs = lhs_expr.eval(vars, rng.as_deref_mut());
150 let rhs = rhs_expr.eval(vars, rng);
151 lhs.max(rhs)
152 }
153 }
154 }
155
156 pub(crate) fn map<W: Clone>(self, map: &dyn Fn(V) -> W) -> NaturalExpr<W> {
157 match self {
158 NaturalExpr::Const(n) => NaturalExpr::Const(n),
159 NaturalExpr::Var(var) => NaturalExpr::Var(map(var)),
160 NaturalExpr::Rand(bounds) => {
161 let (lower_bound, upper_bound) = *bounds;
162 NaturalExpr::Rand(Box::new((lower_bound.map(map), upper_bound.map(map))))
163 }
164 NaturalExpr::Sum(natural_exprs) => NaturalExpr::Sum(
165 natural_exprs
166 .into_iter()
167 .map(|expr| expr.map(map))
168 .collect(),
169 ),
170 NaturalExpr::Product(natural_exprs) => NaturalExpr::Product(
171 natural_exprs
172 .into_iter()
173 .map(|expr| expr.map(map))
174 .collect(),
175 ),
176 NaturalExpr::Rem(args) => {
177 let (lhs, rhs) = *args;
178 NaturalExpr::Rem(Box::new((lhs.map(map), rhs.map(map))))
179 }
180 NaturalExpr::Div(args) => {
181 let (lhs, rhs) = *args;
182 NaturalExpr::Div(Box::new((lhs.map(map), rhs.map(map))))
183 }
184 NaturalExpr::Abs(integer_expr) => NaturalExpr::Abs(Box::new(integer_expr.map(map))),
185 NaturalExpr::Ite(args) => {
186 let (r#if, then, r#else) = *args;
187 NaturalExpr::Ite(Box::new((r#if.map(map), then.map(map), r#else.map(map))))
188 }
189 NaturalExpr::Min(args) => {
190 let (lhs, rhs) = *args;
191 NaturalExpr::Min(Box::new((lhs.map(map), rhs.map(map))))
192 }
193 NaturalExpr::Max(args) => {
194 let (lhs, rhs) = *args;
195 NaturalExpr::Max(Box::new((lhs.map(map), rhs.map(map))))
196 }
197 }
198 }
199
200 pub(crate) fn context(&self, vars: &dyn Fn(V) -> Option<Type>) -> Result<(), TypeError> {
201 match self {
202 NaturalExpr::Const(_) => Ok(()),
203 NaturalExpr::Var(v) => matches!(vars(*v), Some(Type::Natural))
204 .then_some(())
205 .ok_or(TypeError::TypeMismatch),
206 NaturalExpr::Rand(exprs)
207 | NaturalExpr::Div(exprs)
208 | NaturalExpr::Rem(exprs)
209 | NaturalExpr::Min(exprs)
210 | NaturalExpr::Max(exprs) => exprs.0.context(vars).and_then(|()| exprs.1.context(vars)),
211 NaturalExpr::Sum(integer_exprs) | NaturalExpr::Product(integer_exprs) => {
212 integer_exprs.iter().try_for_each(|expr| expr.context(vars))
213 }
214 NaturalExpr::Ite(exprs) => exprs
215 .0
216 .context(vars)
217 .and_then(|()| exprs.1.context(vars))
218 .and_then(|()| exprs.2.context(vars)),
219 NaturalExpr::Abs(integer_expr) => integer_expr.context(vars),
220 }
221 }
222}
223
224impl<V: Clone> From<Natural> for NaturalExpr<V> {
225 fn from(value: Natural) -> Self {
226 NaturalExpr::Const(value)
227 }
228}
229
230impl<V> TryFrom<Expression<V>> for NaturalExpr<V>
231where
232 V: Clone,
233{
234 type Error = TypeError;
235
236 fn try_from(value: Expression<V>) -> Result<Self, Self::Error> {
237 if let Expression::Natural(nat_expr) = value {
238 Ok(nat_expr)
239 } else {
240 Err(TypeError::TypeMismatch)
241 }
242 }
243}
244
245impl<V> Add for NaturalExpr<V>
246where
247 V: Clone,
248{
249 type Output = Self;
250
251 fn add(mut self, mut rhs: Self) -> Self::Output {
252 if let NaturalExpr::Sum(ref mut exprs) = self {
253 if let NaturalExpr::Sum(rhs_exprs) = rhs {
254 exprs.extend(rhs_exprs);
255 } else {
256 exprs.push(rhs);
257 }
258 self
259 } else if let NaturalExpr::Sum(ref mut rhs_exprs) = rhs {
260 rhs_exprs.push(self);
261 rhs
262 } else {
263 NaturalExpr::Sum(vec![self, rhs])
264 }
265 }
266}
267
268impl<V> Mul for NaturalExpr<V>
269where
270 V: Clone,
271{
272 type Output = Self;
273
274 fn mul(mut self, mut rhs: Self) -> Self::Output {
275 if let NaturalExpr::Product(ref mut exprs) = self {
276 if let NaturalExpr::Product(rhs_exprs) = rhs {
277 exprs.extend(rhs_exprs);
278 } else {
279 exprs.push(rhs);
280 }
281 self
282 } else if let NaturalExpr::Product(ref mut rhs_exprs) = rhs {
283 rhs_exprs.push(self);
284 rhs
285 } else {
286 NaturalExpr::Product(vec![self, rhs])
287 }
288 }
289}
290
291impl<V> Div for NaturalExpr<V>
292where
293 V: Clone,
294{
295 type Output = Self;
296
297 fn div(self, rhs: Self) -> Self::Output {
298 NaturalExpr::Div(Box::new((self, rhs)))
299 }
300}
301
302impl<V> Rem for NaturalExpr<V>
303where
304 V: Clone,
305{
306 type Output = Self;
307
308 fn rem(self, rhs: Self) -> Self::Output {
309 NaturalExpr::Rem(Box::new((self, rhs)))
310 }
311}