1use std::cmp::max;
2
3use crate::{
4 expr::{Expr, VarId},
5 prettify::prettify,
6 utils::bimap::BiMap,
7 utils::cache::{LinearCache, LocalCache, MbaCache},
8 varint::make_mask,
9};
10
11use log::debug;
12
13pub use crate::utils::error::SolveError;
14
15mod lambda;
16mod merge_hidden;
17
18use lambda::{find_lambda_int, find_two_lambdas_int};
19
20pub(crate) const MAX_VARS: usize = 20;
23
24const MAX_SIMPLIFICATION_PASSES: usize = 8;
25
26fn sub_coeff(tt: &mut [u64], coeff: u64, index: usize, sublist: &[usize]) {
27 let are_vars_true = |i: usize| sublist[1..].iter().copied().all(|v| ((i >> v) & 1) == 1);
28
29 let gp_size = 1usize << sublist[0];
30 let period = 2 * gp_size;
31
32 let mut start = index;
33 while start < tt.len() {
34 for (i, e) in tt.iter_mut().enumerate().skip(start).take(gp_size) {
35 if sublist.len() == 1 || are_vars_true(i) {
36 *e = e.wrapping_sub(coeff);
37 }
38 }
39 start += period;
40 }
41}
42
43fn reduce_vars(e: Expr, var_map: &mut BiMap<VarId, VarId>, t: &mut usize) -> Expr {
45 match e {
46 Expr::Var(v) => {
47 let vv = if let Some(v) = var_map.get_by_left(&v) {
48 *v
49 } else {
50 let vv = (*t).into();
51 var_map.insert(v, vv);
52 *t += 1;
53 vv
54 };
55 Expr::Var(vv)
56 }
57
58 _ => e.map(|e| reduce_vars(e, var_map, t)),
59 }
60}
61
62fn restore_vars(e: Expr, var_map: &BiMap<VarId, VarId>) -> Result<Expr, SolveError> {
64 match e {
65 Expr::Var(v) => {
66 let vv = var_map
67 .get_by_right(&v)
68 .copied()
69 .ok_or(SolveError::UnknownVariable(v))?;
70 Ok(Expr::Var(vv))
71 }
72
73 _ => e.try_map(|e| restore_vars(e, var_map)),
74 }
75}
76
77struct MBASolver<'a, C: LinearCache> {
78 n: u8,
80
81 mask: u64,
83
84 non_linear_components: BiMap<VarId, Expr>,
86
87 t: usize,
89
90 degree: usize,
92
93 l_cache: &'a C,
95
96 options: SimplifyOptions,
98}
99
100impl<'a, C: LinearCache> MBASolver<'a, C> {
101 fn new(l_cache: &'a C, e: &Expr, n: u8, options: SimplifyOptions) -> Self {
103 Self {
104 non_linear_components: BiMap::new(),
105 t: e.get_vars().iter().copied().map(|v| v.0).max().unwrap_or(0) + 1,
106 degree: 1,
107 n,
108 mask: make_mask(n),
109 l_cache,
110 options,
111 }
112 }
113
114 fn poly_to_nonpoly(&self, e: Expr) -> Expr {
116 match &e {
117 Expr::Var(v) => {
118 if let Some(e) = self.non_linear_components.get_by_left(v) {
119 e.clone()
120 } else {
121 e
122 }
123 }
124
125 _ => e.map(|e| self.poly_to_nonpoly(e)),
126 }
127 }
128
129 fn solve(&mut self, e: Expr) -> Result<Expr, SolveError> {
131 let e = e.reduce_masked(self.mask);
133
134 let p = self.make_polynomial(e)?;
135 let p = self.merge_equal_hidden_components(p);
136 self.degree = 1;
137 let p = self.make_polynomial(p)?;
138 let p = self.solve_polynomial(p)?;
139
140 let e = if self.non_linear_components.len() != 0 {
142 let e = self.poly_to_nonpoly(p);
143 debug!("After adding non linear components, found: {}", e);
144 e.reduce_masked(self.mask)
145 } else {
146 p
147 };
148
149 Ok(if self.options.patterns {
159 crate::patterns::apply_patterns(e, self.mask)
160 } else {
161 e
162 })
163 }
164
165 fn calc_signature(&self, e: &Expr, t: usize) -> Vec<u64> {
167 e.truth_table_masked(t, self.mask)
168 }
169
170 fn make_conjunction_sum(&self, mut signature: Vec<u64>, t: usize) -> Expr {
172 let mut terms: Vec<Expr> = vec![];
173
174 let constant = signature[0];
176
177 if constant != 0 {
178 terms.push(Expr::Const(constant));
179
180 for v in &mut signature {
181 *v = v.wrapping_sub(constant);
182 }
183 }
184
185 let mut sublist = Vec::with_capacity(t);
186 for index in 1..(1usize << t) {
187 let coeff = signature[index] & self.mask;
188
189 if coeff == 0 {
190 continue;
191 }
192
193 sublist.clear();
194 for i in 0..t {
195 if ((index >> i) & 1) == 1 {
196 sublist.push(i);
197 }
198 }
199 let conjunction = Expr::And(
200 sublist
201 .iter()
202 .copied()
203 .map(|v| Expr::Var(v.into()))
204 .collect(),
205 );
206
207 terms.push(match coeff {
208 1 => conjunction,
209 c => c * conjunction,
210 });
211
212 sub_coeff(&mut signature, coeff, index, &sublist);
213 }
214
215 match terms.len() {
216 0 => Expr::zero(),
217 1 => terms.remove(0),
218 _ => Expr::Add(terms),
219 }
220 }
221
222 fn solve_linear_inner(&self, e: Expr, t: usize, from_poly: bool) -> Expr {
224 let signature = self.calc_signature(&e, t);
225
226 if from_poly {
227 return self.make_conjunction_sum(signature, t);
229 }
230
231 self.make_conjunction_sum(signature, t)
233 }
234
235 fn solve_linear(&mut self, e: Expr, from_poly: bool) -> Result<Expr, SolveError> {
237 let mut var_map = BiMap::<VarId, VarId>::new();
238 let mut t = 0;
239
240 debug!("Solving linear MBA: {}", e);
241
242 let e = reduce_vars(e, &mut var_map, &mut t);
244 debug!("Reduced number of variables to equivalent problem: {}", e);
245
246 let e = if let Some(simplified) = self.l_cache.get(&e) {
247 debug!("Found linear MBA in cache");
248 simplified
249 } else {
250 debug!("Solving linear MBA");
251
252 if t > MAX_VARS {
253 return Err(SolveError::TooManyVariables {
254 found: t,
255 max: MAX_VARS,
256 });
257 }
258
259 let simplified = self.solve_linear_inner(e.clone(), t, from_poly);
260 self.l_cache.insert(e, simplified.clone());
261 simplified
262 };
263
264 let e = restore_vars(e, &var_map)?;
265
266 debug!("Found solution to linear MBA: {}", e);
267
268 Ok(e)
269 }
270
271 fn poly_to_linear(&self, e: Expr, deg: usize) -> Expr {
273 match e {
274 Expr::Var(v) => Expr::Var(((deg - 1) * self.t + v.0).into()),
275
276 Expr::Mul(terms) => {
277 let s = if (self.degree - terms.len()) & 1 == 0 {
279 1
280 } else {
281 u64::MAX
282 };
283
284 s * Expr::And(
285 terms
286 .into_iter()
287 .enumerate()
288 .map(|(i, e)| self.poly_to_linear(e, deg + i + 1))
289 .collect(),
290 )
291 }
292
293 Expr::Const(c) => {
294 if deg != 0 {
296 e
297 } else {
298 let s: u64 = if self.degree & 1 == 0 { u64::MAX } else { 1 };
300 Expr::Const(s.wrapping_mul(c))
301 }
302 }
303
304 _ => e.map(|e| self.poly_to_linear(e, deg)),
305 }
306 }
307
308 fn linear_to_poly(&self, e: Expr) -> Result<Expr, SolveError> {
310 match &e {
311 Expr::And(terms) => {
312 let mut s = 1u64;
314
315 let mut grouped: Vec<Vec<usize>> = vec![vec![]; self.degree];
316
317 for t in terms {
318 if let Expr::Var(v) = t {
319 let d = v.0 / self.t;
320 grouped[d].push(v.0 % self.t);
321 } else {
322 return Err(SolveError::UnrecognizedForm(e.clone()));
323 }
324 }
325
326 let mut terms: Vec<Expr> = vec![];
327
328 for g in grouped {
329 if g.is_empty() {
330 s = s.wrapping_mul(u64::MAX);
331 continue;
332 }
333
334 terms.push(Expr::And(
335 g.into_iter().map(|v| Expr::Var(v.into())).collect(),
336 ));
337 }
338
339 Ok(s * Expr::Mul(terms))
340 }
341
342 Expr::Add(_) | Expr::Scale(_, _) => e.try_map(|e| self.linear_to_poly(e)),
343
344 Expr::Var(v) => {
345 let s = if self.degree & 1 == 0 { u64::MAX } else { 1 };
347 Ok(s * Expr::Var((v.0 % self.t).into()))
348 }
349
350 Expr::Const(_) => {
351 let s = if self.degree & 1 == 0 { u64::MAX } else { 1 };
353 Ok(s * e)
354 }
355
356 _ => Err(SolveError::UnrecognizedForm(e.clone())),
357 }
358 }
359
360 fn solve_polynomial(&mut self, e: Expr) -> Result<Expr, SolveError> {
362 debug!("Solving polynomial MBA: {}", e);
363
364 if self.degree == 1 {
366 debug!("This is a linear MBA");
367 return self.solve_linear(e, false);
368 }
369
370 let e = self.poly_to_linear(e, 0);
371 let e: Expr = self.solve_linear(e, true)?;
372 let e = self.linear_to_poly(e)?.reduce_masked(self.mask);
373
374 debug!("Found polynomial solution: {}", e);
375
376 Ok(e)
377 }
378
379 fn hide_in_var(&mut self, e: Expr, mask: u64) -> Result<Expr, SolveError> {
381 debug!("e={} is not linear and will be replaced by a variable", e);
382
383 let e = match e {
384 Expr::Const(_) => e,
385 _ => simplify_mba_inner(self.l_cache, e, mask.count_ones() as u8, self.options)?
386 .reduce_masked(mask),
387 };
388
389 let note = (-e.clone() - Expr::make_const(1)).reduce_masked(mask);
390 debug!("!e would be {}", note);
391
392 if let Some(v) = self.non_linear_components.get_by_right(&e) {
393 debug!("Found variable v{} for e", v);
394 Ok(Expr::Var(*v))
395 } else if let Some(v) = self.non_linear_components.get_by_right(¬e) {
396 debug!("Found variable v{} for !e", v);
397 Ok(!Expr::Var(*v))
398 } else {
399 let v = self.t.into();
400 debug!("Creating variable v{} for e={}", v, e);
401 self.t += 1;
402 self.non_linear_components.insert(v, e);
403 Ok(Expr::Var(v))
404 }
405 }
406
407 fn is_signature_bitwise(&self, s: &Vec<u64>, mask: u64) -> bool {
408 let minus_one = mask;
409 let minus_two = mask - 1;
410
411 if s[0] == 0 {
412 if s.iter().all(|&x| x == 0 || x == 1) {
413 debug!("Signature is {:?} in [0, 1].", s);
414 true
415 } else {
416 false
417 }
418 } else if s[0] == minus_one {
419 if s.iter().all(|&x| x == minus_one || x == minus_two) {
420 debug!("Signature is {:?} in [-1, 2]", s);
421 true
422 } else {
423 false
424 }
425 } else {
426 false
427 }
428 }
429
430 fn variable_substitution(&self, e: Expr) -> Option<Expr> {
432 let mut sub_vars = vec![];
433
434 for v in e.get_vars() {
435 if let Some(definition) = self.non_linear_components.get_by_left(&v)
436 && self.is_linear(definition)
437 {
438 sub_vars.push((v, definition));
439 }
440 }
441 sub_vars.sort_unstable_by_key(|(variable, _)| variable.0);
442
443 if sub_vars.is_empty() || sub_vars.len() > 2 {
444 return None;
445 }
446
447 debug!("While checking if {} is linear", e);
448 debug!("Proceding with advanced variable substitution");
449 let mut zero_expressions = Vec::with_capacity(sub_vars.len());
450 for (sub_var, definition) in sub_vars {
451 debug!("Found substitution v{} = {}", sub_var, definition);
452 let zero = Expr::Var(sub_var) - definition.clone();
453 debug!("Using zero expression {}", zero);
454 zero_expressions.push(zero);
455 }
456
457 let mut var_map = BiMap::new();
458 let mut t = 0;
459
460 let reduced_e = reduce_vars(e.clone(), &mut var_map, &mut t);
461 let reduced_zeros: Vec<_> = zero_expressions
462 .iter()
463 .cloned()
464 .map(|zero| reduce_vars(zero, &mut var_map, &mut t))
465 .collect();
466 if t > 10 {
467 return None;
468 }
469
470 let se = self.calc_signature(&reduced_e, t);
471 debug!("Using signature {:?}", se);
472 let zero_signatures: Vec<_> = reduced_zeros
473 .iter()
474 .map(|zero| self.calc_signature(zero, t))
475 .collect();
476
477 if zero_expressions.len() == 1 {
478 let see = &zero_signatures[0];
479 debug!("Using zero signature {:?}", see);
480 if let Some(lambda) = find_lambda_int(&se, see, 0, 1, self.n) {
481 debug!("Found lambda that creates a [0, 1] signature: {:?}", lambda);
482 return Some(e - lambda * zero_expressions.remove(0));
483 }
484 if let Some(lambda) = find_lambda_int(&se, see, -1, -2, self.n) {
485 debug!(
486 "Found lambda that creates a [-1, -2] signature: {:?}",
487 lambda
488 );
489 return Some(e - lambda * zero_expressions.remove(0));
490 }
491 return None;
492 }
493
494 for (a, b) in [(0, 1), (-1, -2)] {
495 if let Some((left, right)) =
496 find_two_lambdas_int(&se, &zero_signatures[0], &zero_signatures[1], a, b, self.n)
497 {
498 debug!("Found two substitution lambdas: {}, {}", left, right);
499 return Some(
500 e - left * zero_expressions[0].clone() - right * zero_expressions[1].clone(),
501 );
502 }
503 }
504 None
505 }
506
507 fn is_linear_bitwise(&self, l: Expr, mask: u64) -> Option<Expr> {
509 let mut t = 0;
510 let mut var_map = BiMap::new();
511 let e = reduce_vars(l.clone(), &mut var_map, &mut t);
512
513 if t > 10 {
514 return None;
516 }
517
518 let s = e.truth_table_masked(t, mask);
519
520 if self.is_signature_bitwise(&s, mask) {
521 debug!("Will treat {} as a bitwise expression", e);
522 Some(l)
523 } else {
524 self.variable_substitution(l)
526 }
527 }
528
529 fn make_bitwise(&mut self, e: Expr, mut mask: u64) -> Result<Expr, SolveError> {
531 match e {
532 Expr::Const(c) => {
534 if (c & mask) == 0 || (c & mask) == self.mask {
535 Ok(e)
536 } else {
537 self.hide_in_var(e, mask)
538 }
539 }
540
541 Expr::Var(_) => Ok(e),
543
544 Expr::Not(_) | Expr::Or(_) | Expr::Xor(_) => {
545 e.try_map(|e| self.make_bitwise(e, mask))
550 }
551
552 Expr::And(terms) => {
554 for e in &terms {
555 if let Expr::Const(c) = e {
556 let c = c & mask;
557 if c & (c.wrapping_add(1)) == 0 {
558 debug!("Found dynamic mask {} = 2^{} -1", c, c.count_ones());
559 mask = c;
560 }
561 }
562 }
563
564 if mask == 0 {
565 return Ok(Expr::zero());
566 }
567
568 Ok(Expr::And(
569 terms
570 .into_iter()
571 .map(|e| self.make_bitwise(e, mask))
572 .collect::<Result<Vec<_>, _>>()?,
573 ))
574 }
575
576 Expr::Add(_) | Expr::Scale(_, _) => {
578 let previous = e.clone();
579 let l = e.try_map(|e| self.make_linear(e, mask))?;
580 if let Some(l) = self.is_linear_bitwise(l, mask) {
581 Ok(l)
582 } else {
583 self.hide_in_var(previous, mask)
584 }
585 }
586
587 _ => self.hide_in_var(e, mask),
588 }
589 }
590
591 fn make_product(&mut self, e: Expr) -> Result<Expr, SolveError> {
593 match &e {
594 Expr::Mul(terms) => {
595 self.degree = max(self.degree, terms.len());
596 e.try_map(|e| self.make_bitwise(e, self.mask))
597 }
598
599 _ => Ok(Expr::Mul(vec![self.make_bitwise(e, self.mask)?])),
601 }
602 }
603
604 fn make_scaled_product(&mut self, e: Expr) -> Result<Expr, SolveError> {
606 match e {
607 Expr::Const(_) => Ok(e),
608 Expr::Scale(_, _) => e.try_map(|e| self.make_product(e)),
609 _ => self.make_product(e),
610 }
611 }
612
613 fn make_polynomial(&mut self, e: Expr) -> Result<Expr, SolveError> {
615 match e {
616 Expr::Add(_) => e.try_map(|e| self.make_scaled_product(e)),
617 _ => self.make_scaled_product(e),
618 }
619 }
620
621 fn make_scaled_bitwise(&mut self, e: Expr, mask: u64) -> Result<Expr, SolveError> {
623 match e {
624 Expr::Const(_) => Ok(e),
625 Expr::Scale(_, _) => e.try_map(|e| self.make_bitwise(e, mask)),
626 _ => self.make_bitwise(e, mask),
627 }
628 }
629
630 fn make_linear(&mut self, e: Expr, mask: u64) -> Result<Expr, SolveError> {
632 match e {
633 Expr::Add(_) => e.try_map(|e| self.make_scaled_bitwise(e, mask)),
634 _ => self.make_scaled_bitwise(e, mask),
635 }
636 }
637
638 fn is_linear(&self, e: &Expr) -> bool {
639 fn is_bitwise(e: &Expr, mask: u64) -> bool {
640 match e {
641 Expr::Const(c) => ((c & mask) == 0) || ((c & mask) == mask),
643
644 Expr::Var(_) => true,
646
647 Expr::Not(e) => is_bitwise(e, mask),
649
650 Expr::And(es) | Expr::Or(es) | Expr::Xor(es) => {
651 es.iter().all(|e| is_bitwise(e, mask))
652 }
653
654 _ => false,
655 }
656 }
657
658 fn is_scaled_bitwise(e: &Expr, mask: u64) -> bool {
659 match e {
660 Expr::Const(_) => true,
661 Expr::Scale(_, e) => is_bitwise(e, mask),
662 _ => is_bitwise(e, mask),
663 }
664 }
665
666 match e {
667 Expr::Add(terms) => terms.iter().all(|e| is_scaled_bitwise(e, self.mask)),
668 _ => is_scaled_bitwise(e, self.mask),
669 }
670 }
671}
672
673fn simplify_mba_inner<C: LinearCache>(
674 l_cache: &C,
675 e: Expr,
676 n: u8,
677 options: SimplifyOptions,
678) -> Result<Expr, SolveError> {
679 let mut solver = MBASolver::new(l_cache, &e, n, options);
680 solver.solve(e)
681}
682
683fn simplify_to_fixed_point<F>(mut e: Expr, mut simplify: F) -> Result<Expr, SolveError>
684where
685 F: FnMut(Expr) -> Result<Expr, SolveError>,
686{
687 let mut size = usize::MAX;
688
689 for _ in 0..MAX_SIMPLIFICATION_PASSES {
690 let next = simplify(e.clone())?;
691 debug!("e: {}", next);
692
693 let next_size = next.size();
694 let settled = next == e || next_size >= size;
701 e = next;
702 if settled {
703 return Ok(e);
704 }
705 size = next_size;
706 }
707
708 debug!("simplification pass limit reached");
709 Ok(e)
710}
711
712#[derive(Debug, Default)]
720pub struct SimplifyCache(MbaCache);
721
722impl SimplifyCache {
723 pub fn new() -> Self {
725 Self::default()
726 }
727}
728
729#[derive(Debug, Clone, Copy, PartialEq, Eq)]
734pub struct SimplifyOptions {
735 pub patterns: bool,
739}
740
741impl Default for SimplifyOptions {
742 fn default() -> Self {
743 Self { patterns: true }
744 }
745}
746
747pub fn simplify_mba(e: Expr, n: u8) -> Result<Expr, SolveError> {
749 simplify_mba_with(e, n, SimplifyOptions::default())
750}
751
752pub fn simplify_mba_with(e: Expr, n: u8, options: SimplifyOptions) -> Result<Expr, SolveError> {
754 simplify_mba_with_cache(&LocalCache::new(), e, n, options)
755}
756
757pub fn simplify_mba_cached(e: Expr, n: u8, cache: &SimplifyCache) -> Result<Expr, SolveError> {
760 simplify_mba_with_cache(&cache.0, e, n, SimplifyOptions::default())
761}
762
763fn simplify_mba_with_cache<C: LinearCache>(
764 cache: &C,
765 e: Expr,
766 n: u8,
767 options: SimplifyOptions,
768) -> Result<Expr, SolveError> {
769 let mask = make_mask(n);
770 let e = e.reduce_masked(mask);
771 let e = simplify_to_fixed_point(e, |e| simplify_mba_inner(cache, e, n, options))?;
772
773 Ok(prettify(e, n))
776}
777
778#[cfg(test)]
779mod tests {
780 use super::*;
781 use std::cell::Cell;
782
783 #[test]
784 fn disabling_patterns_skips_pattern_only_simplifications() {
785 let e = (Expr::Var(0.into()) & (-Expr::Var(0.into())) & (2u64 * Expr::Var(0.into())))
788 .reduce(64);
789
790 assert_eq!(simplify_mba(e.clone(), 64), Ok(Expr::zero()));
791
792 let without_patterns = SimplifyOptions { patterns: false };
793 assert_ne!(simplify_mba_with(e, 64, without_patterns), Ok(Expr::zero()),);
794 }
795
796 #[test]
797 fn inverse_pct_round_trip_preserves_variable_index() {
798 let cache = LocalCache::new();
799 let mut solver =
800 MBASolver::new(&cache, &Expr::Var(2.into()), 8, SimplifyOptions::default());
801 solver.degree = 2;
802 let original = Expr::Var(2.into());
803 let encoded = solver.poly_to_linear(original.clone(), 2);
804
805 assert_eq!(encoded, Expr::Var(5.into()));
806 assert_eq!(solver.linear_to_poly(encoded), Ok(u64::MAX * original));
807 }
808
809 #[test]
817 fn fixed_point_stops_at_the_first_pass_that_does_not_shrink() {
818 let result = simplify_to_fixed_point(Expr::Var(0.into()), |e| {
819 Ok(match e {
820 Expr::Var(VarId(0)) => !Expr::Var(1.into()),
821 Expr::Not(inner) if *inner == Expr::Var(1.into()) => 2u64 * Expr::Var(2.into()),
822 Expr::Scale(c, inner) if c == 2 && *inner == Expr::Var(2.into()) => {
823 Expr::Var(3.into())
824 }
825 stable => stable,
826 })
827 });
828
829 assert_eq!(result, Ok(2u64 * Expr::Var(2.into())));
830 }
831
832 #[test]
836 fn fixed_point_cycle_stops_at_the_first_non_shrinking_pass() {
837 let start = !Expr::Var(0.into());
838 let result = simplify_to_fixed_point(start.clone(), |e| {
839 if e == start {
840 Ok(Expr::Var(1.into()))
841 } else {
842 Ok(start.clone())
843 }
844 });
845
846 assert_eq!(result, Ok(start));
847 }
848
849 #[test]
850 fn fixed_point_has_a_pass_limit() {
851 let terms: Vec<Expr> = (0..20).map(|i| Expr::Var(i.into())).collect();
853 let calls = Cell::new(0usize);
854 let result = simplify_to_fixed_point(Expr::Add(terms), |e| {
855 calls.set(calls.get() + 1);
856 Ok(match e {
857 Expr::Add(mut terms) if terms.len() > 1 => {
858 terms.pop();
859 Expr::Add(terms)
860 }
861 other => other,
862 })
863 });
864
865 assert_eq!(calls.get(), MAX_SIMPLIFICATION_PASSES);
866 let Ok(Expr::Add(remaining)) = result else {
867 panic!("expected a sum, got {result:?}");
868 };
869 assert_eq!(remaining.len(), 20 - MAX_SIMPLIFICATION_PASSES);
870 }
871}