Skip to main content

pounce_nl/
nl_quadratic.rs

1//! Degree-≤2 recognition over an [`Expr`] DAG.
2//!
3//! This is the classifier's "is this row a quadratic, and which one?"
4//! question, answered once and reused. `pounce-cli`'s `dispatch` owns the
5//! *routing* decision (which `ProblemClass`, which solver); what lives here
6//! is only the algebra, so the consumers that are not the CLI —
7//! `NlTnlp`'s constant-structure evaluation, the parse-time recognizer, the
8//! `QcqpProblem` extractor — can reach it without depending on the
9//! command-line driver. It moved out of `pounce-cli/src/dispatch.rs` in
10//! Q3 of the #588 series; see
11//! `dev-notes/quadratic-structure-exploitation.md`.
12//!
13//! ## Two properties this module is built around
14//!
15//! **It is iterative.** The walk carries its own work stack rather than
16//! recursing, because the trees it is handed are not shallow: a `.nl`
17//! writer that emits `o0` (binary `+`) chains for a long sum — Pyomo does —
18//! produces a left-deep `Add` tree one level per term. The recursive
19//! predecessor aborted the process somewhere between 4 000 and 6 000 terms
20//! on a 2 MB thread (which is what a test gets, and where the crash was
21//! first reproduced) and between 16 000 and 24 000 on the CLI's 8 MB main
22//! thread. A stack overflow is an abort, not an error return, so the depth a
23//! *recognizer* survives must not depend on which thread called it.
24//!
25//! It is worth knowing what this does **not** fix: `nl_reader`'s parser
26//! recurses too, with a fatter frame — it gives out at ~6 000 on that same
27//! 8 MB thread — so a deep `.nl` file still fails to load, and it fails
28//! before reaching this module. What is fixed is every path where the tree
29//! is already built (`NlProblem::from_expressions`, a model handed across
30//! threads) and the ceiling this module used to impose on the parser's
31//! successor.
32//!
33//! **It never allocates per monomial.** The predecessor keyed monomials on
34//! `BTreeMap<Vec<usize>, f64>`, so every term cost a heap allocation and
35//! every merge cloned one (`entry(m.clone())`). A degree-≤2 form has only
36//! three shapes of term, so [`Quad2`] stores them in three fields with
37//! inline keys and the allocation disappears. The same change removes the
38//! `O(N²)` accumulation on `Add` chains: the old `add` re-scanned the whole
39//! accumulated map for zeros on *every* merge, which is quadratic down a
40//! left-deep chain. Zeros can only appear where a merge touched, so that is
41//! all this one looks at.
42
43use crate::nl_reader::{BinOp, Expr, UnaryOp};
44use std::collections::{BTreeMap, BTreeSet};
45
46/// The symmetric Hessian of a quadratic form, stored as a sparse upper-
47/// triangular (i ≤ j) map of `(i, j) -> ∂²/∂xᵢ∂xⱼ`. Empty means the
48/// expression is (at most) linear.
49pub type QuadHessian = BTreeMap<(usize, usize), f64>;
50
51/// Full quadratic read-out: `(Hessian, [(var, linear coef), …], constant)`.
52/// The linear and constant parts are the pieces AMPL/Pyomo fold into the
53/// nonlinear objective tree (see [`analyze_quadratic_full`]).
54pub type QuadForm = (QuadHessian, Vec<(usize, f64)>, f64);
55
56/// A polynomial of total degree ≤ 2 in its own shape: a constant, the
57/// linear coefficients keyed by variable, and the quadratic coefficients
58/// keyed by the (i ≤ j) variable pair.
59///
60/// This replaces a general `BTreeMap<Vec<usize>, f64>` polynomial. Degree
61/// is a property of the *type* here rather than of the data, so the
62/// "is it still quadratic?" test is three `is_empty()` calls instead of a
63/// scan, and no monomial key is ever allocated or cloned.
64///
65/// ### Zero coefficients
66///
67/// Stored coefficients are nonzero: [`add`](Quad2::add) and
68/// [`mul`](Quad2::mul) drop any entry they leave at exactly zero, which is
69/// what makes [`degree`](Quad2::degree) and [`as_constant`](Quad2::as_constant)
70/// answerable in `O(1)`. `constant` is the one exception and needs none:
71/// both `0.0` and `-0.0` in that field mean "no constant term", every
72/// consumer guards on `!= 0.0`, and [`analyze_quadratic_full`] normalizes
73/// the sign on the way out.
74///
75/// ### …and when dropping one loses something
76///
77/// Dropping is right for the *storage* question and can be wrong for the
78/// *degree* question, and [`lost_terms`](Quad2::lost_terms) is the
79/// difference (gh #683, sharpened by gh #687). A coefficient that reaches
80/// zero is a coefficient that was **summed**, and it is the arithmetic of
81/// that sum — not the drop — that decides whether anything went missing:
82///
83/// * `x − x` folds `fl(1) + fl(−1)` to `0.0`, and that add is **exact**.
84///   The term really is absent, the body really is degree 0, and the maps
85///   are not a lower bound on anything.
86/// * `2⁵³·x² + x² − 2⁵³·x²` loses the `x²` at `fl(2⁵³ + 1) = 2⁵³`, which is
87///   an **inexact** add, and only then does `2⁵³ − 2⁵³` drop the survivor.
88///   The body is `x²`; the maps say degree 0.
89///
90/// The loss happens at the inexact fold; the drop is only where it becomes
91/// visible. So the flag is set from the *fold*: a form whose arithmetic
92/// never rounded — and never flushed a live coefficient to zero — carries
93/// coefficients that are exactly what real arithmetic on the writer's own
94/// literals would give, and a term missing from its maps is a term that is
95/// genuinely absent. Flagging the drop instead refused both of the above
96/// alike, which was sound and cost reach on the first (gh #687).
97#[derive(Debug, Clone, Default, PartialEq)]
98pub struct Quad2 {
99    constant: f64,
100    linear: BTreeMap<usize, f64>,
101    quadratic: QuadHessian,
102    /// Set when a term may be **missing** from the maps — see the type
103    /// docs and [`lost_terms`](Quad2::lost_terms). Sticky: it survives
104    /// every operation the form takes part in, because a term that went
105    /// missing three additions ago is no less missing now.
106    lost_terms: bool,
107    /// Set when some arithmetic in this form's construction **rounded**:
108    /// an add that could not represent its sum, or a product that could
109    /// not represent itself. Sticky, and for the same reason.
110    ///
111    /// Clear ⇒ every stored coefficient is exactly the real-arithmetic
112    /// value of the terms folded into it, so a coefficient that reached
113    /// zero reached it because the real value is zero. That is the whole
114    /// warrant for `lost_terms` staying clear across a cancellation, and
115    /// it is why an inexact *fold* sets this bit even when the coefficient
116    /// it produced is nonzero and stays stored: what it rounded away can
117    /// be cancelled out of sight later, by an add that is itself exact.
118    inexact: bool,
119}
120
121impl Quad2 {
122    /// The degree-0 term.
123    pub fn constant(&self) -> f64 {
124        self.constant
125    }
126
127    /// The degree-1 terms, ascending by variable index.
128    pub fn linear(&self) -> &BTreeMap<usize, f64> {
129        &self.linear
130    }
131
132    /// The degree-2 terms as *polynomial* coefficients keyed `(i ≤ j)` —
133    /// the coefficient of `xᵢxⱼ`, **not** the Hessian entry (they differ by
134    /// a factor of 2 on the diagonal; [`analyze_quadratic_full`] applies it).
135    pub fn quadratic(&self) -> &QuadHessian {
136        &self.quadratic
137    }
138
139    /// Whether a term may be **missing** from this form's maps: something
140    /// was dropped on the way here *and* the arithmetic that produced it
141    /// had already rounded, or a live coefficient was flushed to zero (or
142    /// to `NaN`) outright.
143    ///
144    /// This is what turns the term maps from an answer about degree into a
145    /// lower bound on it (gh #683). When it is set, an empty
146    /// [`quadratic`](Quad2::quadratic) map is no longer evidence that the
147    /// body is affine, and a consumer asking a *degree* question — as
148    /// opposed to a storage or evaluation question — must say "not
149    /// established" instead of "affine". The two questions used to share
150    /// one predicate, which is how a genuinely quadratic row came to be
151    /// reported as proved affine and had its Jacobian frozen for a whole
152    /// solve.
153    ///
154    /// It was originally set by the **drop** alone, which conflated an
155    /// exact cancellation with a lossy one and refused both (gh #687):
156    /// `x − x` is degree 0 by an exact add, and a consumer that treats it
157    /// as "not established" gives up a proved degree, a matrix evaluation
158    /// and — once the classifier gates on this too — the convex route, for
159    /// a body that lost nothing. The gate is now the fold, so the exact
160    /// case keeps its fast paths and the lossy one is refused for the
161    /// reason it deserves. [`inexact`](Quad2::inexact) is the other half of
162    /// that answer.
163    ///
164    /// It is deliberately one bit for the whole form rather than a record
165    /// per monomial: the consumer's question is about the form, and Q3's
166    /// point was that a `Quad2` allocates nothing per term (gh #588, Q3).
167    /// Two bits, now, and still nothing per term — the cost of a per-key
168    /// provenance record is what keeps this a *conservative* answer: an
169    /// inexact fold on one monomial makes a cancellation on an unrelated
170    /// one count as lossy.
171    ///
172    /// A *linear* or constant term that goes missing sets it too, even
173    /// though losing one cannot understate an affine body's degree by
174    /// itself: it can once the form is **multiplied**, because
175    /// [`mul`](Quad2::mul)'s degree guard reads the same maps —
176    /// `(2⁵³ + 1 − 2⁵³)·x · x²` is degree 3 and folds to an apparent
177    /// degree 0. Distinguishing the two would take a third flag to buy
178    /// back reach that the corpus does not contain.
179    pub fn lost_terms(&self) -> bool {
180        self.lost_terms
181    }
182
183    /// Whether any arithmetic behind this form's coefficients **rounded**.
184    ///
185    /// Clear means every stored coefficient is exact, which is what makes a
186    /// coefficient of zero a proof of absence rather than a lower bound —
187    /// see [`lost_terms`](Quad2::lost_terms), which is this bit read at the
188    /// moment a term is dropped. Public because it is the thing a test
189    /// about the sharpened gate has to be able to see; no consumer routes
190    /// on it.
191    pub fn inexact(&self) -> bool {
192        self.inexact
193    }
194
195    pub(crate) fn of_constant(c: f64) -> Self {
196        Quad2 {
197            // `-0.0` and `0.0` alike mean "no constant term"; normalizing
198            // here keeps `Neg(Const(0.0))` from reporting `-0.0`.
199            constant: if c != 0.0 { c } else { 0.0 },
200            ..Quad2::default()
201        }
202    }
203
204    pub(crate) fn of_var(i: usize) -> Self {
205        let mut q = Quad2::default();
206        q.linear.insert(i, 1.0);
207        q
208    }
209
210    /// Total degree: 0, 1, or 2.
211    pub(crate) fn degree(&self) -> usize {
212        if !self.quadratic.is_empty() {
213            2
214        } else if !self.linear.is_empty() {
215            1
216        } else {
217            0
218        }
219    }
220
221    /// The value, when this form has no variables in it.
222    pub(crate) fn as_constant(&self) -> Option<f64> {
223        (self.degree() == 0).then_some(self.constant)
224    }
225
226    /// Number of stored (nonzero) variable terms.
227    fn width(&self) -> usize {
228        self.linear.len() + self.quadratic.len()
229    }
230
231    /// `a + b`.
232    ///
233    /// Two things here are not incidental. **Only the entries the smaller
234    /// side contributes are re-checked for zero** — the predecessor
235    /// re-scanned the whole accumulated map on every merge, which is `O(k²)`
236    /// down a `k`-term `Add` chain, and an `o0` chain is how a long sum
237    /// reaches this code. And **the smaller side is merged into the larger**,
238    /// whichever operand it is, so the same chain costs `O(k log k)` leaning
239    /// either way rather than only when it leans left. Choosing the direction
240    /// is free of arithmetic consequence: IEEE addition is commutative bit
241    /// for bit, `-0.0` included.
242    pub(crate) fn add(a: Quad2, b: Quad2) -> Quad2 {
243        let (mut acc, small) = if a.width() >= b.width() {
244            (a, b)
245        } else {
246            (b, a)
247        };
248        // Whether anything either side already knew it had rounded. Read
249        // once, before any merge, so the verdict below does not depend on
250        // the order the maps happen to iterate in.
251        let carried_inexact = acc.inexact || small.inexact;
252        let (mut dropped, mut inexact) = (false, false);
253        if small.constant != 0.0 {
254            let was = acc.constant;
255            acc.constant += small.constant;
256            // Same rule as a coefficient, for the same reason: the degree-0
257            // term is exempt from the "no stored zeros" invariant but not
258            // from arithmetic, and `mul` reads `constant == 0.0` as
259            // *annihilates*. A constant that cancelled *inexactly* is not
260            // one that was proven absent (gh #683); one that cancelled
261            // exactly is (gh #687).
262            inexact |= !add_is_exact(was, small.constant, acc.constant);
263            dropped |= was != 0.0 && acc.constant == 0.0;
264        }
265        for (i, c) in &small.linear {
266            let m = merge(&mut acc.linear, *i, *c);
267            dropped |= m.dropped;
268            inexact |= m.inexact;
269        }
270        for (k, c) in &small.quadratic {
271            let m = merge(&mut acc.quadratic, *k, *c);
272            dropped |= m.dropped;
273            inexact |= m.inexact;
274        }
275        // A term is only *lost* if something was dropped and some fold on
276        // the way here could not represent itself. Neither half is enough
277        // alone: an exact cancellation drops a term that was really zero,
278        // and an inexact fold that keeps its coefficient has hidden
279        // nothing — yet.
280        acc.lost_terms |= small.lost_terms || (dropped && (carried_inexact || inexact));
281        acc.inexact = carried_inexact || inexact;
282        acc
283    }
284
285    pub(crate) fn neg(mut self) -> Quad2 {
286        self.constant = -self.constant;
287        for c in self.linear.values_mut() {
288            *c = -*c;
289        }
290        for c in self.quadratic.values_mut() {
291            *c = -*c;
292        }
293        self
294    }
295
296    /// `self · s`, for a scalar `s`.
297    ///
298    /// The `prune` is not decoration, and neither is the flag it sets. A
299    /// nonzero coefficient times a nonzero `s` can still land on zero by
300    /// underflow — `1e-300 x0²` divided by `1e300` is reachable from a real
301    /// `.nl` body, through [`Op::Div`]'s reciprocal — so this is the
302    /// gh #683 shape again, arrived at by scaling rather than by summing.
303    /// Leaving the flushed entry stored would make an arithmetically
304    /// constant form report degree 2 to [`Quad2::degree`] and be refused as
305    /// degree 3 the moment anything multiplied it; dropping it without
306    /// recording the loss would make a genuinely degree-2 body look affine
307    /// to [`Quad2::lost_terms`]'s consumers. `prune` does the first and
308    /// this does the second. (`neg` needs neither: negation cannot reach
309    /// zero from a coefficient that was not already there, and it never
310    /// rounds.)
311    ///
312    /// A coefficient here can only reach zero by **underflow** — both
313    /// factors are nonzero — so unlike a cancelling sum there is no exact
314    /// case to spare (gh #687): `10⁻³⁰⁰ · 10⁻³⁰⁰` is a term the form can no
315    /// longer see, whatever the flags said before.
316    pub(crate) fn scale(mut self, s: f64) -> Quad2 {
317        if s == 0.0 {
318            // An exact zero annihilates every term, so this is a proof of
319            // degree 0 rather than a loss — nothing is being discarded that
320            // survived the multiplication, and `x · 0` is exact. The flags
321            // still ride along: the one route here is division by an
322            // infinity, and `0 · NaN` is not zero.
323            return Quad2 {
324                lost_terms: self.lost_terms,
325                inexact: self.inexact,
326                ..Quad2::default()
327            };
328        }
329        let was = self.constant;
330        self.constant *= s;
331        if was != 0.0 {
332            self.lost_terms |= self.constant == 0.0;
333            self.inexact |= !mul_is_exact(was, s, self.constant);
334        }
335        let mut inexact = false;
336        for c in self.linear.values_mut() {
337            let was = *c;
338            *c *= s;
339            inexact |= !mul_is_exact(was, s, *c);
340        }
341        for c in self.quadratic.values_mut() {
342            let was = *c;
343            *c *= s;
344            inexact |= !mul_is_exact(was, s, *c);
345        }
346        self.inexact |= inexact;
347        // A scale can underflow a live coefficient to zero, which is a lost
348        // term and used to be stored as a zero.
349        self.lost_terms |= self.prune();
350        self
351    }
352
353    /// Restore the "no stored zeros" invariant after an operation wrote
354    /// coefficients in bulk. Returns whether anything was dropped; what
355    /// that *means* is the caller's to say, because it depends on how the
356    /// coefficients got there (see [`scale`](Quad2::scale) and
357    /// [`mul`](Quad2::mul)).
358    fn prune(&mut self) -> bool {
359        let before = self.width();
360        self.linear.retain(|_, c| is_live(*c));
361        self.quadratic.retain(|_, c| is_live(*c));
362        self.width() != before
363    }
364
365    /// `self / d`, for a constant `d` the caller has already checked is
366    /// nonzero.
367    ///
368    /// Scales by the **reciprocal** (not `c / d`) so the arithmetic matches
369    /// what the recursive predecessor produced bit for bit; what is new is
370    /// that `fl(1/d)` is the reciprocal only when `d` is a power of two,
371    /// and `fma(r, d, −1) == 0` says so exactly. Recording it matters
372    /// because a coefficient no real arithmetic produced can still cancel
373    /// *exactly* against another one later — and that cancellation would
374    /// otherwise be read as a proof of absence (gh #687).
375    pub(crate) fn div_by_constant(self, d: f64) -> Quad2 {
376        let r = 1.0 / d;
377        let exact = r.is_normal() && d.is_normal() && r.mul_add(d, -1.0) == 0.0;
378        let mut out = self.scale(r);
379        out.inexact |= !exact;
380        out
381    }
382
383    /// Take on another form's [`lost_terms`](Quad2::lost_terms) and
384    /// [`inexact`](Quad2::inexact) history.
385    ///
386    /// The operand a `Div` or a `Pow` reads as a *constant* is a [`Quad2`]
387    /// like any other, and its arithmetic is part of this form's: dividing
388    /// by `fl(10²⁰⁰ · 10²⁰⁰)` is dividing by an infinity, and dividing by a
389    /// constant that swallowed a variable term is worse than that. The
390    /// value is read out through [`as_constant`](Quad2::as_constant), which
391    /// keeps neither fact, so the caller hands the form itself over here.
392    pub(crate) fn absorb_flags(&mut self, other: &Quad2) {
393        self.lost_terms |= other.lost_terms;
394        self.inexact |= other.inexact;
395    }
396
397    /// `self · other`, or `None` when the product would exceed total
398    /// degree 2 — past that the recognizer gives up and the caller routes
399    /// to the general NLP path.
400    pub(crate) fn mul(&self, other: &Quad2) -> Option<Quad2> {
401        if self.degree() + other.degree() > 2 {
402            return None;
403        }
404        let mut out = Quad2::default();
405        // Either operand's missing term is missing from the product too,
406        // and so is either operand's rounding.
407        let mut lost = self.lost_terms || other.lost_terms;
408        let carried_inexact = self.inexact || other.inexact;
409        let mut inexact = false;
410        // A product of two live coefficients that lands on zero underflowed
411        // — `(10⁻²⁰⁰·x)·(10⁻²⁰⁰·x)` is one monomial whose coefficient is not
412        // representable (gh #683) — and unlike a cancelling *sum* there is
413        // no exact case to spare, so `lost` is set on the spot rather than
414        // left to the cancellation rule below.
415        let product = |a: f64, b: f64, lost: &mut bool, inexact: &mut bool| -> f64 {
416            let t = a * b;
417            *lost |= !is_live(t);
418            *inexact |= !mul_is_exact(a, b, t);
419            t
420        };
421        if self.constant != 0.0 && other.constant != 0.0 {
422            // `10⁻²⁰⁰ · 10⁻²⁰⁰` is not zero, and the branches below read a
423            // zero constant as an annihilating one — which would take a
424            // degree-2 product down to nothing without a trace (gh #683).
425            out.constant = product(self.constant, other.constant, &mut lost, &mut inexact);
426        }
427        let mut dropped = false;
428        // constant × (linear, quadratic), both ways round.
429        for (a, b) in [(self, other), (other, self)] {
430            if a.constant == 0.0 {
431                continue;
432            }
433            for (i, c) in &b.linear {
434                let t = product(a.constant, *c, &mut lost, &mut inexact);
435                let m = accumulate(&mut out.linear, *i, t);
436                dropped |= m.dropped;
437                inexact |= m.inexact;
438            }
439            for (k, c) in &b.quadratic {
440                let t = product(a.constant, *c, &mut lost, &mut inexact);
441                let m = accumulate(&mut out.quadratic, *k, t);
442                dropped |= m.dropped;
443                inexact |= m.inexact;
444            }
445        }
446        // linear × linear. The degree guard above means at most one of the
447        // two operands carries quadratic terms, so this runs only when
448        // neither does and no ordering question arises.
449        for (i, a) in &self.linear {
450            for (j, b) in &other.linear {
451                let key = (*i.min(j), *i.max(j));
452                let t = product(*a, *b, &mut lost, &mut inexact);
453                let m = accumulate(&mut out.quadratic, key, t);
454                dropped |= m.dropped;
455                inexact |= m.inexact;
456            }
457        }
458        // What `prune` drops here that `accumulate` did not already see is a
459        // key whose *first* contribution was a zero — an underflowed product,
460        // and `lost` is set for it above. So the cancellation rule is the
461        // same one `add` applies: a drop counts as a loss only when some fold
462        // behind it rounded.
463        dropped |= out.prune();
464        out.lost_terms = lost || (dropped && (carried_inexact || inexact));
465        out.inexact = carried_inexact || inexact;
466        Some(out)
467    }
468}
469
470/// What folding one coefficient into a map did — the two facts
471/// [`Quad2::lost_terms`] is decided from.
472#[derive(Clone, Copy)]
473struct Merged {
474    /// The key is no longer stored: the fold reached exactly zero, or
475    /// `NaN`.
476    dropped: bool,
477    /// The fold **rounded**: what is stored (or what cancelled) is not what
478    /// exact arithmetic on the same two numbers would have left.
479    inexact: bool,
480}
481
482/// Add `c` to `map[key]`, keeping the "no stored zeros" invariant. Returns
483/// what that did, as [`Merged`].
484///
485/// Only the touched key can have become zero, which is what keeps a merge
486/// proportional to what it merged rather than to what it merged *into*.
487fn merge<K: Ord>(map: &mut BTreeMap<K, f64>, key: K, c: f64) -> Merged {
488    use std::collections::btree_map::Entry;
489    match map.entry(key) {
490        Entry::Occupied(mut e) => {
491            let a = *e.get();
492            let v = a + c;
493            let inexact = !add_is_exact(a, c, v);
494            if is_live(v) {
495                e.insert(v);
496                Merged {
497                    dropped: false,
498                    inexact,
499                }
500            } else {
501                e.remove();
502                Merged {
503                    dropped: true,
504                    inexact,
505                }
506            }
507        }
508        // Nothing was added to anything: a live `c` is stored as it stands,
509        // and a dead one is only reachable from a caller that already knows
510        // what its own zero means (`mul`'s underflowed products; `add` only
511        // ever passes stored — hence live — coefficients).
512        Entry::Vacant(e) => {
513            if is_live(c) {
514                e.insert(c);
515                Merged {
516                    dropped: false,
517                    inexact: false,
518                }
519            } else {
520                Merged {
521                    dropped: true,
522                    inexact: c.is_nan(),
523                }
524            }
525        }
526    }
527}
528
529/// Accumulate `t` into `map[key]`, the way [`Quad2::mul`] builds a product
530/// up term by term, and report what the fold did.
531///
532/// Unlike [`merge`] this leaves a zero stored — `mul` prunes once at the
533/// end — because a key it lands on twice must see the first contribution as
534/// a plain `0.0 + t`, bit for bit what the `+=` this replaced produced.
535/// Only a key that was already nonzero can be said to have *dropped*
536/// anything.
537fn accumulate<K: Ord>(map: &mut BTreeMap<K, f64>, key: K, t: f64) -> Merged {
538    let slot = map.entry(key).or_insert(0.0);
539    let was = *slot;
540    *slot = was + t;
541    Merged {
542        dropped: was != 0.0 && !is_live(*slot),
543        inexact: !add_is_exact(was, t, *slot),
544    }
545}
546
547/// The exactness predicates this recognizer's fold is decided from live in
548/// `pounce-common` since gh #673, because the `.nl` pipeline grew a second
549/// fold — `Σ 2wₖbₖbₖᵀ` in `QuadraticStructure::push_factored_form` — in a
550/// crate that cannot name this one. See [`pounce_common::exact`] for what
551/// they answer and why it is an exact question rather than a tolerance.
552use pounce_common::exact::{add_is_exact, is_live, mul_is_exact};
553
554/// One entry on the recognizer's explicit work stack.
555enum Step<'a> {
556    /// Lower this subexpression onto the value stack.
557    Visit(&'a Expr),
558    /// Combine values already on the value stack.
559    Apply(Op),
560}
561
562/// A pending combination, popped once its operands have been lowered.
563enum Op {
564    Neg,
565    Add,
566    Sub,
567    Mul,
568    Div,
569    Pow,
570    /// n-ary sum over the top `n` values.
571    Sum(usize),
572    /// Not an operation: the value now on top of the stack is the lowering
573    /// of the `Cse` body with this address, so record it before moving on.
574    CacheCse(*const Expr),
575}
576
577/// Lower an [`Expr`] to a [`Quad2`], or `None` if it contains anything the
578/// recognizer cannot prove is a degree-≤2 polynomial (transcendental ops,
579/// division by a non-constant, `Pow` with an exponent ∉ {0, 1, 2}, products
580/// of degree > 2, external calls, comparisons, `if-then-else`, `min`/`max`,
581/// …). `None` ⇒ treat as general nonlinear.
582///
583/// `Cse` nodes are inlined: a reference is mathematically its body, and
584/// every reference is an independent occurrence. A body is nevertheless
585/// lowered **once**, and its [`Quad2`] reused at every later reference
586/// (keyed on `Arc` identity). That is what makes the walk `Θ(nodes)` on a
587/// shared DAG instead of `Θ(2^depth)`; Q4 had to refuse a re-referenced
588/// body outright to bound the cost, and this is the memoization that
589/// refusal was waiting for (gh #588, Q5). It is bitwise neutral — the
590/// lowering of a body is a function of the body, so the reused value is
591/// the one a second lowering would have produced, bit for bit.
592///
593/// The walk is iterative. See the module docs for why that is a
594/// correctness property and not a style choice.
595pub fn recognize_expr(e: &Expr) -> Option<Quad2> {
596    let mut work: Vec<Step<'_>> = vec![Step::Visit(e)];
597    let mut vals: Vec<Quad2> = Vec::new();
598    // Lowered `Cse` bodies, by address. Only successes land here: a body
599    // that fails aborts the whole walk, so there is no negative result to
600    // remember.
601    let mut cse: std::collections::HashMap<*const Expr, Quad2> = std::collections::HashMap::new();
602
603    while let Some(step) = work.pop() {
604        match step {
605            Step::Visit(e) => match e {
606                Expr::Const(c) => vals.push(Quad2::of_constant(*c)),
607                Expr::Var(i) => vals.push(Quad2::of_var(*i)),
608                Expr::Cse(body) => {
609                    let key = std::sync::Arc::as_ptr(body);
610                    match cse.get(&key) {
611                        Some(q) => vals.push(q.clone()),
612                        None => {
613                            work.push(Step::Apply(Op::CacheCse(key)));
614                            work.push(Step::Visit(body));
615                        }
616                    }
617                }
618                Expr::Sum(items) => {
619                    work.push(Step::Apply(Op::Sum(items.len())));
620                    // Pushed forward, so they pop back to front and land on
621                    // the value stack with item 0 on top — which puts item
622                    // 0 at the *end* of the region `Op::Sum` drains, and
623                    // item n-1 at its start. The fold there walks that
624                    // region in reverse so the sum still accumulates front
625                    // to back, the order the recursive version summed in.
626                    for it in items {
627                        work.push(Step::Visit(it));
628                    }
629                }
630                Expr::Unary(UnaryOp::Neg, a) => {
631                    work.push(Step::Apply(Op::Neg));
632                    work.push(Step::Visit(a));
633                }
634                // Every other unary op is transcendental.
635                Expr::Unary(..) => return None,
636                Expr::Binary(op, a, b) => {
637                    let op = match op {
638                        BinOp::Add => Op::Add,
639                        BinOp::Sub => Op::Sub,
640                        BinOp::Mul => Op::Mul,
641                        BinOp::Div => Op::Div,
642                        BinOp::Pow => Op::Pow,
643                        // atan2 and any other binary opcode.
644                        _ => return None,
645                    };
646                    work.push(Step::Apply(op));
647                    // `b` under `a`: `a` pops first and is lowered first.
648                    work.push(Step::Visit(b));
649                    work.push(Step::Visit(a));
650                }
651                // External calls are opaque; comparisons, logicals,
652                // conditionals and n-ary min/max are the control-flow `.nl`
653                // opcodes. None is provably polynomial ⇒ route to NLP.
654                _ => return None,
655            },
656            Step::Apply(Op::CacheCse(key)) => {
657                // The body's value is already on the stack and stays there;
658                // this only records it for the next reference.
659                cse.insert(key, vals.last()?.clone());
660            }
661            Step::Apply(op) => {
662                let combined = match op {
663                    Op::CacheCse(_) => unreachable!("handled above"),
664                    Op::Sum(n) => {
665                        // The items are the top `n` values, item 0 on top,
666                        // so item 0 is the *last* thing `drain` yields and
667                        // item n-1 the first. Reverse it: floating-point
668                        // addition is not associative, and summing back to
669                        // front would disagree with the recursive
670                        // predecessor — and with the AD tape — by an ulp
671                        // whenever two items share a monomial key.
672                        let at = vals.len().checked_sub(n)?;
673                        let mut acc = Quad2::default();
674                        for p in vals.drain(at..).rev() {
675                            acc = Quad2::add(acc, p);
676                        }
677                        acc
678                    }
679                    Op::Neg => vals.pop()?.neg(),
680                    Op::Add => {
681                        let (a, b) = pop2(&mut vals)?;
682                        Quad2::add(a, b)
683                    }
684                    Op::Sub => {
685                        let (a, b) = pop2(&mut vals)?;
686                        Quad2::add(a, b.neg())
687                    }
688                    Op::Mul => {
689                        let (a, b) = pop2(&mut vals)?;
690                        a.mul(&b)?
691                    }
692                    Op::Div => {
693                        // Division is polynomial only by a nonzero constant,
694                        // and scales by the reciprocal (not `c / d`) so the
695                        // arithmetic matches what the recursive predecessor
696                        // produced bit for bit.
697                        let (a, b) = pop2(&mut vals)?;
698                        let d = b.as_constant()?;
699                        if d == 0.0 {
700                            return None;
701                        }
702                        let mut out = a.div_by_constant(d);
703                        // The divisor's own history comes along: `as_constant`
704                        // reads a number out of a form and leaves the flags
705                        // behind.
706                        out.absorb_flags(&b);
707                        out
708                    }
709                    Op::Pow => {
710                        // Polynomial only for constant exponents in {0, 1, 2}.
711                        let (a, b) = pop2(&mut vals)?;
712                        let exp = b.as_constant()?;
713                        let mut out = if exp == 0.0 {
714                            Quad2::of_constant(1.0)
715                        } else if exp == 1.0 {
716                            a
717                        } else if exp == 2.0 {
718                            a.mul(&a)?
719                        } else {
720                            return None;
721                        };
722                        // Same as `Div`: the exponent is read out of a form,
723                        // and one that lost a term is not the exponent it
724                        // looks like.
725                        out.absorb_flags(&b);
726                        out
727                    }
728                };
729                vals.push(combined);
730            }
731        }
732    }
733
734    debug_assert_eq!(vals.len(), 1, "one value per lowered expression");
735    vals.pop()
736}
737
738/// Pop a binary operator's two operands, left first.
739fn pop2(vals: &mut Vec<Quad2>) -> Option<(Quad2, Quad2)> {
740    let b = vals.pop()?;
741    let a = vals.pop()?;
742    Some((a, b))
743}
744
745/// Attempt to read an expression as a polynomial of total degree ≤ 2 and
746/// return its Hessian (constant, since the form is quadratic). `None` if
747/// the expression is not provably quadratic ⇒ treat as general nonlinear.
748pub fn analyze_quadratic(e: &Expr) -> Option<QuadHessian> {
749    analyze_quadratic_full(e).map(|(h, _, _)| h)
750}
751
752/// Like [`analyze_quadratic`] but also returns the degree-1 (linear)
753/// coefficients *and* the degree-0 (constant) term of the form:
754/// `(Hessian, [(var, coef), …], constant)`.
755///
756/// AMPL folds the linear part of a nonlinear term into the objective's
757/// nonlinear expression tree (the `−6·x₀` of `(x₀−3)²`, say) rather than
758/// the linear section. Callers building the QP objective vector `c` must
759/// add these in, exactly as the NLP path's `eval_f` sums the linear
760/// section *and* the nonlinear tree — otherwise the linear shift is
761/// silently dropped and the convex solve minimizes the wrong objective.
762///
763/// The **constant** is returned for the same reason: AMPL/Pyomo also fold
764/// the objective's degree-0 term into the nonlinear tree (the `+9` of
765/// `(x₀−3)²`), where it does *not* land in `NlProblem::obj_constant`. It
766/// is irrelevant to the minimizer but is part of the *reported objective
767/// value*; dropping it makes the convex solve report an objective off by
768/// that constant versus the NLP path.
769pub fn analyze_quadratic_full(e: &Expr) -> Option<QuadForm> {
770    Some(quad_form_readout(&recognize_expr(e)?))
771}
772
773/// The `(Hessian, linear, constant)` read-out of an already-recognized
774/// form — the second half of [`analyze_quadratic_full`], split out so a
775/// caller holding a [`Quad2`] the *parser* produced (gh #588, Q5) reaches
776/// the identical numbers by the identical route. There is exactly one
777/// conversion in the crate; a second one is a second thing to keep in step.
778pub fn quad_form_readout(q: &Quad2) -> QuadForm {
779    // ∂²(c·xᵢxⱼ)/∂xᵢ∂xⱼ = c for i≠j; ∂²(c·xᵢ²)/∂xᵢ² = 2c.
780    let mut h: QuadHessian = q
781        .quadratic
782        .iter()
783        .map(|(&(i, j), c)| ((i, j), if i == j { 2.0 * c } else { *c }))
784        .collect();
785    // Drop explicit zeros so `is_empty()` means "linear".
786    h.retain(|_, v| v.abs() > 0.0);
787    let lin: Vec<(usize, f64)> = q.linear.iter().map(|(i, c)| (*i, *c)).collect();
788    // `0.0 +` normalizes `-0.0`, which is how this form spells "absent".
789    (h, lin, 0.0 + q.constant)
790}
791
792/// Is this expression **already** a flat sum of monomials — that is, would
793/// reading it as `½xᵀHx + aᵀx + c` reproduce the same additions the writer
794/// wrote, rather than algebraically expanding something it did not?
795///
796/// This is the gate on the constant-structure evaluator (gh #588, Q4), and
797/// it exists because expanding a quadratic is not an accuracy-neutral
798/// rewrite. `(xᵢ − xⱼ)²` evaluated as written squares a small residual;
799/// evaluated as `xᵢ² − 2xᵢxⱼ + xⱼ²` it cancels two large numbers, and on
800/// `airport.nl` — 84 coordinates around 10³, every row a squared distance —
801/// that difference is enough to take an adaptive-μ solve from stopping at a
802/// tiny step in 16 iterations to grinding out the 300-iteration cap at the
803/// same objective. Which is precisely the gh #544 failure mode: the right
804/// answer, slowly.
805///
806/// So the rule is *exactness*, not magnitude: a form is admitted only when
807/// the recognizer's read-out does no algebra the tape would not have done.
808/// A sum of monomials qualifies — which is exactly how AMPL emits the
809/// `qcqp*` family (`o54` over `o2 n0.5 o2 o2 n<c> v<i> v<j>`), so the target
810/// of the phase is unaffected — and a `Pow` or a `Mul` over a non-atomic
811/// operand does not.
812///
813/// What this predicate is **not** is a verdict on whether a body can be
814/// evaluated from constant structure at all. It is a verdict on one
815/// *representation*. Q4 shipped as though the two were the same, and the
816/// consequence was that a model written in factored form kept its tape and
817/// gained nothing — 41 of `airport.nl`'s 42 rows, and every sum of squared
818/// residuals ever written. gh #673 is that gap, and
819/// [`recognize_factored_quadratic`] closes it by keeping the writer's own
820/// grouping instead of expanding it: same constant structure, no algebra
821/// the tape did not do. A body this predicate refuses is offered to that
822/// one before it falls back to a tape.
823///
824/// ## A re-referenced `Cse` body is *skipped*, not refused
825///
826/// This walk and [`recognize_expr`] behind it both inline a `Cse` body at
827/// **every** reference, so a DAG whose bodies are shared cost `Θ(2^depth)`
828/// rather than `Θ(nodes)` before either was memoized. `nl_reader`'s
829/// `shared_dag_walks_are_memoized_not_exponential` builds exactly that shape
830/// — 30 levels, each a `Cse` referenced twice — and taking it through this
831/// gate in Q4 measured **0.00 s → 172 s** on a model the tape path loads
832/// instantly. At depth 40 it would not have returned at all.
833///
834/// Q4 bounded that by ending the walk in `false` at the second reference,
835/// which cost reach to buy termination and was explicitly left for Q5 to
836/// replace. It is replaced here: a body reached a second time in the same
837/// mode is **skipped**, which is sound for exactly the reason
838/// `nl_reader::validate_expr` documents for the same trick — this walk
839/// aborts on the first violation, so reaching a body a second time proves
840/// the first visit found none. The verdict is unchanged for every DAG that
841/// is a tree, and a shared DAG that *is* an expanded quadratic is now
842/// admitted instead of refused. The mode is part of the key: a body is
843/// legal on the sum spine if it is a sum of monomials, and legal inside a
844/// monomial only if it is itself a monomial, so the two answers are cached
845/// apart.
846///
847/// Iterative for the same reason [`recognize_expr`] is.
848pub fn is_expanded_quadratic(e: &Expr) -> bool {
849    // The sum spine: `Add`/`Sub`/`Neg`/`Sum` may nest freely, and every
850    // leaf of that spine must be a monomial.
851    let mut seen: BTreeSet<(*const Expr, bool)> = BTreeSet::new();
852    let mut spine: Vec<&Expr> = vec![e];
853    while let Some(e) = spine.pop() {
854        match e {
855            Expr::Sum(items) => spine.extend(items.iter()),
856            Expr::Binary(BinOp::Add | BinOp::Sub, a, b) => {
857                spine.push(a);
858                spine.push(b);
859            }
860            Expr::Unary(UnaryOp::Neg, a) => spine.push(a),
861            Expr::Cse(body) => {
862                if seen.insert((std::sync::Arc::as_ptr(body), false)) {
863                    spine.push(body);
864                }
865            }
866            other => {
867                if !is_monomial(other, &mut seen) {
868                    return false;
869                }
870            }
871        }
872    }
873    true
874}
875
876/// Is this expression a single monomial — the leaf shape
877/// [`is_expanded_quadratic`] admits on its sum spine?
878///
879/// Public because the parse-time recognizer has to answer the same
880/// question about a `V`-segment body it has already parsed, and must
881/// answer it with *this* code rather than a second copy of the rule.
882pub fn is_monomial_expr(e: &Expr) -> bool {
883    let mut seen: BTreeSet<(*const Expr, bool)> = BTreeSet::new();
884    is_monomial(e, &mut seen)
885}
886
887/// A single product term: constants and variables multiplied together, with
888/// no addition anywhere inside it. `xᵢxⱼ`, `0.5·c·xᵢ·xⱼ`, `xᵢ²` and `xᵢ/3`
889/// all qualify; `(xᵢ − xⱼ)²` and `(xᵢ + 1)·xⱼ` do not.
890///
891/// Degree is not checked here — [`recognize_expr`] already refused anything
892/// past 2 by the time this runs, and duplicating the rule would only give
893/// the two a way to disagree.
894fn is_monomial(e: &Expr, seen: &mut BTreeSet<(*const Expr, bool)>) -> bool {
895    let mut work: Vec<&Expr> = vec![e];
896    while let Some(e) = work.pop() {
897        match e {
898            Expr::Const(_) | Expr::Var(_) => {}
899            // Shared with the spine walk, and skipped on revisit for the
900            // same reason — see [`is_expanded_quadratic`]. The `true` in the
901            // key is "seen in monomial mode": a body cleared on the spine
902            // has not been cleared here.
903            Expr::Cse(body) => {
904                if seen.insert((std::sync::Arc::as_ptr(body), true)) {
905                    work.push(body);
906                }
907            }
908            Expr::Unary(UnaryOp::Neg, a) => work.push(a),
909            Expr::Binary(BinOp::Mul | BinOp::Div, a, b) => {
910                work.push(a);
911                work.push(b);
912            }
913            // `x^2` is one monomial; `(x - y)^2` is an expansion. The
914            // exponent itself may be any constant expression — the
915            // recognizer has already restricted it to {0, 1, 2}.
916            Expr::Binary(BinOp::Pow, a, b) => {
917                if !matches!(
918                    a.as_ref(),
919                    Expr::Const(_) | Expr::Var(_) | Expr::Unary(UnaryOp::Neg, _)
920                ) {
921                    return false;
922                }
923                work.push(a);
924                work.push(b);
925            }
926            _ => return false,
927        }
928    }
929    true
930}
931
932/// True if the expression is the literal constant zero the `.nl` reader
933/// uses for "no nonlinear part".
934pub fn is_trivially_zero(e: &Expr) -> bool {
935    matches!(e, Expr::Const(c) if *c == 0.0)
936}
937
938// ---------------------------------------------------------------------
939// Factored forms — the writer's own grouping, kept
940// ---------------------------------------------------------------------
941
942/// One `w·(bᵀx + d)²` term of a [`FactoredQuadratic`].
943///
944/// `coefs` is `b`, ascending by variable index and free of stored zeros —
945/// the same convention [`Quad2::linear`] keeps, because that is where it
946/// comes from. It may be **empty**, which is a square of a constant: the
947/// term is then `w·d²`, contributes nothing to the gradient or the
948/// Hessian, and is kept rather than folded so that the value is still
949/// computed the way it was written.
950#[derive(Debug, Clone, PartialEq)]
951pub struct SquaredAffine {
952    /// The constant the square is multiplied by, sign of the sum spine
953    /// already folded in.
954    pub weight: f64,
955    /// `b`, ascending by variable index.
956    pub coefs: Vec<(usize, f64)>,
957    /// `d`.
958    pub constant: f64,
959}
960
961/// A degree-2 form kept as `Σ wₖ(bₖᵀx + dₖ)² + aᵀx + c` — the shape the
962/// `.nl` writer wrote, rather than its expansion about the origin.
963///
964/// This is the representation gh #673 asks for and
965/// [`is_expanded_quadratic`] exists to refuse the alternative to. Reading
966/// `(x − 500000)²` back as `x² − 10⁶x + 2.5·10¹¹` cancels five digits;
967/// reading it back as one squared residual repeats exactly the
968/// multiplication the tape performs, so admitting it costs no accuracy at
969/// all.
970///
971/// The degree-≤1 leftovers (`linear`, `constant`) are the monomials the
972/// writer folded into the same tree — the `+ 3y` of `(x − 1)² + 3y`. They
973/// are summed as coefficients rather than kept term by term, which is the
974/// same reassociation the expanded path already makes and is gated the
975/// same way (see [`Quad2::lost_terms`]).
976#[derive(Debug, Clone, PartialEq)]
977pub struct FactoredQuadratic {
978    /// The squared affine terms, in the order the spine walk met them.
979    pub squares: Vec<SquaredAffine>,
980    /// `a`, ascending by variable index.
981    pub linear: Vec<(usize, f64)>,
982    /// `c`.
983    pub constant: f64,
984}
985
986/// Recognize a degree-2 body as a sum of **squared affine forms** plus
987/// degree-≤1 leftovers, keeping the squares factored (gh #673).
988///
989/// This is the second half of the gate [`is_expanded_quadratic`] opens.
990/// That predicate admits a body whose read-out repeats the additions the
991/// writer wrote; a factored body fails it, and used to keep its AD tape for
992/// good reason — expanding it to the origin cancels. What this returns is
993/// the *third* option: a read-out that is not an expansion either, because
994/// it stores the writer's own linear forms and squares them at evaluation
995/// time exactly as the tape does.
996///
997/// ## What is admitted
998///
999/// The sum spine (`Add`/`Sub`/`Neg`/`Sum`) may nest freely. Every leaf of
1000/// it must be one of:
1001///
1002/// * a **square**: a product of constants with exactly one `Pow(base, 2)`,
1003///   where `base` is itself an [`is_expanded_quadratic`] body of degree
1004///   ≤ 1 — `(xᵢ − xⱼ)²`, `0.5·(2x − y + 3)²`, `x²`;
1005/// * a **degree-2 monomial on the diagonal**, `c·xᵢ²`, which is the same
1006///   term wearing a different opcode and is stored as `c·(xᵢ)²`;
1007/// * a **degree-≤1 monomial**, which folds into `a` and `c`.
1008///
1009/// Anything else — a cross monomial `c·xᵢxⱼ`, a product of two *different*
1010/// affine forms, a transcendental — refuses the whole body, which then
1011/// keeps its tape exactly as it did before this function existed. The
1012/// refusal is deliberate rather than incidental: a mixed form would need
1013/// both an expanded and a factored quadratic part stored side by side, and
1014/// nothing in the corpus asks for one.
1015///
1016/// At least one square with a variable in it is required. Without that the
1017/// body is affine or already expanded, and [`is_expanded_quadratic`] is
1018/// the path that serves it.
1019///
1020/// ## Exactness
1021///
1022/// Each admitted square evaluates as `w·(d + Σbᵢxᵢ)²`. Against the tape
1023/// that is a reassociation of an affine sum and a folding of the writer's
1024/// own constants into `w` — the same latitude [`Quad2::add`] and
1025/// [`Quad2::scale`] already take on the expanded path — and *not* an
1026/// algebraic expansion. The `2.4e-5` disagreement gh #673 records for
1027/// `(x − 500000)²` comes from the expansion; there is none here. Measured
1028/// end to end: with the outer sum left naive, `airport.nl`'s 42 rows
1029/// evaluate bit-identically to the 42 tapes they replace.
1030///
1031/// The *terms* are the tape's; the sum over them is not, because
1032/// `QuadraticStructure::value` compensates it (gh #702). That makes the
1033/// read-out slightly better than the tape rather than equal to it, which is
1034/// gh #702's deliberate trade and the one line the fixture sweep moves.
1035///
1036/// The degree-≤1 leftovers are held to the same [`Quad2::lost_terms`] gate
1037/// the expanded path uses (gh #685), so a leftover that cancelled inexactly
1038/// refuses the body rather than dropping a term out of it.
1039///
1040/// ## A `Cse` on the sum spine refuses the body
1041///
1042/// [`is_expanded_quadratic`] may *skip* a re-referenced body because it
1043/// answers yes/no and the first visit already decided it. This one
1044/// accumulates terms, so skipping would silently drop them and visiting
1045/// every reference is `Θ(2^depth)` on a shared DAG — the shape
1046/// `nl_reader`'s `shared_dag_walks_are_memoized_not_exponential` builds.
1047/// So a `Cse` reached on the spine ends the walk, and the body keeps its
1048/// tape, which is what `ConHybrid` is for. Inside a square's base or a
1049/// monomial there is no such problem: [`recognize_expr`] memoizes on `Arc`
1050/// identity and the two predicates answer yes/no.
1051///
1052/// Iterative, for the reason the module docs give.
1053pub fn recognize_factored_quadratic(e: &Expr) -> Option<FactoredQuadratic> {
1054    let mut seen: BTreeSet<(*const Expr, bool)> = BTreeSet::new();
1055    // The spine, each entry carrying the sign the enclosing `Sub`/`Neg`
1056    // chain gives it. `±1` exactly, so folding it into a weight or negating
1057    // a form is not arithmetic.
1058    let mut spine: Vec<(&Expr, f64)> = vec![(e, 1.0)];
1059    let mut squares: Vec<SquaredAffine> = Vec::new();
1060    let mut rest = Quad2::default();
1061
1062    while let Some((e, sign)) = spine.pop() {
1063        match e {
1064            // Pushed back to front so the leaves *pop* in source order.
1065            // The degree-≤1 leftovers are folded into one `Quad2` by
1066            // floating-point addition, and source order is the order a
1067            // reader can reason about — not, to be exact about what this
1068            // buys, the same association `recognize_expr` uses: that folds
1069            // the *tree*, and `Quad2::add` additionally swaps its operands
1070            // by map width. So the two can still differ in the last ulp on
1071            // a leaning spine. What keeps that from mattering is
1072            // `lost_terms`, which refuses any body where the difference
1073            // could be a dropped term rather than a rounded one.
1074            Expr::Sum(items) => spine.extend(items.iter().rev().map(|it| (it, sign))),
1075            Expr::Binary(BinOp::Add, a, b) => {
1076                spine.push((b, sign));
1077                spine.push((a, sign));
1078            }
1079            Expr::Binary(BinOp::Sub, a, b) => {
1080                spine.push((b, -sign));
1081                spine.push((a, sign));
1082            }
1083            Expr::Unary(UnaryOp::Neg, a) => spine.push((a, -sign)),
1084            // See the docs: accumulating terms is what makes skipping
1085            // unsound here, and visiting is what makes it exponential.
1086            Expr::Cse(_) => return None,
1087            leaf => {
1088                if let Some((weight, base)) = peel_square(leaf) {
1089                    squares.push(admit_square(sign * weight, base)?);
1090                    continue;
1091                }
1092                if !is_monomial(leaf, &mut seen) {
1093                    return None;
1094                }
1095                let q = recognize_expr(leaf)?;
1096                match diagonal_square(&q) {
1097                    Some((i, c)) => squares.push(SquaredAffine {
1098                        weight: sign * c,
1099                        coefs: vec![(i, 1.0)],
1100                        constant: 0.0,
1101                    }),
1102                    // A cross term `c·xᵢxⱼ` lands here and refuses: it is
1103                    // not a square, and storing it expanded alongside the
1104                    // squares is the mixed representation this refuses.
1105                    None if !q.quadratic().is_empty() => return None,
1106                    None => rest = Quad2::add(rest, if sign < 0.0 { q.neg() } else { q }),
1107                }
1108            }
1109        }
1110    }
1111
1112    // Nothing factored to keep ⇒ this is not our case. `is_expanded_quadratic`
1113    // serves an expanded body and an affine one needs no form at all.
1114    if !squares.iter().any(|t| !t.coefs.is_empty()) {
1115        return None;
1116    }
1117    // The leftovers are held to the expanded path's gate, for its reasons
1118    // (gh #685): a term that went missing inexactly is a term the form no
1119    // longer evaluates.
1120    if !rest.quadratic().is_empty() || rest.lost_terms() {
1121        return None;
1122    }
1123    Some(FactoredQuadratic {
1124        squares,
1125        linear: rest.linear().iter().map(|(&i, &c)| (i, c)).collect(),
1126        // `0.0 +` normalizes `-0.0`, which is how a `Quad2` spells "absent".
1127        constant: 0.0 + rest.constant(),
1128    })
1129}
1130
1131/// `c·xᵢ²` read off a recognized monomial: the one degree-2 shape that is
1132/// a square without being written as one. `None` for anything else,
1133/// including a cross term and anything carrying a linear or constant part.
1134///
1135/// The coefficient is the **polynomial** one — `Quad2::quadratic` is not
1136/// the Hessian — so `c·xᵢ²` is stored with weight `c` and not `2c`.
1137fn diagonal_square(q: &Quad2) -> Option<(usize, f64)> {
1138    if q.lost_terms() || !q.linear().is_empty() || q.constant() != 0.0 {
1139        return None;
1140    }
1141    match q.quadratic().iter().next() {
1142        Some((&(i, j), &c)) if i == j && q.quadratic().len() == 1 => Some((i, c)),
1143        _ => None,
1144    }
1145}
1146
1147/// Turn `weight · base²` into a stored term, or refuse the body.
1148///
1149/// The base has to clear both of the expanded path's gates — the shape gate
1150/// [`is_expanded_quadratic`] (so reading its coefficients back repeats the
1151/// writer's own additions) and the [`Quad2::lost_terms`] gate (so none of
1152/// them went missing) — and be degree ≤ 1, which is what makes the product
1153/// a square rather than a quartic.
1154fn admit_square(weight: f64, base: &Expr) -> Option<SquaredAffine> {
1155    if !is_expanded_quadratic(base) {
1156        return None;
1157    }
1158    let q = recognize_expr(base)?;
1159    if !q.quadratic().is_empty() || q.lost_terms() {
1160        return None;
1161    }
1162    Some(SquaredAffine {
1163        weight,
1164        coefs: q.linear().iter().map(|(&i, &c)| (i, c)).collect(),
1165        constant: 0.0 + q.constant(),
1166    })
1167}
1168
1169/// Split a leaf into `(constant weight, squared base)`, or `None` when it
1170/// is not a constant multiple of exactly one square.
1171///
1172/// The product tree may nest `Mul`, `Div` and `Neg` freely; every leaf of
1173/// it must be a constant except for the single `Pow(base, 2)`, which may
1174/// not sit under a division. Folding several constants into one `weight`
1175/// reassociates the writer's own constants and nothing else.
1176///
1177/// The order is worth stating precisely, because it is not source order:
1178/// this pops an explicit stack, so `a·(b·(c·s))` folds as `((1·c)·b)·a`.
1179/// One multiply's worth of rounding either way, on constants the writer
1180/// wrote next to each other — the same latitude `Quad2::scale` takes on the
1181/// expanded arm — but a reader reconstructing a last-ulp difference by hand
1182/// needs the real order, not the written one.
1183fn peel_square(e: &Expr) -> Option<(f64, &Expr)> {
1184    let mut work: Vec<(&Expr, bool)> = vec![(e, false)];
1185    let mut weight = 1.0f64;
1186    let mut base: Option<&Expr> = None;
1187    while let Some((e, recip)) = work.pop() {
1188        match e {
1189            Expr::Const(c) => weight = if recip { weight / c } else { weight * c },
1190            Expr::Unary(UnaryOp::Neg, a) => {
1191                weight = -weight;
1192                work.push((a, recip));
1193            }
1194            Expr::Binary(BinOp::Mul, a, b) => {
1195                work.push((a, recip));
1196                work.push((b, recip));
1197            }
1198            Expr::Binary(BinOp::Div, a, b) => {
1199                work.push((a, recip));
1200                work.push((b, !recip));
1201            }
1202            Expr::Binary(BinOp::Pow, a, b) if matches!(b.as_ref(), Expr::Const(c) if *c == 2.0) => {
1203                // A square under a division is `1/(…)²`, which is not one.
1204                // A second square makes the leaf degree 4.
1205                if recip || base.is_some() {
1206                    return None;
1207                }
1208                base = Some(a);
1209            }
1210            _ => return None,
1211        }
1212    }
1213    // A weight that is not finite (`x²/0`) is not a form anyone should
1214    // evaluate from stored coefficients, and one that is exactly zero has
1215    // annihilated a term the tape still computes.
1216    let base = base?;
1217    (weight.is_finite() && weight != 0.0).then_some((weight, base))
1218}
1219
1220#[cfg(test)]
1221mod tests {
1222    use super::*;
1223
1224    fn sq(i: usize) -> Expr {
1225        Expr::Binary(
1226            BinOp::Pow,
1227            Box::new(Expr::Var(i)),
1228            Box::new(Expr::Const(2.0)),
1229        )
1230    }
1231
1232    #[test]
1233    fn quadratic_diagonal() {
1234        // (x0 - 1)^2  =>  x0^2 - 2 x0 + 1
1235        let e = Expr::Binary(
1236            BinOp::Pow,
1237            Box::new(Expr::Binary(
1238                BinOp::Sub,
1239                Box::new(Expr::Var(0)),
1240                Box::new(Expr::Const(1.0)),
1241            )),
1242            Box::new(Expr::Const(2.0)),
1243        );
1244        let (h, lin, c) = analyze_quadratic_full(&e).expect("degree-2 polynomial");
1245        assert_eq!(h.get(&(0, 0)), Some(&2.0));
1246        assert_eq!(lin, vec![(0, -2.0)]);
1247        assert_eq!(c, 1.0);
1248    }
1249
1250    #[test]
1251    fn cross_term_hessian() {
1252        // x0 · x1 => H[0,1] = 1
1253        let e = Expr::Binary(BinOp::Mul, Box::new(Expr::Var(0)), Box::new(Expr::Var(1)));
1254        let h = analyze_quadratic(&e).expect("degree-2");
1255        assert_eq!(h.get(&(0, 1)), Some(&1.0));
1256    }
1257
1258    #[test]
1259    fn rejects_transcendental_and_cubic() {
1260        assert!(analyze_quadratic(&Expr::Unary(UnaryOp::Sin, Box::new(Expr::Var(0)))).is_none());
1261        let cubic = Expr::Binary(
1262            BinOp::Pow,
1263            Box::new(Expr::Var(0)),
1264            Box::new(Expr::Const(3.0)),
1265        );
1266        assert!(analyze_quadratic(&cubic).is_none());
1267        // x0² · x1 — degree 3 by multiplication rather than by exponent.
1268        let deg3 = Expr::Binary(BinOp::Mul, Box::new(sq(0)), Box::new(Expr::Var(1)));
1269        assert!(analyze_quadratic(&deg3).is_none());
1270    }
1271
1272    #[test]
1273    fn division_by_a_constant_scales_by_the_reciprocal() {
1274        // x0² / 3 — the coefficient must be `2 · (1/3)`, which is what the
1275        // Hessian of `x0²/3` came out as before this module existed, and is
1276        // *not* bitwise `2/3`.
1277        let e = Expr::Binary(BinOp::Div, Box::new(sq(0)), Box::new(Expr::Const(3.0)));
1278        let h = analyze_quadratic(&e).expect("degree-2");
1279        assert_eq!(h.get(&(0, 0)), Some(&(2.0 * (1.0 / 3.0))));
1280        // Division by a variable is not polynomial.
1281        let e = Expr::Binary(BinOp::Div, Box::new(sq(0)), Box::new(Expr::Var(1)));
1282        assert!(analyze_quadratic(&e).is_none());
1283    }
1284
1285    #[test]
1286    fn scaling_a_coefficient_to_zero_drops_it_like_cancellation_does() {
1287        // 1e-300·x0² divided by 1e300. Neither the coefficient nor the
1288        // divisor is zero, so the whole-form `s == 0.0` shortcut does not
1289        // fire — the product underflows instead, and the entry has to go
1290        // the same way a cancelled one does.
1291        let tiny = Expr::Binary(BinOp::Mul, Box::new(Expr::Const(1e-300)), Box::new(sq(0)));
1292        let flushed = Expr::Binary(BinOp::Div, Box::new(tiny), Box::new(Expr::Const(1e300)));
1293        let h = analyze_quadratic(&flushed).expect("degree-2 at worst");
1294        assert!(
1295            h.is_empty(),
1296            "underflowed coefficient was kept as a structural nonzero: {h:?}"
1297        );
1298        // Storage and degree are the two halves of gh #683 and this route
1299        // reaches both: the entry is gone from the map, *and* the form
1300        // says so, so a consumer asking whether the body is affine gets
1301        // "not established" rather than "yes".
1302        let q = recognize_expr(&flushed).expect("degree-2 at worst");
1303        assert!(
1304            q.lost_terms(),
1305            "a coefficient that underflowed in `scale` was dropped silently"
1306        );
1307        // And the degree has to go with it for the *storage* question:
1308        // multiplying by another variable must still be recognized rather
1309        // than refused as degree 3.
1310        let times_x1 = Expr::Binary(BinOp::Mul, Box::new(flushed), Box::new(Expr::Var(1)));
1311        assert!(
1312            analyze_quadratic(&times_x1).is_some(),
1313            "a form scaled to nothing was refused as degree 3"
1314        );
1315    }
1316
1317    #[test]
1318    fn cancellation_drops_the_term_and_the_degree_with_it() {
1319        // x0² − x0² is linear (empty Hessian), not a quadratic with a zero
1320        // coefficient — otherwise `x0²−x0²` times `x1` would be refused as
1321        // degree 3.
1322        let zero = Expr::Binary(BinOp::Sub, Box::new(sq(0)), Box::new(sq(0)));
1323        let h = analyze_quadratic(&zero).expect("degree-2 at worst");
1324        assert!(h.is_empty());
1325        let times_x1 = Expr::Binary(BinOp::Mul, Box::new(zero), Box::new(Expr::Var(1)));
1326        assert!(analyze_quadratic(&times_x1).is_some());
1327    }
1328
1329    /// The gh #687 distinction, at the level it is made: an **exact**
1330    /// cancellation leaves the form complete, and an absorbed term does
1331    /// not. Both bodies drop a coefficient; only one of them lost anything
1332    /// doing it.
1333    #[test]
1334    fn an_exact_cancellation_is_not_a_lost_term() {
1335        // x0² − x0²: `fl(1) + fl(−1)` is exactly `0`, so the body really is
1336        // degree 0 and the maps say so with nothing held back.
1337        let zero = Expr::Binary(BinOp::Sub, Box::new(sq(0)), Box::new(sq(0)));
1338        let q = recognize_expr(&zero).expect("degree-2 at worst");
1339        assert!(q.quadratic().is_empty());
1340        assert!(!q.inexact(), "an exact fold reported rounding");
1341        assert!(
1342            !q.lost_terms(),
1343            "x0² − x0² was refused a fast path it is entitled to",
1344        );
1345
1346        // 2⁵³·x0² + x0² − 2⁵³·x0², front to back: the `x0²` is absorbed by
1347        // `fl(2⁵³ + 1) = 2⁵³` — an inexact add — and the exact `− 2⁵³` only
1348        // makes the loss visible.
1349        let big = (1u64 << 53) as f64;
1350        let scaled = |c: f64| Expr::Binary(BinOp::Mul, Box::new(Expr::Const(c)), Box::new(sq(0)));
1351        let absorbing = Expr::Sum(vec![scaled(big), sq(0), scaled(-big)]);
1352        let q = recognize_expr(&absorbing).expect("degree-2 at worst");
1353        assert!(q.quadratic().is_empty());
1354        assert!(q.inexact(), "the absorbing add was not seen to round");
1355        assert!(
1356            q.lost_terms(),
1357            "a body whose x0² was absorbed was reported complete",
1358        );
1359    }
1360
1361    /// The absorption is recorded where it happens, which is the point of
1362    /// the sharpened gate: `2⁵³·x0² + x0²` has lost the `x0²` already, and
1363    /// the form says so before anything cancels — while its coefficient is
1364    /// still stored and its degree is still 2.
1365    #[test]
1366    fn the_inexact_fold_is_flagged_before_anything_drops() {
1367        let big = (1u64 << 53) as f64;
1368        let scaled = |c: f64| Expr::Binary(BinOp::Mul, Box::new(Expr::Const(c)), Box::new(sq(0)));
1369        let e = Expr::Binary(BinOp::Add, Box::new(scaled(big)), Box::new(sq(0)));
1370        let q = recognize_expr(&e).expect("degree-2");
1371        assert_eq!(q.quadratic().get(&(0, 0)), Some(&big));
1372        assert!(q.inexact(), "fl(2⁵³ + 1) = 2⁵³ was called exact");
1373        // Nothing is missing from the maps *yet*, so the consumers keep
1374        // their fast paths: the degree is not in question here.
1375        assert!(!q.lost_terms());
1376    }
1377
1378    /// The sticky half of the same rule. The rounding and the drop are in
1379    /// different subexpressions — `(2⁵³·x0 + x0)` absorbs, `− 2⁵³·x0`
1380    /// cancels — and the verdict has to survive the trip between them.
1381    #[test]
1382    fn rounding_carried_from_a_subexpression_makes_a_later_drop_a_loss() {
1383        let big = (1u64 << 53) as f64;
1384        let scaled =
1385            |c: f64| Expr::Binary(BinOp::Mul, Box::new(Expr::Const(c)), Box::new(Expr::Var(0)));
1386        let absorbed = Expr::Binary(BinOp::Add, Box::new(scaled(big)), Box::new(Expr::Var(0)));
1387        let e = Expr::Binary(BinOp::Sub, Box::new(absorbed), Box::new(scaled(big)));
1388        let q = recognize_expr(&e).expect("degree-1 at worst");
1389        assert!(q.linear().is_empty());
1390        assert!(
1391            q.lost_terms(),
1392            "the x0 absorbed two additions ago is no less missing now",
1393        );
1394    }
1395
1396    /// Cancellation *inside a product*: `(x0 + x1)·(x0 − x1)` accumulates
1397    /// `+x0x1` and `−x0x1` onto one key and they cancel exactly. The body
1398    /// really has no cross term, so the form must not be demoted for
1399    /// noticing.
1400    #[test]
1401    fn an_exact_cancellation_in_a_product_is_not_a_lost_term() {
1402        let add = |a: Expr, b: Expr| Expr::Binary(BinOp::Add, Box::new(a), Box::new(b));
1403        let sub = |a: Expr, b: Expr| Expr::Binary(BinOp::Sub, Box::new(a), Box::new(b));
1404        let e = Expr::Binary(
1405            BinOp::Mul,
1406            Box::new(add(Expr::Var(0), Expr::Var(1))),
1407            Box::new(sub(Expr::Var(0), Expr::Var(1))),
1408        );
1409        let q = recognize_expr(&e).expect("degree-2");
1410        assert_eq!(q.quadratic().get(&(0, 0)), Some(&1.0));
1411        assert_eq!(q.quadratic().get(&(1, 1)), Some(&-1.0));
1412        assert_eq!(q.quadratic().get(&(0, 1)), None);
1413        assert!(!q.lost_terms(), "x0² − x1² was reported incomplete");
1414    }
1415
1416    /// An underflowing **multiply** has no exact case to spare: both
1417    /// factors are nonzero and the product is a monomial the form can no
1418    /// longer see. `(10⁻²⁰⁰·x0)·(10⁻²⁰⁰·x0)` stays refused, which is gh
1419    /// #683's second reproduction and gh #687's second acceptance case.
1420    #[test]
1421    fn an_underflowing_product_is_still_a_lost_term() {
1422        let tiny = |i: usize| {
1423            Expr::Binary(
1424                BinOp::Mul,
1425                Box::new(Expr::Const(1e-200)),
1426                Box::new(Expr::Var(i)),
1427            )
1428        };
1429        let e = Expr::Binary(BinOp::Mul, Box::new(tiny(0)), Box::new(tiny(0)));
1430        let q = recognize_expr(&e).expect("degree-2 at worst");
1431        assert!(q.quadratic().is_empty());
1432        assert!(
1433            q.lost_terms(),
1434            "a coefficient that underflowed on the multiply was reported absent",
1435        );
1436    }
1437
1438    /// A constant folded away exactly is degree 0 for real — and `mul`
1439    /// reads a zero constant as annihilating, so this is the one drop whose
1440    /// verdict a *product* depends on.
1441    #[test]
1442    fn an_exactly_cancelled_constant_is_not_a_lost_term() {
1443        let e = Expr::Binary(
1444            BinOp::Sub,
1445            Box::new(Expr::Const(3.0)),
1446            Box::new(Expr::Const(3.0)),
1447        );
1448        let q = recognize_expr(&e).expect("degree 0");
1449        assert_eq!(q.as_constant(), Some(0.0));
1450        assert!(!q.lost_terms());
1451
1452        // The same shape one ulp off: `fl(2⁵³ + 1) − 2⁵³` is `0.0` where the
1453        // real value is `1`.
1454        let big = (1u64 << 53) as f64;
1455        let absorbed = Expr::Binary(
1456            BinOp::Add,
1457            Box::new(Expr::Const(big)),
1458            Box::new(Expr::Const(1.0)),
1459        );
1460        let e = Expr::Binary(BinOp::Sub, Box::new(absorbed), Box::new(Expr::Const(big)));
1461        let q = recognize_expr(&e).expect("degree 0");
1462        assert_eq!(q.as_constant(), Some(0.0));
1463        assert!(q.lost_terms(), "an absorbed constant was reported absent");
1464    }
1465
1466    /// Rounding a coefficient is recorded even when nothing is dropped,
1467    /// because that is what a later exact cancellation would be hiding:
1468    /// `x0²/3` cannot be represented, and the form has to remember it.
1469    #[test]
1470    fn a_rounded_scale_is_inexact_but_loses_nothing() {
1471        let e = Expr::Binary(BinOp::Div, Box::new(sq(0)), Box::new(Expr::Const(3.0)));
1472        let q = recognize_expr(&e).expect("degree-2");
1473        assert_eq!(q.quadratic().get(&(0, 0)), Some(&(1.0 / 3.0)));
1474        assert!(q.inexact(), "1 · (1/3) was called exact");
1475        assert!(
1476            !q.lost_terms(),
1477            "a rounded coefficient is not a missing one"
1478        );
1479
1480        // Halving is exact, so the same shape by a power of two is not.
1481        let e = Expr::Binary(BinOp::Div, Box::new(sq(0)), Box::new(Expr::Const(2.0)));
1482        let q = recognize_expr(&e).expect("degree-2");
1483        assert!(!q.inexact(), "x0²/2 rounds nothing");
1484    }
1485
1486    /// The two-sum and the two-product, against the cases the recognizer
1487    /// actually meets. Exactness is not a tolerance, so these are `==`.
1488    #[test]
1489    fn the_exactness_tests_agree_with_the_arithmetic() {
1490        let big = (1u64 << 53) as f64;
1491        assert!(add_is_exact(big, -big, big + -big));
1492        assert!(add_is_exact(1.0, 2.0, 3.0));
1493        assert!(!add_is_exact(big, 1.0, big + 1.0), "2⁵³ + 1 loses the 1");
1494        assert!(!add_is_exact(0.1, 0.2, 0.1 + 0.2));
1495        assert!(!add_is_exact(f64::INFINITY, 1.0, f64::INFINITY));
1496        assert!(!add_is_exact(f64::NAN, 1.0, f64::NAN + 1.0));
1497
1498        assert!(mul_is_exact(3.0, 1.0, 3.0), "the ±1 shortcut");
1499        assert!(mul_is_exact(3.0, 0.5, 1.5));
1500        assert!(mul_is_exact(0.1, 4.0, 0.4));
1501        assert!(!mul_is_exact(3.0, 1.0 / 3.0, 3.0 * (1.0 / 3.0)));
1502        assert!(!mul_is_exact(1e-200, 1e-200, 1e-200 * 1e-200), "underflow");
1503        assert!(!mul_is_exact(1e300, 1e300, 1e300 * 1e300), "overflow");
1504    }
1505
1506    #[test]
1507    fn cse_bodies_are_inlined_at_every_reference() {
1508        // c = x0; c · c is x0².
1509        let body = std::sync::Arc::new(Expr::Var(0));
1510        let e = Expr::Binary(
1511            BinOp::Mul,
1512            Box::new(Expr::Cse(body.clone())),
1513            Box::new(Expr::Cse(body)),
1514        );
1515        let h = analyze_quadratic(&e).expect("degree-2");
1516        assert_eq!(h.get(&(0, 0)), Some(&2.0));
1517    }
1518
1519    #[test]
1520    fn a_wide_nary_sum_does_not_recurse() {
1521        // Σ xᵢ² over 5000 terms as one `o54` node. Distinct keys, so this
1522        // says nothing about *order* — see the test below for that.
1523        const N: usize = 5000;
1524        let e = Expr::Sum((0..N).map(sq).collect());
1525        let h = analyze_quadratic(&e).expect("sum of squares is a QP");
1526        assert_eq!(h.len(), N);
1527        assert_eq!(h.get(&(N - 1, N - 1)), Some(&2.0));
1528    }
1529
1530    /// Summation order is only observable when two items land on the *same*
1531    /// monomial key with magnitudes far enough apart that the addition
1532    /// rounds. `1e16·x0² + x0² + x0²` folds to `1e16` front to back (each
1533    /// `+1` falls below the ulp) and to `1e16 + 2` back to front. The tape
1534    /// — the reference every non-recognized row is still evaluated with —
1535    /// sums front to back, so that is the answer required here.
1536    ///
1537    /// The 5000-square test above is named for this property but cannot see
1538    /// it: 5000 distinct keys never re-add anything.
1539    #[test]
1540    fn a_repeated_monomial_in_an_nary_sum_accumulates_front_to_back() {
1541        let scaled =
1542            |c: f64, i: usize| Expr::Binary(BinOp::Mul, Box::new(Expr::Const(c)), Box::new(sq(i)));
1543        let e = Expr::Sum(vec![scaled(1.0e16, 0), scaled(1.0, 0), scaled(1.0, 0)]);
1544        let h = analyze_quadratic(&e).expect("sum of squares is a QP");
1545        let got = h[&(0, 0)];
1546        assert_eq!(
1547            got.to_bits(),
1548            (2.0 * 1.0e16_f64).to_bits(),
1549            "expected the front-to-back fold, got {got:e}"
1550        );
1551        assert_ne!(got.to_bits(), (2.0 * (1.0e16_f64 + 2.0)).to_bits());
1552    }
1553
1554    /// The gate that decides whether a form may be evaluated from its
1555    /// coefficients. The two directions cost different things, so both are
1556    /// pinned: admitting a factored form loses digits (and, on
1557    /// `airport.nl`, 284 iterations), refusing an expanded one loses the
1558    /// phase's whole point on the family it was built for.
1559    #[test]
1560    fn only_already_expanded_forms_are_admitted() {
1561        // AMPL's `qcqp*` emission: `o54` over `0.5·((c·xᵢ)·xⱼ)`.
1562        let monomial = |c: f64, i: usize, j: usize| {
1563            Expr::Binary(
1564                BinOp::Mul,
1565                Box::new(Expr::Const(0.5)),
1566                Box::new(Expr::Binary(
1567                    BinOp::Mul,
1568                    Box::new(Expr::Binary(
1569                        BinOp::Mul,
1570                        Box::new(Expr::Const(c)),
1571                        Box::new(Expr::Var(i)),
1572                    )),
1573                    Box::new(Expr::Var(j)),
1574                )),
1575            )
1576        };
1577        let row = Expr::Sum(vec![monomial(2.0, 0, 0), monomial(3.0, 0, 1)]);
1578        assert!(is_expanded_quadratic(&row));
1579        // `xᵢ²` is a single monomial, not an expansion.
1580        assert!(is_expanded_quadratic(&Expr::Sum(vec![sq(0), sq(1)])));
1581        // A left-deep `Add` chain is still a sum.
1582        let chain = Expr::Binary(BinOp::Add, Box::new(sq(0)), Box::new(sq(1)));
1583        assert!(is_expanded_quadratic(&chain));
1584        // Division by a constant, and a negated term.
1585        assert!(is_expanded_quadratic(&Expr::Binary(
1586            BinOp::Div,
1587            Box::new(sq(0)),
1588            Box::new(Expr::Const(3.0)),
1589        )));
1590        assert!(is_expanded_quadratic(&Expr::Unary(
1591            UnaryOp::Neg,
1592            Box::new(monomial(1.0, 0, 1))
1593        )));
1594
1595        // `(x₀ − x₁)²` — the `airport.nl` shape. Reading it as
1596        // `x₀² − 2x₀x₁ + x₁²` cancels, so it stays on the tape.
1597        let diff = Expr::Binary(BinOp::Sub, Box::new(Expr::Var(0)), Box::new(Expr::Var(1)));
1598        let factored = Expr::Binary(BinOp::Pow, Box::new(diff), Box::new(Expr::Const(2.0)));
1599        assert!(analyze_quadratic(&factored).is_some(), "it is quadratic");
1600        assert!(
1601            !is_expanded_quadratic(&factored),
1602            "but not already expanded"
1603        );
1604
1605        // `(x₀ + 1)·x₁` — expansion by multiplication rather than by power.
1606        let product = Expr::Binary(
1607            BinOp::Mul,
1608            Box::new(Expr::Binary(
1609                BinOp::Add,
1610                Box::new(Expr::Var(0)),
1611                Box::new(Expr::Const(1.0)),
1612            )),
1613            Box::new(Expr::Var(1)),
1614        );
1615        assert!(!is_expanded_quadratic(&product));
1616
1617        // One factored term anywhere in a long expanded sum disqualifies
1618        // the whole row — the cancellation is in that term, not in the sum.
1619        let mixed = Expr::Sum(vec![monomial(1.0, 0, 0), factored]);
1620        assert!(!is_expanded_quadratic(&mixed));
1621
1622        // A `Cse` referenced once is inlined and judged on its body.
1623        let once = Expr::Binary(
1624            BinOp::Mul,
1625            Box::new(Expr::Cse(std::sync::Arc::new(Expr::Var(0)))),
1626            Box::new(Expr::Var(1)),
1627        );
1628        assert!(is_expanded_quadratic(&once));
1629    }
1630
1631    /// A `Cse` body reached twice is walked once, not `2^depth` times, and
1632    /// is **admitted** rather than refused.
1633    ///
1634    /// Q4 refused it, and refused it for cost rather than algebra — this
1635    /// walk and `recognize_expr` behind it both inline a body per
1636    /// reference, so a shared DAG was `Θ(2^depth)`. Q5's memoization is
1637    /// what makes admitting it affordable, so both halves are asserted
1638    /// here: the verdict, and the fact that the test returns at all. The
1639    /// shape is `nl_reader`'s `shared_dag_walks_are_memoized_not_exponential`
1640    /// scaled up; at depth 60 an exponential walk does not finish inside
1641    /// the lifetime of this test run, and `2^60` sums do not fit in an
1642    /// `f64` count either.
1643    #[test]
1644    fn a_shared_cse_body_is_walked_once_and_admitted() {
1645        let mut e = Expr::Var(0);
1646        for _ in 0..60 {
1647            let shared = std::sync::Arc::new(e);
1648            e = Expr::Binary(
1649                BinOp::Add,
1650                Box::new(Expr::Cse(std::sync::Arc::clone(&shared))),
1651                Box::new(Expr::Cse(shared)),
1652            );
1653        }
1654        assert!(is_expanded_quadratic(&e));
1655        // And the algebra agrees: `x + x` sixty times over is `2^60 · x`.
1656        let q = recognize_expr(&e).expect("a sum of one monomial is degree 1");
1657        assert_eq!(q.linear().get(&0).copied(), Some(2.0_f64.powi(60)));
1658        std::mem::forget(e);
1659    }
1660
1661    /// The same shape inside a *monomial*, which is a different question
1662    /// with a different answer and therefore a separately keyed memo: a
1663    /// body that is a sum of monomials is legal on the spine and illegal
1664    /// under a `*`.
1665    #[test]
1666    fn a_shared_body_is_judged_per_context_not_once_and_for_all() {
1667        let sum = std::sync::Arc::new(Expr::Binary(
1668            BinOp::Add,
1669            Box::new(Expr::Var(0)),
1670            Box::new(Expr::Var(1)),
1671        ));
1672        // On the spine, twice: fine.
1673        let spine = Expr::Binary(
1674            BinOp::Add,
1675            Box::new(Expr::Cse(std::sync::Arc::clone(&sum))),
1676            Box::new(Expr::Cse(std::sync::Arc::clone(&sum))),
1677        );
1678        assert!(is_expanded_quadratic(&spine));
1679        // The same body on the spine and then under a product: the product
1680        // is `(x0 + x1) · x1`, a factored form, and the earlier clean visit
1681        // on the spine must not clear it.
1682        let mixed = Expr::Binary(
1683            BinOp::Add,
1684            Box::new(Expr::Cse(std::sync::Arc::clone(&sum))),
1685            Box::new(Expr::Binary(
1686                BinOp::Mul,
1687                Box::new(Expr::Cse(sum)),
1688                Box::new(Expr::Var(1)),
1689            )),
1690        );
1691        assert!(!is_expanded_quadratic(&mixed));
1692    }
1693
1694    /// Deep, and on a default-sized test thread, for the same reason
1695    /// [`recognize_expr`] is iterative: the gate runs on every row of every
1696    /// model, including the ones that overflow a recursive walk.
1697    #[test]
1698    fn the_expansion_gate_does_not_overflow_the_stack() {
1699        const K: usize = 250_000;
1700        let mut e = sq(0);
1701        for i in 1..K {
1702            e = Expr::Binary(BinOp::Add, Box::new(e), Box::new(sq(i)));
1703        }
1704        assert!(is_expanded_quadratic(&e));
1705        std::mem::forget(e);
1706    }
1707
1708    /// The reason this module is iterative.
1709    ///
1710    /// A `.nl` writer that emits `o0` (binary `+`) chains for a long sum
1711    /// hands the recognizer a left-deep `Add` tree one level deep per term.
1712    /// The recursive predecessor aborted the process — a stack overflow is
1713    /// not a catchable error — somewhere under 8 000 terms on a 2 MB thread
1714    /// and under 24 000 on the CLI's 8 MB main thread. The depth below is
1715    /// far past both, and the test runs on a **default-sized test thread**
1716    /// deliberately: what a recognizer survives must not depend on who
1717    /// called it.
1718    ///
1719    /// The tree is leaked rather than dropped, because `Expr`'s derived
1720    /// `Drop` is still recursive and would overflow tearing this down.
1721    /// That is a real and separate defect (pounce#472 works around it in
1722    /// the Python bindings with a big-stack worker thread); it is not what
1723    /// this test is measuring.
1724    #[test]
1725    fn deep_add_chain_does_not_overflow_the_stack() {
1726        const K: usize = 250_000;
1727        let mut e = sq(0);
1728        for i in 1..K {
1729            e = Expr::Binary(BinOp::Add, Box::new(e), Box::new(sq(i)));
1730        }
1731        let h = analyze_quadratic(&e).expect("a sum of squares is a QP at any depth");
1732        assert_eq!(h.len(), K, "every xᵢ² contributes one diagonal entry");
1733        assert_eq!(h.get(&(K - 1, K - 1)), Some(&2.0));
1734        std::mem::forget(e);
1735    }
1736
1737    /// Same shape, right-deep — the value stack, not the work stack, is
1738    /// what grows here. Both live on the heap.
1739    #[test]
1740    fn deep_right_leaning_chain_does_not_overflow_the_stack() {
1741        const K: usize = 250_000;
1742        let mut e = sq(K - 1);
1743        for i in (0..K - 1).rev() {
1744            e = Expr::Binary(BinOp::Add, Box::new(sq(i)), Box::new(e));
1745        }
1746        let h = analyze_quadratic(&e).expect("a sum of squares is a QP at any depth");
1747        assert_eq!(h.len(), K);
1748        std::mem::forget(e);
1749    }
1750
1751    /// A non-quadratic node deep inside a deep tree must return `None`
1752    /// rather than unwind — the bail-out path drops the work and value
1753    /// stacks, and neither is recursive.
1754    #[test]
1755    fn deep_chain_with_a_transcendental_bails_without_overflowing() {
1756        const K: usize = 250_000;
1757        let mut e = Expr::Unary(UnaryOp::Sin, Box::new(Expr::Var(0)));
1758        for i in 1..K {
1759            e = Expr::Binary(BinOp::Add, Box::new(e), Box::new(sq(i)));
1760        }
1761        assert!(analyze_quadratic(&e).is_none());
1762        std::mem::forget(e);
1763    }
1764
1765    // -----------------------------------------------------------------
1766    // Factored forms (gh #673)
1767    // -----------------------------------------------------------------
1768
1769    /// `(base)^2`.
1770    fn sq_of(base: Expr) -> Expr {
1771        Expr::Binary(BinOp::Pow, Box::new(base), Box::new(Expr::Const(2.0)))
1772    }
1773
1774    fn var_minus(i: usize, c: f64) -> Expr {
1775        Expr::Binary(BinOp::Sub, Box::new(Expr::Var(i)), Box::new(Expr::Const(c)))
1776    }
1777
1778    /// The motivating case. `(x − 500000)²` is exactly what
1779    /// `feasible_x0_extreme_row.nl` writes, and expanding it is the 2.4e-5
1780    /// disagreement gh #673 is named after.
1781    #[test]
1782    fn a_shifted_square_is_kept_factored() {
1783        let e = sq_of(var_minus(0, 500_000.0));
1784        assert!(
1785            !is_expanded_quadratic(&e),
1786            "the expanded gate must refuse it"
1787        );
1788        let f = recognize_factored_quadratic(&e).expect("a square of an affine form");
1789        assert_eq!(f.squares.len(), 1);
1790        assert_eq!(f.squares[0].weight, 1.0);
1791        assert_eq!(f.squares[0].coefs, vec![(0, 1.0)]);
1792        assert_eq!(f.squares[0].constant, -500_000.0);
1793        assert!(f.linear.is_empty());
1794        assert_eq!(f.constant, 0.0);
1795    }
1796
1797    /// The accuracy claim, stated as a number rather than as prose: at
1798    /// `x = 500000 + 1e-4` the expansion loses five digits and the factored
1799    /// read-out loses nothing.
1800    #[test]
1801    fn the_factored_read_out_does_not_cancel_where_the_expansion_does() {
1802        let e = sq_of(var_minus(0, 500_000.0));
1803        let x = 500_000.0 + 1e-4;
1804        // What the tape computes: square the residual as written.
1805        let r = x - 500_000.0;
1806        let taped = r * r;
1807
1808        let (h, lin, c) = analyze_quadratic_full(&e).expect("degree 2");
1809        let expanded = 0.5 * h[&(0, 0)] * x * x + lin[0].1 * x + c;
1810        assert!(
1811            (expanded - taped).abs() / taped > 1e-6,
1812            "the expansion is supposed to cancel here, got {expanded} for {taped}"
1813        );
1814
1815        let f = recognize_factored_quadratic(&e).expect("a square");
1816        let t = &f.squares[0];
1817        let l = t.constant + t.coefs[0].1 * x;
1818        // Bit for bit, not within a tolerance.
1819        assert_eq!(t.weight * l * l, taped);
1820    }
1821
1822    /// `airport.nl`'s shape: a row that is a sum of squared coordinate
1823    /// differences, with the writer's grouping kept term by term.
1824    #[test]
1825    fn a_sum_of_squared_differences_is_admitted() {
1826        let diff = |i: usize, j: usize| {
1827            Expr::Binary(BinOp::Sub, Box::new(Expr::Var(i)), Box::new(Expr::Var(j)))
1828        };
1829        let e = Expr::Binary(
1830            BinOp::Add,
1831            Box::new(sq_of(diff(0, 1))),
1832            Box::new(sq_of(diff(2, 3))),
1833        );
1834        let f = recognize_factored_quadratic(&e).expect("two squares");
1835        assert_eq!(f.squares.len(), 2);
1836        let mut sup: Vec<Vec<(usize, f64)>> = f.squares.iter().map(|t| t.coefs.clone()).collect();
1837        sup.sort_by_key(|c| c[0].0);
1838        assert_eq!(
1839            sup,
1840            vec![vec![(0, 1.0), (1, -1.0)], vec![(2, 1.0), (3, -1.0)]]
1841        );
1842        assert!(
1843            f.squares
1844                .iter()
1845                .all(|t| t.weight == 1.0 && t.constant == 0.0)
1846        );
1847    }
1848
1849    /// The sign of the sum spine reaches the weight, and a constant factor
1850    /// folds into it.
1851    #[test]
1852    fn spine_signs_and_constant_factors_fold_into_the_weight() {
1853        // 3·(x₀ − 1)² − 0.5·(x₁ + 2)²
1854        let a = Expr::Binary(
1855            BinOp::Mul,
1856            Box::new(Expr::Const(3.0)),
1857            Box::new(sq_of(var_minus(0, 1.0))),
1858        );
1859        let b = Expr::Binary(
1860            BinOp::Mul,
1861            Box::new(Expr::Const(0.5)),
1862            Box::new(sq_of(Expr::Binary(
1863                BinOp::Add,
1864                Box::new(Expr::Var(1)),
1865                Box::new(Expr::Const(2.0)),
1866            ))),
1867        );
1868        let f = recognize_factored_quadratic(&Expr::Binary(BinOp::Sub, Box::new(a), Box::new(b)))
1869            .expect("two squares");
1870        let mut got: Vec<(f64, f64)> = f.squares.iter().map(|t| (t.weight, t.constant)).collect();
1871        got.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
1872        assert_eq!(got, vec![(-0.5, 2.0), (3.0, -1.0)]);
1873    }
1874
1875    /// A degree-≤1 leftover in the same tree is not a reason to refuse the
1876    /// row; it folds into `a`/`c` and is evaluated there.
1877    #[test]
1878    fn degree_one_leftovers_fold_into_the_linear_part() {
1879        // (x₀ − 1)² + 3·x₁ + 7
1880        let e = Expr::Sum(vec![
1881            sq_of(var_minus(0, 1.0)),
1882            Expr::Binary(
1883                BinOp::Mul,
1884                Box::new(Expr::Const(3.0)),
1885                Box::new(Expr::Var(1)),
1886            ),
1887            Expr::Const(7.0),
1888        ]);
1889        let f = recognize_factored_quadratic(&e).expect("a square plus leftovers");
1890        assert_eq!(f.squares.len(), 1);
1891        assert_eq!(f.linear, vec![(1, 3.0)]);
1892        assert_eq!(f.constant, 7.0);
1893    }
1894
1895    /// A bare `c·xᵢ²` monomial is the same term wearing a different opcode,
1896    /// and is stored as `c·(xᵢ)²` so that a row mixing the two shapes is
1897    /// still served.
1898    #[test]
1899    fn a_diagonal_monomial_is_stored_as_a_square() {
1900        let e = Expr::Binary(
1901            BinOp::Add,
1902            Box::new(sq_of(var_minus(0, 1.0))),
1903            Box::new(Expr::Binary(
1904                BinOp::Mul,
1905                Box::new(Expr::Const(4.0)),
1906                Box::new(Expr::Binary(
1907                    BinOp::Mul,
1908                    Box::new(Expr::Var(1)),
1909                    Box::new(Expr::Var(1)),
1910                )),
1911            )),
1912        );
1913        let f = recognize_factored_quadratic(&e).expect("square + diagonal monomial");
1914        assert_eq!(f.squares.len(), 2);
1915        let diag = f
1916            .squares
1917            .iter()
1918            .find(|t| t.coefs == vec![(1, 1.0)])
1919            .unwrap();
1920        // The *polynomial* coefficient, so `4x₁²` and not `8x₁²` —
1921        // `push_factored_form` is what doubles it into a Hessian entry.
1922        assert_eq!((diag.weight, diag.constant), (4.0, 0.0));
1923    }
1924
1925    /// A cross monomial has no square to be stored as, and storing it
1926    /// expanded next to the squares is the mixed representation this
1927    /// deliberately does not have. The row keeps its tape.
1928    #[test]
1929    fn a_cross_monomial_alongside_a_square_is_refused() {
1930        let e = Expr::Binary(
1931            BinOp::Add,
1932            Box::new(sq_of(var_minus(0, 1.0))),
1933            Box::new(Expr::Binary(
1934                BinOp::Mul,
1935                Box::new(Expr::Var(1)),
1936                Box::new(Expr::Var(2)),
1937            )),
1938        );
1939        assert!(recognize_factored_quadratic(&e).is_none());
1940    }
1941
1942    /// An already-expanded body is `is_expanded_quadratic`'s to serve, and
1943    /// this must not offer a second, slower answer for it.
1944    #[test]
1945    fn an_expanded_body_is_not_claimed_here() {
1946        let e = Expr::Binary(BinOp::Add, Box::new(sq(0)), Box::new(sq(1)));
1947        assert!(is_expanded_quadratic(&e));
1948        // `x²` *is* a square, so this one is genuinely recognized — the
1949        // caller's ordering is what keeps it on the cheaper path.
1950        assert!(recognize_factored_quadratic(&e).is_some());
1951        // A body with no square at all is refused outright.
1952        let cross = Expr::Binary(BinOp::Mul, Box::new(Expr::Var(0)), Box::new(Expr::Var(1)));
1953        assert!(recognize_factored_quadratic(&cross).is_none());
1954        assert!(recognize_factored_quadratic(&Expr::Var(0)).is_none());
1955    }
1956
1957    /// The product of two *different* affine forms is degree 2 and is not a
1958    /// square. Admitting it would mean storing `l₁·l₂`, which this
1959    /// representation cannot express.
1960    #[test]
1961    fn a_product_of_two_different_affine_forms_is_refused() {
1962        let e = Expr::Binary(
1963            BinOp::Mul,
1964            Box::new(var_minus(0, 1.0)),
1965            Box::new(var_minus(1, 2.0)),
1966        );
1967        assert!(recognize_factored_quadratic(&e).is_none());
1968    }
1969
1970    /// A square of a *quadratic* is degree 4, and a transcendental is not a
1971    /// polynomial at all.
1972    #[test]
1973    fn quartics_and_transcendentals_are_refused() {
1974        assert!(recognize_factored_quadratic(&sq_of(sq(0))).is_none());
1975        let s = Expr::Unary(UnaryOp::Sin, Box::new(Expr::Var(0)));
1976        assert!(recognize_factored_quadratic(&sq_of(s)).is_none());
1977    }
1978
1979    /// A base that is not itself a flat sum of monomials is refused: the
1980    /// coefficients read out of `2·(x + 1)` are not the additions the
1981    /// writer wrote, which is the same rule `is_expanded_quadratic` states.
1982    #[test]
1983    fn a_base_the_expanded_gate_refuses_is_refused_here_too() {
1984        let inner = Expr::Binary(
1985            BinOp::Mul,
1986            Box::new(Expr::Const(2.0)),
1987            Box::new(Expr::Binary(
1988                BinOp::Add,
1989                Box::new(Expr::Var(0)),
1990                Box::new(Expr::Const(1.0)),
1991            )),
1992        );
1993        assert!(!is_expanded_quadratic(&inner));
1994        assert!(recognize_factored_quadratic(&sq_of(inner)).is_none());
1995    }
1996
1997    /// A leftover that lost a term inexactly refuses the body, for the
1998    /// reason gh #685 gives: the form would evaluate a row its own tape
1999    /// does not agree with.
2000    #[test]
2001    fn a_lost_leftover_refuses_the_body() {
2002        let big = 9_007_199_254_740_992.0f64; // 2⁵³
2003        let scaled = |c: f64, i: usize| {
2004            Expr::Binary(BinOp::Mul, Box::new(Expr::Const(c)), Box::new(Expr::Var(i)))
2005        };
2006        // (x₀ − 1)² + 2⁵³·x₁ + x₁ − 2⁵³·x₁ — the `x₁` is folded away.
2007        let e = Expr::Sum(vec![
2008            sq_of(var_minus(0, 1.0)),
2009            scaled(big, 1),
2010            Expr::Var(1),
2011            scaled(-big, 1),
2012        ]);
2013        assert!(recognize_factored_quadratic(&e).is_none());
2014    }
2015
2016    /// A `Cse` on the sum spine ends the walk: skipping a second reference
2017    /// would drop its terms and visiting every reference is exponential.
2018    #[test]
2019    fn a_shared_body_on_the_spine_is_refused() {
2020        let shared = std::sync::Arc::new(sq_of(var_minus(0, 1.0)));
2021        let e = Expr::Binary(
2022            BinOp::Add,
2023            Box::new(Expr::Cse(shared.clone())),
2024            Box::new(Expr::Cse(shared)),
2025        );
2026        assert!(recognize_factored_quadratic(&e).is_none());
2027    }
2028
2029    /// A weight that annihilates or overflows is not evaluated from stored
2030    /// coefficients: `0·(x−1)²` has dropped a term the tape still walks,
2031    /// and `(x−1)²/0` is not a number.
2032    #[test]
2033    fn degenerate_weights_are_refused() {
2034        let zero = Expr::Binary(
2035            BinOp::Mul,
2036            Box::new(Expr::Const(0.0)),
2037            Box::new(sq_of(var_minus(0, 1.0))),
2038        );
2039        assert!(recognize_factored_quadratic(&zero).is_none());
2040        let div0 = Expr::Binary(
2041            BinOp::Div,
2042            Box::new(sq_of(var_minus(0, 1.0))),
2043            Box::new(Expr::Const(0.0)),
2044        );
2045        assert!(recognize_factored_quadratic(&div0).is_none());
2046    }
2047
2048    /// Deep, on a default-sized test thread, for the reason every walk in
2049    /// this module is iterative: a least-squares model is a long `o0` chain
2050    /// of squared residuals, which is exactly the shape that aborts a
2051    /// recursive walk — and now exactly the shape this recognizer is for.
2052    ///
2053    /// Leaked rather than dropped: `Expr`'s derived `Drop` is still
2054    /// recursive. See `deep_add_chain_does_not_overflow_the_stack`.
2055    #[test]
2056    fn a_deep_chain_of_squares_does_not_overflow_the_stack() {
2057        const K: usize = 250_000;
2058        let mut e = sq_of(var_minus(0, 1.0));
2059        for i in 1..K {
2060            e = Expr::Binary(
2061                BinOp::Add,
2062                Box::new(e),
2063                Box::new(sq_of(var_minus(i, i as f64))),
2064            );
2065        }
2066        let f = recognize_factored_quadratic(&e).expect("K squares");
2067        assert_eq!(f.squares.len(), K);
2068        std::mem::forget(e);
2069    }
2070}