Skip to main content

rumba_core/
simplify.rs

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
20/// The largest number of variables a linear MBA may carry into the truth-table
21/// solve. The signature is `2^t` wide, so this bounds one solve at 1M entries.
22pub(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
43/// Reduces the number of variables present in the MBA
44fn 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
62/// Restores the varialbes in the mba
63fn 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    /// The number of bits being considered
79    n: u8,
80
81    /// The mask that corresponds to the given bitsize
82    mask: u64,
83
84    /// The map between non linear components and variables
85    non_linear_components: BiMap<VarId, Expr>,
86
87    /// The number of variables in the expression
88    t: usize,
89
90    /// The degree of the polynomial expression
91    degree: usize,
92
93    /// A cache for simplifying linear MBAs
94    l_cache: &'a C,
95
96    /// The options controlling this run.
97    options: SimplifyOptions,
98}
99
100impl<'a, C: LinearCache> MBASolver<'a, C> {
101    /// Create a new Solver
102    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    /// Replaces non polynomial variables by their hidden expressions
115    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    /// Solves a non polynomial MBA
130    fn solve(&mut self, e: Expr) -> Result<Expr, SolveError> {
131        // TODO: Remove this only needs to be done once
132        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        // This was a non linear MBA
141        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        // Structural pattern rewrites, applied to the *solved* expression
150        // rather than the raw input. The residues these patterns target are
151        // what the solver leaves behind, and by this point the expression is
152        // a fraction of the size it arrived as — the matchers walk every node
153        // of it, so running them on the obfuscated input is both slower and
154        // less likely to match. A hit changes the expression, so
155        // `simplify_to_fixed_point` runs another pass and the downstream
156        // stages still get to exploit the collapse. Skipped when the caller
157        // disables the pattern engine via [`SimplifyOptions`].
158        Ok(if self.options.patterns {
159            crate::patterns::apply_patterns(e, self.mask)
160        } else {
161            e
162        })
163    }
164
165    /// Calcluates the signature of a linear MBA
166    fn calc_signature(&self, e: &Expr, t: usize) -> Vec<u64> {
167        e.truth_table_masked(t, self.mask)
168    }
169
170    /// Creates a conjuction sum for the given signature
171    fn make_conjunction_sum(&self, mut signature: Vec<u64>, t: usize) -> Expr {
172        let mut terms: Vec<Expr> = vec![];
173
174        // The constant term
175        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    /// Solves a linear MBA
223    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            // We necessarily want a sum of conjunctions
228            return self.make_conjunction_sum(signature, t);
229        }
230
231        // TODO: add a refined solution to identify xor etc
232        self.make_conjunction_sum(signature, t)
233    }
234
235    /// Simplifies a linear MBA
236    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        // Reduce the number of variables in the expression
243        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    /// Turns a polynomial MBA to a linear one using PCT
272    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                // The sign correction -> see paper
278                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                // This shouldn't be done in multiplications, only constants in the addition
295                if deg != 0 {
296                    e
297                } else {
298                    // The sign correction -> see paper
299                    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    /// Turns a linear MBA into a polynomial one using the inverse PCT
309    fn linear_to_poly(&self, e: Expr) -> Result<Expr, SolveError> {
310        match &e {
311            Expr::And(terms) => {
312                // The sign correction -> see paper
313                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                // The sign correction -> see paper
346                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                // The sign correction -> see paper
352                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    /// Solves a polynomial MBA
361    fn solve_polynomial(&mut self, e: Expr) -> Result<Expr, SolveError> {
362        debug!("Solving polynomial MBA: {}", e);
363
364        // This is a linear MBA
365        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    /// Hides a non linear element behind a variable
380    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(&note) {
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    // Read paper
431    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    // A Linear MBA might "hide" a bitwise expression
508    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            // This would be too expensive
515            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            // Attempt to "fix" the signature with a variable substitution
525            self.variable_substitution(l)
526        }
527    }
528
529    /// Turns an expression into a bitwise expression
530    fn make_bitwise(&mut self, e: Expr, mut mask: u64) -> Result<Expr, SolveError> {
531        match e {
532            // -1 and 0 are bitwise
533            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            // Variables are bitwise
542            Expr::Var(_) => Ok(e),
543
544            Expr::Not(_) | Expr::Or(_) | Expr::Xor(_) => {
545                // if only bitwise
546                // self
547                // if has negatives, try and fix the "biphased" problem
548                // worst case
549                e.try_map(|e| self.make_bitwise(e, mask))
550            }
551
552            // Dynamic masking: if we and with a constant that constant will be are new mask
553            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            // A Linear MBA might "hide" a bitwise expression
577            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    /// Turns an expression into a bitwise product
592    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            // This makes life better on much easier
600            _ => Ok(Expr::Mul(vec![self.make_bitwise(e, self.mask)?])),
601        }
602    }
603
604    /// Turns an expression into a scaled bitwise product
605    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    /// Turns an expression into a polynomial expression
614    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    /// Turns an expression into a scaled bitwise expression
622    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    /// Turns an expression into a linear expression
631    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                // -1 and 0 are bitwise
642                Expr::Const(c) => ((c & mask) == 0) || ((c & mask) == mask),
643
644                // Variables are bitwise
645                Expr::Var(_) => true,
646
647                // A boolean expression is bitwise if all its sub expressions are bitwise
648                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        // Always keep the pass just made. Size is not a proxy for quality here:
695        // a pass that grows the expression is usually the canonicalization the
696        // ground truth is compared against, and discarding it collapses the OK
697        // rate (loki_tiny 24997 -> 13522). The trade is that a cycle is
698        // indistinguishable from such a pass, so the result may be larger than
699        // an intermediate; the corpus shows no case where that costs anything.
700        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/// A reusable memo of solved linear MBAs.
713///
714/// Simplifying an expression solves many linear sub-MBAs; a `SimplifyCache`
715/// remembers those solutions so that a caller simplifying many expressions — or
716/// the same ones repeatedly across rounds of an analysis — pays for each distinct
717/// linear solve once. Create one with [`SimplifyCache::new`] and hand it to
718/// [`simplify_mba_cached`]. It is safe to share one across threads.
719#[derive(Debug, Default)]
720pub struct SimplifyCache(MbaCache);
721
722impl SimplifyCache {
723    /// An empty cache.
724    pub fn new() -> Self {
725        Self::default()
726    }
727}
728
729/// Options controlling how an expression is simplified.
730///
731/// Use [`SimplifyOptions::default`] for the standard behaviour and pass a
732/// customised value to [`simplify_mba_with`].
733#[derive(Debug, Clone, Copy, PartialEq, Eq)]
734pub struct SimplifyOptions {
735    /// Whether to run the structural pattern-rewrite engine before lifting the
736    /// expression to a polynomial. Enabled by default; disabling it is mainly
737    /// useful for measuring the engine's contribution.
738    pub patterns: bool,
739}
740
741impl Default for SimplifyOptions {
742    fn default() -> Self {
743        Self { patterns: true }
744    }
745}
746
747/// Simplifies a Mixed Boolean-Arithmetic expression on `n` bits.
748pub fn simplify_mba(e: Expr, n: u8) -> Result<Expr, SolveError> {
749    simplify_mba_with(e, n, SimplifyOptions::default())
750}
751
752/// [`simplify_mba`] with explicit [`SimplifyOptions`].
753pub 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
757/// [`simplify_mba`] against a caller-owned [`SimplifyCache`], so linear solves
758/// are reused across calls.
759pub 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    // The only place prettify may run: on the way out, after the fixed point has
774    // settled. See the module docs for why it must stay out of the loop.
775    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        // `X & -X & 2·X` collapses to 0 only through the structural pattern
786        // engine; the polynomial machinery alone leaves it untouched.
787        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    /// The loop stops at the first pass that fails to shrink the expression,
810    /// even when continuing would eventually reach something smaller.
811    ///
812    /// `Var(0)` -> `!Var(1)` -> `2·Var(2)` -> `Var(3)` plateaus in size at the
813    /// second step before shrinking again; the chase is abandoned there and the
814    /// plateau result is kept. Following such plateaus measured 16-61% across
815    /// the corpus while leaving every OK/OKZ/NG count unchanged.
816    #[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    /// A two-cycle ends at the first pass that fails to shrink, and that pass
833    /// is kept even though an intermediate was smaller — see the note in
834    /// [`simplify_to_fixed_point`] on why the last pass always wins.
835    #[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        // Shrink on every pass, so the cap is what ends the run.
852        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}