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
11pub type Integer = i64;
13
14#[derive(Debug, Clone, GetSize)]
16pub enum IntegerExpr<V>
17where
18 V: Clone,
19{
20 Const(Integer),
25 Var(V),
27 Nat(NaturalExpr<V>),
29 Rand(Box<(IntegerExpr<V>, IntegerExpr<V>)>),
34 Opposite(Box<IntegerExpr<V>>),
39 Sum(Vec<IntegerExpr<V>>),
41 Product(Vec<IntegerExpr<V>>),
43 Div(Box<(IntegerExpr<V>, IntegerExpr<V>)>),
45 Rem(Box<(IntegerExpr<V>, IntegerExpr<V>)>),
47 Floor(Box<FloatExpr<V>>),
49 Ceil(Box<FloatExpr<V>>),
51 Min(Box<(IntegerExpr<V>, IntegerExpr<V>)>),
56 Max(Box<(IntegerExpr<V>, IntegerExpr<V>)>),
58 Ite(Box<(BooleanExpr<V>, IntegerExpr<V>, IntegerExpr<V>)>),
65}
66
67impl<V> IntegerExpr<V>
68where
69 V: Copy,
70{
71 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 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}