Skip to main content

scan_core/grammar/
integer.rs

1use std::ops::{Add, Div, Mul, Neg};
2
3use get_size2::GetSize;
4use rand::{Rng, RngExt};
5
6use crate::{
7    Expression, Type, TypeError, Val,
8    grammar::{BooleanExpr, FloatExpr, NaturalExpr},
9};
10
11/// Integer values.
12pub type Integer = i64;
13
14/// Integer expressions.
15#[derive(Debug, Clone, GetSize)]
16pub enum IntegerExpr<V>
17where
18    V: Clone,
19{
20    // -------------------
21    // General expressions
22    // -------------------
23    /// A constant value.
24    Const(Integer),
25    /// A typed variable.
26    Var(V),
27    /// Conversion from Natural
28    Nat(NaturalExpr<V>),
29    // -------------
30    // Random values
31    // -------------
32    /// A random integer between a lower bound (included) and an upper bound (excluded).
33    Rand(Box<(IntegerExpr<V>, IntegerExpr<V>)>),
34    // --------------------
35    // Arithmetic operators
36    // --------------------
37    /// Opposite of a numerical expression.
38    Opposite(Box<IntegerExpr<V>>),
39    /// Arithmetic n-ary sum.
40    Sum(Vec<IntegerExpr<V>>),
41    /// Arithmetic n-ary multiplication.
42    Product(Vec<IntegerExpr<V>>),
43    /// Div operation
44    Div(Box<(IntegerExpr<V>, IntegerExpr<V>)>),
45    /// Rem operation
46    Rem(Box<(IntegerExpr<V>, IntegerExpr<V>)>),
47    /// Floor
48    Floor(Box<FloatExpr<V>>),
49    /// Ceil
50    Ceil(Box<FloatExpr<V>>),
51    // --------------------
52    // Order operators
53    // --------------------
54    /// Min of two values
55    Min(Box<(IntegerExpr<V>, IntegerExpr<V>)>),
56    /// Max of two values
57    Max(Box<(IntegerExpr<V>, IntegerExpr<V>)>),
58    // -----
59    // Flow
60    // -----
61    /// If-Then-Else construct, where If must be a boolean expression,
62    /// Then and Else must have the same type,
63    /// and this is also the type of the whole expression.
64    Ite(Box<(BooleanExpr<V>, IntegerExpr<V>, IntegerExpr<V>)>),
65}
66
67impl<V> IntegerExpr<V>
68where
69    V: Copy,
70{
71    /// Returns `true` if the expression is constant, i.e., it contains no variables, and `false` otherwise.
72    pub fn is_constant(&self) -> bool {
73        match self {
74            IntegerExpr::Const(_) => true,
75            IntegerExpr::Var(_) | IntegerExpr::Rand(_) => false,
76            IntegerExpr::Nat(natural_expr) => natural_expr.is_constant(),
77            IntegerExpr::Opposite(integer_expr) => integer_expr.is_constant(),
78            IntegerExpr::Sum(integer_exprs) | IntegerExpr::Product(integer_exprs) => {
79                integer_exprs.iter().all(IntegerExpr::is_constant)
80            }
81            IntegerExpr::Div(args)
82            | IntegerExpr::Rem(args)
83            | IntegerExpr::Min(args)
84            | IntegerExpr::Max(args) => {
85                let (lhs, rhs) = args.as_ref();
86                lhs.is_constant() && rhs.is_constant()
87            }
88            IntegerExpr::Floor(float_expr) | IntegerExpr::Ceil(float_expr) => {
89                float_expr.is_constant()
90            }
91            IntegerExpr::Ite(args) => {
92                let (ite, lhs, rhs) = args.as_ref();
93                ite.is_constant() && lhs.is_constant() && rhs.is_constant()
94            }
95        }
96    }
97
98    /// Returns the [`Integer`] value computed from the expression,
99    /// given the variable evaluation.
100    /// It panics if the evaluation is not possible, including:
101    ///
102    /// - If a variable is not included in the evaluation;
103    /// - If a variable included in the evaluation is not of [`Integer`] type;
104    /// - Division by 0;
105    /// - Overflow.
106    pub fn eval<R: Rng>(&self, vars: &dyn Fn(V) -> Val, mut rng: Option<&mut R>) -> Integer {
107        match self {
108            IntegerExpr::Const(int) => *int,
109            IntegerExpr::Var(var) => {
110                if let Val::Integer(int) = vars(*var) {
111                    int
112                } else {
113                    panic!("type mismatch: expected natural variable")
114                }
115            }
116            IntegerExpr::Nat(natural_expr) => natural_expr.eval(vars, rng) as Integer,
117            IntegerExpr::Rand(bounds) => {
118                let (lower_bound_expr, upper_bound_expr) = bounds.as_ref();
119                let lower_bound = lower_bound_expr.eval(vars, rng.as_deref_mut());
120                let upper_bound = upper_bound_expr.eval(vars, rng.as_deref_mut());
121                rng.expect("rng").random_range(lower_bound..upper_bound)
122            }
123            IntegerExpr::Opposite(integer_expr) => integer_expr.eval(vars, rng).strict_neg(),
124            IntegerExpr::Sum(integer_exprs) => integer_exprs.iter().fold(0, |acc, expr| {
125                acc.strict_add(expr.eval(vars, rng.as_deref_mut()))
126            }),
127            IntegerExpr::Product(integer_exprs) => integer_exprs.iter().fold(1, |acc, expr| {
128                acc.strict_mul(expr.eval(vars, rng.as_deref_mut()))
129            }),
130            IntegerExpr::Div(args) => {
131                let (lhs_expr, rhs_expr) = args.as_ref();
132                let lhs = lhs_expr.eval(vars, rng.as_deref_mut());
133                let rhs = rhs_expr.eval(vars, rng);
134                lhs.strict_div(rhs)
135            }
136            IntegerExpr::Rem(args) => {
137                let (lhs_expr, rhs_expr) = args.as_ref();
138                let lhs = lhs_expr.eval(vars, rng.as_deref_mut());
139                let rhs = rhs_expr.eval(vars, rng);
140                lhs.strict_rem_euclid(rhs)
141            }
142            IntegerExpr::Floor(float_expr) => float_expr.eval(vars, rng).floor() as Integer,
143            IntegerExpr::Ceil(float_expr) => float_expr.eval(vars, rng).ceil() as Integer,
144            IntegerExpr::Ite(args) => {
145                let (ite, lhs, rhs) = args.as_ref();
146                if ite.eval(vars, rng.as_deref_mut()) {
147                    lhs.eval(vars, rng)
148                } else {
149                    rhs.eval(vars, rng)
150                }
151            }
152            IntegerExpr::Min(args) => {
153                let (lhs_expr, rhs_expr) = args.as_ref();
154                let lhs = lhs_expr.eval(vars, rng.as_deref_mut());
155                let rhs = rhs_expr.eval(vars, rng);
156                lhs.min(rhs)
157            }
158            IntegerExpr::Max(args) => {
159                let (lhs_expr, rhs_expr) = args.as_ref();
160                let lhs = lhs_expr.eval(vars, rng.as_deref_mut());
161                let rhs = rhs_expr.eval(vars, rng);
162                lhs.max(rhs)
163            }
164        }
165    }
166
167    pub(crate) fn map<W: Clone>(self, map: &dyn Fn(V) -> W) -> IntegerExpr<W> {
168        match self {
169            IntegerExpr::Const(i) => IntegerExpr::Const(i),
170            IntegerExpr::Var(var) => IntegerExpr::Var(map(var)),
171            IntegerExpr::Nat(nat_expr) => IntegerExpr::Nat(nat_expr.map(map)),
172            IntegerExpr::Rand(bounds) => {
173                let (lower_bound, upper_bound) = *bounds;
174                IntegerExpr::Rand(Box::new((lower_bound.map(map), upper_bound.map(map))))
175            }
176            IntegerExpr::Opposite(integer_expr) => {
177                IntegerExpr::Opposite(Box::new(integer_expr.map(map)))
178            }
179            IntegerExpr::Sum(integer_exprs) => IntegerExpr::Sum(
180                integer_exprs
181                    .into_iter()
182                    .map(|expr| expr.map(map))
183                    .collect(),
184            ),
185            IntegerExpr::Product(integer_exprs) => IntegerExpr::Product(
186                integer_exprs
187                    .into_iter()
188                    .map(|expr| expr.map(map))
189                    .collect(),
190            ),
191            IntegerExpr::Div(args) => {
192                let (lhs, rhs) = *args;
193                IntegerExpr::Div(Box::new((lhs.map(map), rhs.map(map))))
194            }
195            IntegerExpr::Rem(args) => {
196                let (lhs, rhs) = *args;
197                IntegerExpr::Rem(Box::new((lhs.map(map), rhs.map(map))))
198            }
199            IntegerExpr::Floor(float_expr) => IntegerExpr::Floor(Box::new(float_expr.map(map))),
200            IntegerExpr::Ceil(float_expr) => IntegerExpr::Ceil(Box::new(float_expr.map(map))),
201            IntegerExpr::Ite(args) => {
202                let (r#if, then, r#else) = *args;
203                IntegerExpr::Ite(Box::new((r#if.map(map), then.map(map), r#else.map(map))))
204            }
205            IntegerExpr::Min(args) => {
206                let (lhs, rhs) = *args;
207                IntegerExpr::Min(Box::new((lhs.map(map), rhs.map(map))))
208            }
209            IntegerExpr::Max(args) => {
210                let (lhs, rhs) = *args;
211                IntegerExpr::Max(Box::new((lhs.map(map), rhs.map(map))))
212            }
213        }
214    }
215
216    pub(crate) fn context(&self, vars: &dyn Fn(V) -> Option<Type>) -> Result<(), TypeError> {
217        match self {
218            IntegerExpr::Const(_) => Ok(()),
219            IntegerExpr::Var(v) => matches!(vars(*v), Some(Type::Integer))
220                .then_some(())
221                .ok_or(TypeError::TypeMismatch),
222            IntegerExpr::Nat(natural_expr) => natural_expr.context(vars),
223            IntegerExpr::Rand(exprs)
224            | IntegerExpr::Div(exprs)
225            | IntegerExpr::Rem(exprs)
226            | IntegerExpr::Min(exprs)
227            | IntegerExpr::Max(exprs) => exprs.0.context(vars).and_then(|()| exprs.1.context(vars)),
228            IntegerExpr::Opposite(integer_expr) => integer_expr.context(vars),
229            IntegerExpr::Sum(integer_exprs) | IntegerExpr::Product(integer_exprs) => {
230                integer_exprs.iter().try_for_each(|expr| expr.context(vars))
231            }
232            IntegerExpr::Floor(float_expr) | IntegerExpr::Ceil(float_expr) => {
233                float_expr.context(vars)
234            }
235            IntegerExpr::Ite(exprs) => exprs
236                .0
237                .context(vars)
238                .and_then(|()| exprs.1.context(vars))
239                .and_then(|()| exprs.2.context(vars)),
240        }
241    }
242}
243
244impl<V> From<Integer> for IntegerExpr<V>
245where
246    V: Clone,
247{
248    fn from(value: Integer) -> Self {
249        Self::Const(value)
250    }
251}
252
253impl<V> TryFrom<Expression<V>> for IntegerExpr<V>
254where
255    V: Clone,
256{
257    type Error = TypeError;
258
259    fn try_from(value: Expression<V>) -> Result<Self, Self::Error> {
260        match value {
261            Expression::Boolean(_) | Expression::Float(_) => Err(TypeError::TypeMismatch),
262            Expression::Natural(natural_expr) => Ok(IntegerExpr::Nat(natural_expr)),
263            Expression::Integer(integer_expr) => Ok(integer_expr),
264        }
265    }
266}
267
268impl<V> From<NaturalExpr<V>> for IntegerExpr<V>
269where
270    V: Clone,
271{
272    fn from(value: NaturalExpr<V>) -> Self {
273        Self::Nat(value)
274    }
275}
276
277impl<V> Add for IntegerExpr<V>
278where
279    V: Clone,
280{
281    type Output = Self;
282
283    fn add(mut self, mut rhs: Self) -> Self::Output {
284        if let IntegerExpr::Sum(ref mut exprs) = self {
285            if let IntegerExpr::Sum(rhs_exprs) = rhs {
286                exprs.extend(rhs_exprs);
287            } else {
288                exprs.push(rhs);
289            }
290            self
291        } else if let IntegerExpr::Sum(ref mut rhs_exprs) = rhs {
292            rhs_exprs.push(self);
293            rhs
294        } else {
295            IntegerExpr::Sum(vec![self, rhs])
296        }
297    }
298}
299
300impl<V> Mul for IntegerExpr<V>
301where
302    V: Clone,
303{
304    type Output = Self;
305
306    fn mul(mut self, mut rhs: Self) -> Self::Output {
307        if let IntegerExpr::Product(ref mut exprs) = self {
308            if let IntegerExpr::Product(rhs_exprs) = rhs {
309                exprs.extend(rhs_exprs);
310            } else {
311                exprs.push(rhs);
312            }
313            self
314        } else if let IntegerExpr::Product(ref mut rhs_exprs) = rhs {
315            rhs_exprs.push(self);
316            rhs
317        } else {
318            IntegerExpr::Product(vec![self, rhs])
319        }
320    }
321}
322
323impl<V> Div for IntegerExpr<V>
324where
325    V: Clone,
326{
327    type Output = Self;
328
329    fn div(self, rhs: Self) -> Self::Output {
330        IntegerExpr::Div(Box::new((self, rhs)))
331    }
332}
333
334impl<V> Neg for IntegerExpr<V>
335where
336    V: Clone,
337{
338    type Output = Self;
339
340    fn neg(self) -> Self::Output {
341        if let IntegerExpr::Opposite(expr) = self {
342            *expr
343        } else {
344            IntegerExpr::Opposite(Box::new(self))
345        }
346    }
347}