Skip to main content

scan_core/grammar/
natural.rs

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
11/// Natural (unsigned) values.
12pub type Natural = u64;
13
14/// A [`Natural`] number expression
15#[derive(Debug, Clone, GetSize)]
16pub enum NaturalExpr<V>
17where
18    V: Clone,
19{
20    // -------------------
21    // General expressions
22    // -------------------
23    /// A constant value.
24    Const(Natural),
25    /// A typed variable.
26    Var(V),
27    // -------------
28    // Random values
29    // -------------
30    /// A random integer between a lower bound (included) and an upper bound (excluded).
31    Rand(Box<(NaturalExpr<V>, NaturalExpr<V>)>),
32    // --------------------
33    // Arithmetic operators
34    // --------------------
35    /// Arithmetic n-ary sum.
36    Sum(Vec<NaturalExpr<V>>),
37    /// Arithmetic n-ary multiplication.
38    Product(Vec<NaturalExpr<V>>),
39    /// Mod operation
40    Rem(Box<(NaturalExpr<V>, NaturalExpr<V>)>),
41    /// Div operation
42    Div(Box<(NaturalExpr<V>, NaturalExpr<V>)>),
43    /// Abs
44    Abs(Box<IntegerExpr<V>>),
45    // --------------------
46    // Order operators
47    // --------------------
48    /// Min of two values
49    Min(Box<(NaturalExpr<V>, NaturalExpr<V>)>),
50    /// Max of two values
51    Max(Box<(NaturalExpr<V>, NaturalExpr<V>)>),
52    // -----
53    // Flow
54    // -----
55    /// If-Then-Else construct, where If must be a boolean expression,
56    /// Then and Else must have the same type,
57    /// and this is also the type of the whole expression.
58    Ite(Box<(BooleanExpr<V>, NaturalExpr<V>, NaturalExpr<V>)>),
59}
60
61impl<V> NaturalExpr<V>
62where
63    V: Copy,
64{
65    /// Returns `true` if the expression is constant, i.e., it contains no variables, and `false` otherwise.
66    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    /// Returns the [`Natural`] value computed from the expression,
89    /// given the variable evaluation.
90    /// It panics if the evaluation is not possible, including:
91    ///
92    /// - If a variable is not included in the evaluation;
93    /// - If a variable included in the evaluation is not of [`Natural`] type;
94    /// - Division by 0;
95    /// - Overflow.
96    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}