1use crate::math::maybe_iroot;
4use core::fmt::{Display, Formatter};
5
6#[derive(Debug, Clone, PartialEq, Eq)]
8pub enum DimExpr<'a> {
9 Param {
11 id: usize,
13 },
14
15 Negate {
17 child: &'a DimExpr<'a>,
19 },
20
21 Pow {
23 base: &'a DimExpr<'a>,
25
26 exp: usize,
28 },
29
30 Sum {
32 children: &'a [DimExpr<'a>],
34 },
35
36 Prod {
38 children: &'a [DimExpr<'a>],
40 },
41}
42
43pub struct ExprDisplayAdapter<'a> {
45 pub index: &'a [&'a str],
47
48 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 Value { value: isize },
115
116 UnboundParams {
118 count: usize,
120 },
121}
122
123#[derive(Debug, Clone, PartialEq, Eq)]
131pub enum MatchResult {
132 Match,
134
135 Conflict,
137
138 ParamConstraint {
140 id: usize,
142
143 value: isize,
145 },
146}
147
148impl<'a> DimExpr<'a> {
149 #[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 #[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 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 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 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}