Skip to main content

bimm_contracts/
expressions.rs

1//! # Dimension Expressions.
2
3use crate::math::maybe_iroot;
4use core::fmt::{Display, Formatter};
5
6/// A stack/static expression algebra for dimension sizes.
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub enum DimExpr<'a> {
9    /// A parameter reference.
10    Param {
11        /// The id of the parameter.
12        id: usize,
13    },
14
15    /// Negation of an expression.
16    Negate {
17        /// The child expression.
18        child: &'a DimExpr<'a>,
19    },
20
21    /// Exponentiation of an expression.
22    Pow {
23        /// The child expression.
24        base: &'a DimExpr<'a>,
25
26        /// The exponent.
27        exp: usize,
28    },
29
30    /// Sum of expressions.
31    Sum {
32        /// The child expressions.
33        children: &'a [DimExpr<'a>],
34    },
35
36    /// Product of expressions.
37    Prod {
38        /// The child expressions.
39        children: &'a [DimExpr<'a>],
40    },
41}
42
43/// Display Adapter to format `DimExprs` with a `Index`.
44pub struct ExprDisplayAdapter<'a> {
45    ///  index.
46    pub index: &'a [&'a str],
47
48    /// Expression to format.
49    pub expr: &'a DimExpr<'a>,
50}
51
52impl<'a> Display for ExprDisplayAdapter<'a> {
53    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
54        match self.expr {
55            DimExpr::Param { id } => write!(f, "{}", self.index[*id]),
56            DimExpr::Negate { child } => write!(
57                f,
58                "(-{})",
59                ExprDisplayAdapter {
60                    expr: child,
61                    index: self.index
62                }
63            ),
64            DimExpr::Pow { base: child, exp } => write!(
65                f,
66                "({}^{})",
67                ExprDisplayAdapter {
68                    expr: child,
69                    index: self.index
70                },
71                exp
72            ),
73            DimExpr::Sum { children } => {
74                write!(f, "(")?;
75                for (idx, expr) in children.iter().enumerate() {
76                    if idx > 0 {
77                        write!(f, "+")?;
78                    }
79                    write!(
80                        f,
81                        "{}",
82                        ExprDisplayAdapter {
83                            expr,
84                            index: self.index
85                        }
86                    )?;
87                }
88                write!(f, ")")
89            }
90            DimExpr::Prod { children } => {
91                write!(f, "(")?;
92                for (idx, expr) in children.iter().enumerate() {
93                    if idx > 0 {
94                        write!(f, "*")?;
95                    }
96                    write!(
97                        f,
98                        "{}",
99                        ExprDisplayAdapter {
100                            expr,
101                            index: self.index
102                        }
103                    )?;
104                }
105                write!(f, ")")
106            }
107        }
108    }
109}
110
111#[derive(Debug, Clone, PartialEq, Eq)]
112enum EvalResult {
113    /// The evaluated value of the expression.
114    Value { value: isize },
115
116    /// The count of unbound parameters in the expression.
117    UnboundParams {
118        /// The count of unbound parameters.
119        count: usize,
120    },
121}
122
123/// Result of `SizeExpr::try_match()`.
124///
125/// All values are borrowed from the original expression,
126/// so they are valid as long as the expression is valid.
127///
128/// Runtime errors (malformed expressions, too-many unbound parameters, etc.)
129/// are not represented here; and are returned as `Err(String)` from `try_match`.
130#[derive(Debug, Clone, PartialEq, Eq)]
131pub enum MatchResult {
132    /// All params bound and expression equals target.
133    Match,
134
135    /// Expression value does not match the target.
136    Conflict,
137
138    /// Expression can be solved for a single unbound param.
139    ParamConstraint {
140        /// The id of the parameter.
141        id: usize,
142
143        /// The value the parameter must take to satisfy the expression.
144        value: isize,
145    },
146}
147
148impl<'a> DimExpr<'a> {
149    /// Evaluate an expression.
150    ///
151    /// ## Arguments
152    ///
153    /// - `env` - the binding environment.
154    ///
155    /// ## Returns
156    ///
157    /// A `TryEvalResult`:
158    /// * `Value(value)` - the evaluated value of the expression.
159    /// * `UnboundParams(count)` - the count of unbound parameters.
160    #[must_use]
161    fn try_eval(&self, env: &[Option<isize>]) -> EvalResult {
162        #[inline(always)]
163        fn reduce_children<'a>(
164            exprs: &'a [DimExpr<'a>],
165            env: &[Option<isize>],
166            zero: isize,
167            op: fn(&mut isize, isize),
168        ) -> EvalResult {
169            let mut value = zero;
170            let mut count = 0;
171            for expr in exprs {
172                match expr.try_eval(env) {
173                    EvalResult::Value { value: v } => op(&mut value, v),
174                    EvalResult::UnboundParams { count: c } => count += c,
175                }
176            }
177            if count == 0 {
178                EvalResult::Value { value }
179            } else {
180                EvalResult::UnboundParams { count }
181            }
182        }
183
184        match self {
185            DimExpr::Param { id } => match env[*id] {
186                Some(value) => EvalResult::Value { value },
187                None => EvalResult::UnboundParams { count: 1 },
188            },
189            DimExpr::Negate { child } => match child.try_eval(env) {
190                EvalResult::Value { value } => EvalResult::Value { value: -value },
191                x => x,
192            },
193            DimExpr::Pow { base: child, exp } => match child.try_eval(env) {
194                EvalResult::Value { value } => EvalResult::Value {
195                    value: value.pow(*exp as u32),
196                },
197                x => x,
198            },
199            DimExpr::Sum { children } => {
200                reduce_children(children, env, 0, |tmp, value| *tmp += value)
201            }
202            DimExpr::Prod { children } => {
203                reduce_children(children, env, 1, |tmp, value| *tmp *= value)
204            }
205        }
206    }
207
208    /// Reconcile an expression against a target value.
209    ///
210    /// ## Arguments
211    ///
212    /// * `target`: The target value to match.
213    /// * `env`: The environment containing bindings for parameters.
214    ///
215    /// ## Returns
216    ///
217    /// * `Ok(MatchResult::Match)` if the expression matches the target.
218    /// * `Ok(MatchResult::MissMatch)` if the expression does not match the target.
219    /// * `Ok(MatchResult::Constraint(name, value))` if the expression can be solved for a single unbound parameter.
220    /// * `Ok(MatchResult::UnderConstrained)` if the expression cannot be solved with the current bindings.
221    #[must_use]
222    pub fn try_match(
223        &self,
224        target: isize,
225        env: &[Option<isize>],
226    ) -> Result<MatchResult, &'static str> {
227        #[inline(always)]
228        fn reduce_children<'a>(
229            exprs: &'a [DimExpr<'a>],
230            env: &[Option<isize>],
231            zero: isize,
232            op: fn(&mut isize, isize),
233        ) -> Result<(isize, Option<&'a DimExpr<'a>>), &'static str> {
234            let mut partial_value: isize = zero;
235            let mut rem_expr = None;
236            // At most one child can be unbound, and by only one parameter.
237            for expr in exprs {
238                match expr.try_eval(env) {
239                    EvalResult::Value { value } => op(&mut partial_value, value),
240                    EvalResult::UnboundParams { count } => {
241                        if count == 1 && rem_expr.is_none() {
242                            rem_expr = Some(expr);
243                        } else {
244                            return Err("Too many unbound params.");
245                        }
246                    }
247                }
248            }
249            // If the monoid is fully bound, then return (value, None);
250            // Otherwise, return (partial_value, expr).
251            Ok((partial_value, rem_expr))
252        }
253
254        match self {
255            DimExpr::Param { id } => {
256                let id = *id;
257                if let Some(value) = env[id] {
258                    if value == target {
259                        Ok(MatchResult::Match)
260                    } else {
261                        Ok(MatchResult::Conflict)
262                    }
263                } else {
264                    Ok(MatchResult::ParamConstraint { id, value: target })
265                }
266            }
267            DimExpr::Negate { child } => child.try_match(-target, env),
268            DimExpr::Pow { base: child, exp } => match maybe_iroot(target, *exp) {
269                Some(root) => child.try_match(root, env),
270                None => Err("No integer solution."),
271            },
272            DimExpr::Sum { children } => {
273                let (value, rem) = reduce_children(children, env, 0, |tmp, value| *tmp += value)?;
274                if let Some(expr) = rem {
275                    expr.try_match(target - value, env)
276                } else if value == target {
277                    Ok(MatchResult::Match)
278                } else {
279                    Ok(MatchResult::Conflict)
280                }
281            }
282            DimExpr::Prod { children } => {
283                let (value, rem) = reduce_children(children, env, 1, |tmp, value| *tmp *= value)?;
284                if let Some(expr) = rem {
285                    if target % value != 0 {
286                        // Non-integer solution
287                        return Err("No integer solution.");
288                    }
289                    expr.try_match(target / value, env)
290                } else if value == target {
291                    Ok(MatchResult::Match)
292                } else {
293                    Ok(MatchResult::Conflict)
294                }
295            }
296        }
297    }
298}
299
300#[cfg(test)]
301mod tests {
302    use super::*;
303    use alloc::format;
304    use alloc::string::String;
305
306    #[test]
307    fn test_format() {
308        static INDEX: [&str; 5] = ["a", "b", "c", "d", "e"];
309
310        fn fmt(expr: &DimExpr) -> String {
311            format!(
312                "{}",
313                ExprDisplayAdapter {
314                    expr: &expr,
315                    index: &INDEX
316                }
317            )
318        }
319
320        assert_eq!(
321            fmt(&DimExpr::Param {
322                id: INDEX
323                    .iter()
324                    .enumerate()
325                    .find_map(|(i, k)| if *k == "a" { Some(i) } else { None })
326                    .unwrap()
327            }),
328            "a"
329        );
330
331        let _expr = DimExpr::Prod {
332            children: &[
333                DimExpr::Param { id: 0 },
334                DimExpr::Param { id: 1 },
335                DimExpr::Sum {
336                    children: &[
337                        DimExpr::Param { id: 2 },
338                        DimExpr::Pow {
339                            base: &DimExpr::Param { id: 3 },
340                            exp: 2,
341                        },
342                        DimExpr::Negate {
343                            child: &DimExpr::Param { id: 4 },
344                        },
345                    ],
346                },
347            ],
348        };
349        assert_eq!(fmt(&_expr), "(a*b*(c+(d^2)+(-e)))");
350    }
351
352    #[test]
353    fn test_eval_param() {
354        let env = [Some(5), None];
355
356        let expr = DimExpr::Param { id: 0 };
357        assert_eq!(expr.try_eval(&env), EvalResult::Value { value: 5 });
358        assert_eq!(expr.try_match(5, &env), Ok(MatchResult::Match));
359        assert_eq!(expr.try_match(42, &env), Ok(MatchResult::Conflict));
360
361        let expr = DimExpr::Param { id: 1 };
362        assert_eq!(expr.try_eval(&env), EvalResult::UnboundParams { count: 1 });
363        assert_eq!(
364            expr.try_match(5, &env),
365            Ok(MatchResult::ParamConstraint { id: 1, value: 5 })
366        );
367    }
368
369    #[test]
370    fn try_eval_negate() {
371        let expr = DimExpr::Negate {
372            child: &DimExpr::Param { id: 0 },
373        };
374
375        let env = [Some(5)];
376        assert_eq!(expr.try_eval(&env), EvalResult::Value { value: -5 });
377        assert_eq!(expr.try_match(-5, &env), Ok(MatchResult::Match));
378        assert_eq!(expr.try_match(42, &env), Ok(MatchResult::Conflict));
379
380        let env = [None];
381        assert_eq!(expr.try_eval(&env), EvalResult::UnboundParams { count: 1 });
382        assert_eq!(
383            expr.try_match(-5, &env),
384            Ok(MatchResult::ParamConstraint { id: 0, value: 5 })
385        );
386    }
387
388    #[test]
389    fn try_eval_pow() {
390        let expr = DimExpr::Pow {
391            base: &DimExpr::Param { id: 0 },
392            exp: 3,
393        };
394
395        let env = [Some(5)];
396        assert_eq!(expr.try_eval(&env), EvalResult::Value { value: 125 });
397        assert_eq!(expr.try_match(125, &env), Ok(MatchResult::Match));
398        assert_eq!(expr.try_match(42, &env), Err("No integer solution."));
399
400        let env = [None];
401        assert_eq!(expr.try_eval(&env), EvalResult::UnboundParams { count: 1 });
402        assert_eq!(
403            expr.try_match(125, &env),
404            Ok(MatchResult::ParamConstraint { id: 0, value: 5 })
405        );
406    }
407
408    #[test]
409    fn test_eval_sum() {
410        let expr = DimExpr::Sum {
411            children: &[
412                DimExpr::Param { id: 0 },
413                DimExpr::Param { id: 1 },
414                DimExpr::Param { id: 2 },
415                DimExpr::Param { id: 3 },
416            ],
417        };
418
419        let env = [Some(2), Some(3), Some(4), Some(5)];
420        assert_eq!(expr.try_eval(&env), EvalResult::Value { value: 14 });
421        assert_eq!(expr.try_match(14, &env), Ok(MatchResult::Match));
422        assert_eq!(expr.try_match(42, &env), Ok(MatchResult::Conflict));
423
424        let env = [Some(2), Some(3), None, Some(5)];
425        assert_eq!(expr.try_eval(&env), EvalResult::UnboundParams { count: 1 });
426        assert_eq!(
427            expr.try_match(14, &env),
428            Ok(MatchResult::ParamConstraint { id: 2, value: 4 })
429        );
430
431        let env = [Some(5), Some(3), None, None];
432        assert_eq!(expr.try_eval(&env), EvalResult::UnboundParams { count: 2 });
433        assert_eq!(expr.try_match(14, &env), Err("Too many unbound params."));
434    }
435
436    #[test]
437    fn test_eval_prod() {
438        let expr = DimExpr::Prod {
439            children: &[
440                DimExpr::Param { id: 0 },
441                DimExpr::Param { id: 1 },
442                DimExpr::Param { id: 2 },
443                DimExpr::Param { id: 3 },
444            ],
445        };
446
447        let env = [Some(2), Some(3), Some(4), Some(5)];
448        assert_eq!(expr.try_eval(&env), EvalResult::Value { value: 120 });
449        assert_eq!(expr.try_match(120, &env), Ok(MatchResult::Match));
450        assert_eq!(expr.try_match(42, &env), Ok(MatchResult::Conflict));
451
452        let env = [Some(2), Some(3), None, Some(5)];
453        assert_eq!(expr.try_eval(&env), EvalResult::UnboundParams { count: 1 });
454        assert_eq!(
455            expr.try_match(120, &env),
456            Ok(MatchResult::ParamConstraint { id: 2, value: 4 })
457        );
458
459        let env = [Some(1), Some(5), None, Some(5)];
460        assert_eq!(expr.try_match(40, &env), Err("No integer solution."));
461
462        let env = [Some(5), Some(3), None, None];
463        assert_eq!(expr.try_eval(&env), EvalResult::UnboundParams { count: 2 });
464        assert_eq!(expr.try_match(120, &env), Err("Too many unbound params."));
465    }
466}