Skip to main content

alkahest_cas/kernel/
pool.rs

1use crate::kernel::{
2    domain::Domain,
3    expr::{BigFloat, BigInt, BigRat, ExprData, ExprId},
4};
5use std::fmt;
6
7/// Canonical ∞ symbol name for [`ExprPool::pos_infinity`] / limits (V2-16).
8pub const POS_INFINITY_SYMBOL: &str = "\u{221e}";
9
10// ---------------------------------------------------------------------------
11// Lock-free arena for ExprPool nodes.
12//
13// Strategy:
14//   * The `nodes` array (ExprId → ExprData) is a `boxcar::Vec` — a
15//     lock-free, append-only, reference-stable segmented array.  Reads
16//     (`with`, `get`, `len`) acquire no lock at all; they index directly
17//     into the array via a single atomic load.
18//   * The `index` (ExprData → ExprId) still requires coordination during
19//     insertion to preserve hash-cons uniqueness:
20//     - Under `--features parallel` we use `DashMap::entry` which holds a
21//       per-shard write-lock only for the duration of the insert.  The
22//       closure passed to `or_insert_with` calls `boxcar::push` (lock-free)
23//       while the shard lock is held, so no two threads can insert the same
24//       key.
25//     - Without `parallel` the `Mutex<HashMap>` serialises all inserts as
26//       before; the boxcar push happens while the Mutex is held.
27// ---------------------------------------------------------------------------
28
29#[cfg(feature = "parallel")]
30use dashmap::DashMap;
31
32#[cfg(not(feature = "parallel"))]
33use std::collections::HashMap;
34
35#[cfg(not(feature = "parallel"))]
36use std::sync::Mutex;
37
38// ---------------------------------------------------------------------------
39// PoolState — two variants depending on build features
40// ---------------------------------------------------------------------------
41
42#[cfg(feature = "parallel")]
43struct PoolIndex(DashMap<ExprData, ExprId>);
44
45#[cfg(not(feature = "parallel"))]
46struct PoolIndex(HashMap<ExprData, ExprId>);
47
48#[cfg(feature = "parallel")]
49impl PoolIndex {
50    fn new() -> Self {
51        PoolIndex(DashMap::new())
52    }
53    fn get(&self, data: &ExprData) -> Option<ExprId> {
54        self.0.get(data).map(|v| *v)
55    }
56    /// Atomically return the existing id for `key`, or call `f` to produce one
57    /// and insert it.  The DashMap shard write-lock is held for the duration of
58    /// `f`, guaranteeing at most one call to `f` per unique key.
59    fn or_insert_with(&self, key: ExprData, f: impl FnOnce() -> ExprId) -> ExprId {
60        *self.0.entry(key).or_insert_with(f)
61    }
62}
63
64#[cfg(not(feature = "parallel"))]
65impl PoolIndex {
66    fn new() -> Self {
67        PoolIndex(HashMap::new())
68    }
69    fn get(&self, data: &ExprData) -> Option<ExprId> {
70        self.0.get(data).copied()
71    }
72    fn insert(&mut self, data: ExprData, id: ExprId) {
73        self.0.insert(data, id);
74    }
75}
76
77/// Owns all expression nodes. Every [`ExprId`] is valid only within its pool.
78///
79/// `ExprPool` is `Send + Sync`.
80///
81/// Read operations (`with`, `get`, `len`) are fully lock-free — they index
82/// into a `boxcar::Vec` via a single atomic load with no lock acquisition.
83/// Write operations (`intern`) use a per-shard lock (parallel mode) or a
84/// `Mutex` (non-parallel mode) only during new-node insertion.
85/// A node plus the properties that are cheaper to record once than to recompute.
86struct Node {
87    data: ExprData,
88    /// Whether every generator in this subtree commutes under multiplication.
89    ///
90    /// This is a bottom-up property, and hash-consing guarantees a node's
91    /// children are interned before the node itself, so it is computed once
92    /// here from the children's cached flags — O(arity) — instead of by
93    /// walking the whole subtree on every query.
94    mult_commutative: bool,
95    /// Length of the longest root-to-leaf path in this subtree; a leaf is 1.
96    ///
97    /// Computed exactly like `mult_commutative` — once, at intern time, from
98    /// the children's cached values — so [`ExprPool::depth`] is a single array
99    /// read.  Recomputing it on demand is not an option: the pool is a DAG, so
100    /// an unmemoised depth walk is exponential in the sharing, and a memoised
101    /// one allocates a map per query.  Saturating, so a pathological expression
102    /// pins at `u32::MAX` instead of wrapping to a small value.
103    ///
104    /// This is what lets every recursive consumer refuse a too-deep expression
105    /// in O(1) rather than discovering the problem by overflowing the stack.
106    depth: u32,
107}
108
109pub struct ExprPool {
110    /// Lock-free, append-only, reference-stable node array.
111    nodes: boxcar::Vec<Node>,
112    /// Deduplication index: ExprData → ExprId.
113    #[cfg(feature = "parallel")]
114    index: PoolIndex,
115    #[cfg(not(feature = "parallel"))]
116    index: Mutex<PoolIndex>,
117}
118
119unsafe impl Send for ExprPool {}
120unsafe impl Sync for ExprPool {}
121
122impl ExprPool {
123    pub fn new() -> Self {
124        ExprPool {
125            nodes: boxcar::Vec::new(),
126            #[cfg(feature = "parallel")]
127            index: PoolIndex::new(),
128            #[cfg(not(feature = "parallel"))]
129            index: Mutex::new(PoolIndex::new()),
130        }
131    }
132
133    /// Intern `data`, returning a shared [`ExprId`]. Identical structures
134    /// always return the same id; structural equality ⟺ id equality.
135    pub fn intern(&self, data: ExprData) -> ExprId {
136        #[cfg(feature = "parallel")]
137        {
138            // Fast path: lock-free DashMap read.
139            if let Some(id) = self.index.get(&data) {
140                return id;
141            }
142            // Slow path: DashMap shard write-lock ensures at most one push
143            // per unique key.  `boxcar::push` is lock-free so it can be
144            // called safely while the shard lock is held.
145            self.index.or_insert_with(data.clone(), || {
146                let node = self.make_node(data);
147                ExprId(self.nodes.push(node) as u32)
148            })
149        }
150
151        #[cfg(not(feature = "parallel"))]
152        {
153            let mut idx = self.index.lock().expect("ExprPool index Mutex poisoned");
154            if let Some(id) = idx.get(&data) {
155                return id;
156            }
157            let node = self.make_node(data.clone());
158            let id = ExprId(self.nodes.push(node) as u32);
159            idx.insert(data, id);
160            id
161        }
162    }
163
164    /// Wrap `data` with its cached properties.  Children are already interned,
165    /// so their flags are just array reads.
166    fn make_node(&self, data: ExprData) -> Node {
167        let mult_commutative = self.compute_mult_commutative(&data);
168        let depth = self.compute_depth(&data);
169        Node {
170            data,
171            mult_commutative,
172            depth,
173        }
174    }
175
176    /// One level of the depth recurrence: `1 + max(child depths)`, reading each
177    /// child's cached depth rather than descending into it.
178    fn compute_depth(&self, data: &ExprData) -> u32 {
179        let child = |c: ExprId| self.depth(c);
180        let deepest = match data {
181            ExprData::Symbol { .. }
182            | ExprData::Integer(_)
183            | ExprData::Rational(_)
184            | ExprData::Float(_) => 0,
185            ExprData::Add(args) | ExprData::Mul(args) => {
186                args.iter().copied().map(child).max().unwrap_or(0)
187            }
188            ExprData::Pow { base, exp } => child(*base).max(child(*exp)),
189            ExprData::Func { args, .. } => args.iter().copied().map(child).max().unwrap_or(0),
190            ExprData::Piecewise { branches, default } => branches
191                .iter()
192                .map(|&(c, v)| child(c).max(child(v)))
193                .max()
194                .unwrap_or(0)
195                .max(child(*default)),
196            ExprData::Predicate { args, .. } => args.iter().copied().map(child).max().unwrap_or(0),
197            ExprData::Forall { var, body } | ExprData::Exists { var, body } => {
198                child(*var).max(child(*body))
199            }
200            ExprData::BigO(inner) => child(*inner),
201            ExprData::RootSum { poly, body, .. } => child(*poly).max(child(*body)),
202        };
203        deepest.saturating_add(1)
204    }
205
206    /// One level of the `mult_tree_is_commutative` recurrence, reading each
207    /// child's cached flag rather than descending into it.
208    fn compute_mult_commutative(&self, data: &ExprData) -> bool {
209        let child = |c: ExprId| self.is_mult_commutative(c);
210        match data {
211            ExprData::Symbol { commutative, .. } => *commutative,
212            ExprData::Integer(_) | ExprData::Rational(_) | ExprData::Float(_) => true,
213            ExprData::Add(args) | ExprData::Mul(args) => args.iter().copied().all(child),
214            ExprData::Pow { base, exp } => child(*base) && child(*exp),
215            ExprData::Func { args, .. } => args.iter().copied().all(child),
216            ExprData::Piecewise { branches, default } => {
217                branches.iter().all(|&(c, v)| child(c) && child(v)) && child(*default)
218            }
219            ExprData::Predicate { args, .. } => args.iter().copied().all(child),
220            ExprData::Forall { var, body } | ExprData::Exists { var, body } => {
221                child(*var) && child(*body)
222            }
223            ExprData::BigO(inner) => child(*inner),
224            ExprData::RootSum { poly, body, .. } => child(*poly) && child(*body),
225        }
226    }
227
228    /// Whether every generator in the subtree rooted at `id` commutes under
229    /// multiplication.  O(1): the flag was computed when `id` was interned.
230    pub fn is_mult_commutative(&self, id: ExprId) -> bool {
231        self.node(id).mult_commutative
232    }
233
234    /// Length of the longest root-to-leaf path in the subtree rooted at `id`.
235    ///
236    /// A leaf (symbol or number) has depth 1.  O(1): the value was computed
237    /// when `id` was interned.  Saturates at [`u32::MAX`].
238    ///
239    /// Every recursive consumer of an expression uses this to decline a tree
240    /// too deep for the stack — see
241    /// [`crate::kernel::depth::check_expr_depth`].
242    pub fn depth(&self, id: ExprId) -> u32 {
243        self.node(id).depth
244    }
245
246    fn node(&self, id: ExprId) -> &Node {
247        self.nodes
248            .get(id.0 as usize)
249            .expect("ExprPool: ExprId out of range")
250    }
251
252    /// Borrow a node by id and apply `f` without cloning.  Lock-free.
253    pub fn with<R, F: FnOnce(&ExprData) -> R>(&self, id: ExprId, f: F) -> R {
254        f(&self.node(id).data)
255    }
256
257    /// Clone and return the `ExprData` for `id`.
258    pub fn get(&self, id: ExprId) -> ExprData {
259        self.with(id, |d| d.clone())
260    }
261
262    /// Number of distinct expressions interned so far.  Lock-free.
263    pub fn len(&self) -> usize {
264        self.nodes.count()
265    }
266
267    pub fn is_empty(&self) -> bool {
268        self.nodes.is_empty()
269    }
270
271    // -----------------------------------------------------------------------
272    // Atom constructors
273    // -----------------------------------------------------------------------
274
275    /// Free symbol; multiplication treats it as commuting with every other factor (default).
276    pub fn symbol(&self, name: impl Into<String>, domain: Domain) -> ExprId {
277        self.symbol_commutative(name, domain, true)
278    }
279
280    /// Canonical name of the kernel-blessed imaginary unit `i = √(−1)`.
281    ///
282    /// Reserved: do not create an unrelated free symbol with this name and
283    /// `Domain::Complex` — the simplifier applies the algebraic power rules
284    /// `i² = −1`, `i³ = −i`, `i⁴ = 1`, … to any symbol matching this name and
285    /// domain (see [`ExprPool::is_imaginary_unit`]).
286    pub const IMAGINARY_UNIT_NAME: &'static str = "I";
287
288    /// The first-class imaginary unit `i = √(−1)`.
289    ///
290    /// Represented as the interned, kernel-blessed commuting symbol
291    /// [`IMAGINARY_UNIT_NAME`](Self::IMAGINARY_UNIT_NAME) with
292    /// [`Domain::Complex`]. This is the *canonical* representation: the
293    /// simplifier knows the algebraic identities `i² = −1`, `i³ = −i`,
294    /// `i⁴ = 1`, and more generally `i^(4k+r) → i^r` for literal integer
295    /// exponents (no branch-cut identities — `√(−1) → i`, `log`/`exp` of
296    /// complex arguments etc. are *not* added).
297    ///
298    /// Differentiation treats it as a constant (`d/dx i = 0`, like `π`/`e`)
299    /// and numeric evaluation declines (it has no `f64` value), matching the
300    /// behaviour of other non-real atoms.
301    pub fn imaginary_unit(&self) -> ExprId {
302        self.symbol(Self::IMAGINARY_UNIT_NAME, Domain::Complex)
303    }
304
305    /// Returns `true` iff `id` is the canonical imaginary unit produced by
306    /// [`ExprPool::imaginary_unit`] (an interned `Domain::Complex` symbol named
307    /// [`IMAGINARY_UNIT_NAME`](Self::IMAGINARY_UNIT_NAME)).
308    pub fn is_imaginary_unit(&self, id: ExprId) -> bool {
309        self.with(id, |d| {
310            matches!(
311                d,
312                ExprData::Symbol { name, domain, .. }
313                    if name == Self::IMAGINARY_UNIT_NAME && *domain == Domain::Complex
314            )
315        })
316    }
317
318    /// Free symbol with explicit commutative flag (V3-2). `commutative: false` is for
319    /// matrix or operator generators where `A*B` and `B*A` must remain distinct.
320    pub fn symbol_commutative(
321        &self,
322        name: impl Into<String>,
323        domain: Domain,
324        commutative: bool,
325    ) -> ExprId {
326        self.intern(ExprData::Symbol {
327            name: name.into(),
328            domain,
329            commutative,
330        })
331    }
332
333    pub fn integer(&self, n: impl Into<rug::Integer>) -> ExprId {
334        self.intern(ExprData::Integer(BigInt(n.into())))
335    }
336
337    pub fn rational(
338        &self,
339        numer: impl Into<rug::Integer>,
340        denom: impl Into<rug::Integer>,
341    ) -> ExprId {
342        let r = rug::Rational::from((numer.into(), denom.into()));
343        self.intern(ExprData::Rational(BigRat(r)))
344    }
345
346    pub fn float(&self, value: f64, prec: u32) -> ExprId {
347        let f = rug::Float::with_val(prec, value);
348        self.intern(ExprData::Float(BigFloat { inner: f, prec }))
349    }
350
351    // -----------------------------------------------------------------------
352    // Compound constructors
353    // -----------------------------------------------------------------------
354
355    pub fn add(&self, mut args: Vec<ExprId>) -> ExprId {
356        // Sort children at construction time so that commutativity holds
357        // structurally: `a + b` and `b + a` intern to the same ExprId.
358        // The sort key is the raw ExprId (opaque u32), which gives a stable,
359        // deterministic canonical order.
360        args.sort_unstable();
361        self.intern(ExprData::Add(args))
362    }
363
364    pub fn mul(&self, mut args: Vec<ExprId>) -> ExprId {
365        // Canonical sort only when every subtree is multiplicatively commutative (V3-2).
366        let sort_ok = args
367            .iter()
368            .all(|&a| crate::kernel::expr_props::mult_tree_is_commutative(self, a));
369        if sort_ok {
370            args.sort_unstable();
371        }
372        self.intern(ExprData::Mul(args))
373    }
374
375    pub fn pow(&self, base: ExprId, exp: ExprId) -> ExprId {
376        self.intern(ExprData::Pow { base, exp })
377    }
378
379    pub fn func(&self, name: impl Into<String>, args: Vec<ExprId>) -> ExprId {
380        self.intern(ExprData::Func {
381            name: name.into(),
382            args,
383        })
384    }
385
386    // -----------------------------------------------------------------------
387    // PA-9 — Piecewise / Predicate constructors
388    // -----------------------------------------------------------------------
389
390    /// Build a `Piecewise` expression.
391    ///
392    /// Branches are `(cond, value)` pairs where `cond` must be a
393    /// `Predicate` node.  The `default` value is used when no condition
394    /// matches.
395    pub fn piecewise(&self, branches: Vec<(ExprId, ExprId)>, default: ExprId) -> ExprId {
396        self.intern(ExprData::Piecewise { branches, default })
397    }
398
399    /// Build a `Predicate` node (symbolic boolean condition).
400    pub fn predicate(&self, kind: crate::kernel::expr::PredicateKind, args: Vec<ExprId>) -> ExprId {
401        self.intern(ExprData::Predicate { kind, args })
402    }
403
404    // Convenience constructors for common predicates.
405    pub fn pred_lt(&self, a: ExprId, b: ExprId) -> ExprId {
406        self.predicate(crate::kernel::expr::PredicateKind::Lt, vec![a, b])
407    }
408    pub fn pred_le(&self, a: ExprId, b: ExprId) -> ExprId {
409        self.predicate(crate::kernel::expr::PredicateKind::Le, vec![a, b])
410    }
411    pub fn pred_gt(&self, a: ExprId, b: ExprId) -> ExprId {
412        self.predicate(crate::kernel::expr::PredicateKind::Gt, vec![a, b])
413    }
414    pub fn pred_ge(&self, a: ExprId, b: ExprId) -> ExprId {
415        self.predicate(crate::kernel::expr::PredicateKind::Ge, vec![a, b])
416    }
417    pub fn pred_eq(&self, a: ExprId, b: ExprId) -> ExprId {
418        self.predicate(crate::kernel::expr::PredicateKind::Eq, vec![a, b])
419    }
420    pub fn pred_ne(&self, a: ExprId, b: ExprId) -> ExprId {
421        self.predicate(crate::kernel::expr::PredicateKind::Ne, vec![a, b])
422    }
423    pub fn pred_and(&self, args: Vec<ExprId>) -> ExprId {
424        self.predicate(crate::kernel::expr::PredicateKind::And, args)
425    }
426    pub fn pred_or(&self, args: Vec<ExprId>) -> ExprId {
427        self.predicate(crate::kernel::expr::PredicateKind::Or, args)
428    }
429    pub fn pred_not(&self, a: ExprId) -> ExprId {
430        self.predicate(crate::kernel::expr::PredicateKind::Not, vec![a])
431    }
432    pub fn pred_true(&self) -> ExprId {
433        self.predicate(crate::kernel::expr::PredicateKind::True, vec![])
434    }
435    pub fn pred_false(&self) -> ExprId {
436        self.predicate(crate::kernel::expr::PredicateKind::False, vec![])
437    }
438
439    // V3-3 — first-order quantifiers (first-class `Formula` / FOFormula).
440    /// `∀ var . body`
441    pub fn forall(&self, var: ExprId, body: ExprId) -> ExprId {
442        self.intern(ExprData::Forall { var, body })
443    }
444
445    /// `∃ var . body`
446    pub fn exists(&self, var: ExprId, body: ExprId) -> ExprId {
447        self.intern(ExprData::Exists { var, body })
448    }
449
450    /// `Σ_{c : poly(c)=0} body[var := c]` — a sum over the roots of `poly`.
451    pub fn root_sum(&self, poly: ExprId, var: ExprId, body: ExprId) -> ExprId {
452        self.intern(ExprData::RootSum { poly, var, body })
453    }
454
455    /// `O(arg)` — symbolic big-O bound used in truncated series (V2-15).
456    pub fn big_o(&self, arg: ExprId) -> ExprId {
457        self.intern(ExprData::BigO(arg))
458    }
459
460    /// Canonical `+∞` symbol for limits at infinity (V2-16).
461    pub fn pos_infinity(&self) -> ExprId {
462        self.symbol(POS_INFINITY_SYMBOL, Domain::Positive)
463    }
464
465    // -----------------------------------------------------------------------
466    // Display helper
467    // -----------------------------------------------------------------------
468
469    pub fn display(&self, id: ExprId) -> ExprDisplay<'_> {
470        ExprDisplay { id, pool: self }
471    }
472}
473
474impl Default for ExprPool {
475    fn default() -> Self {
476        Self::new()
477    }
478}
479
480// ---------------------------------------------------------------------------
481// Display — pool-aware recursive formatter
482// ---------------------------------------------------------------------------
483
484/// Wraps an `(ExprId, &ExprPool)` pair so it can implement [`fmt::Display`].
485pub struct ExprDisplay<'a> {
486    pub id: ExprId,
487    pub pool: &'a ExprPool,
488}
489
490impl fmt::Display for ExprDisplay<'_> {
491    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
492        let data = self.pool.get(self.id);
493        fmt_data(&data, self.pool, f)
494    }
495}
496
497impl fmt::Debug for ExprDisplay<'_> {
498    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
499        write!(f, "{}", self)
500    }
501}
502
503/// Format a power base or exponent, parenthesizing compound subexpressions.
504///
505/// `Add`/`Mul` already render with outer parentheses, so wrapping again would
506/// produce `((z + -2))^-1`. Only wrap forms that do not already self-group.
507fn fmt_pow_atom(id: ExprId, pool: &ExprPool) -> String {
508    let s = pool.display(id).to_string();
509    let needs_parens = match pool.get(id) {
510        ExprData::Symbol { .. } | ExprData::Integer(_) | ExprData::Float(_) => false,
511        ExprData::Func { .. } => false,
512        // Already printed as `(…)` by fmt_data.
513        ExprData::Add(_) | ExprData::Mul(_) => false,
514        ExprData::Rational(_)
515        | ExprData::Pow { .. }
516        | ExprData::Piecewise { .. }
517        | ExprData::Predicate { .. }
518        | ExprData::Forall { .. }
519        | ExprData::Exists { .. }
520        | ExprData::BigO(_)
521        | ExprData::RootSum { .. } => true,
522    };
523    if needs_parens {
524        format!("({s})")
525    } else {
526        s
527    }
528}
529
530fn fmt_data(data: &ExprData, pool: &ExprPool, f: &mut fmt::Formatter<'_>) -> fmt::Result {
531    match data {
532        ExprData::Symbol { name, .. } => write!(f, "{}", name),
533        ExprData::Integer(n) => write!(f, "{}", n),
534        ExprData::Rational(r) => write!(f, "{}", r),
535        ExprData::Float(fl) => write!(f, "{}", fl),
536        ExprData::Add(args) => {
537            write!(f, "(")?;
538            for (i, &arg) in args.iter().enumerate() {
539                if i > 0 {
540                    write!(f, " + ")?;
541                }
542                write!(f, "{}", pool.display(arg))?;
543            }
544            write!(f, ")")
545        }
546        ExprData::Mul(args) => {
547            write!(f, "(")?;
548            for (i, &arg) in args.iter().enumerate() {
549                if i > 0 {
550                    write!(f, " * ")?;
551                }
552                write!(f, "{}", pool.display(arg))?;
553            }
554            write!(f, ")")
555        }
556        ExprData::Pow { base, exp } => {
557            // Parenthesize compound bases/exponents so `x^(1/2)^3` cannot be
558            // misread as `x^1 / 2^3`. Prefer `(x^(1/2))^3`.
559            let base_s = fmt_pow_atom(*base, pool);
560            let exp_s = fmt_pow_atom(*exp, pool);
561            write!(f, "{base_s}^{exp_s}")
562        }
563        ExprData::Func { name, args } => {
564            write!(f, "{}(", name)?;
565            for (i, &arg) in args.iter().enumerate() {
566                if i > 0 {
567                    write!(f, ", ")?;
568                }
569                write!(f, "{}", pool.display(arg))?;
570            }
571            write!(f, ")")
572        }
573        ExprData::Piecewise { branches, default } => {
574            write!(f, "Piecewise(")?;
575            for (i, (cond, val)) in branches.iter().enumerate() {
576                if i > 0 {
577                    write!(f, ", ")?;
578                }
579                write!(f, "({}, {})", pool.display(*cond), pool.display(*val))?;
580            }
581            write!(f, "; default={})", pool.display(*default))
582        }
583        ExprData::Predicate { kind, args } => match kind {
584            crate::kernel::expr::PredicateKind::True => write!(f, "True"),
585            crate::kernel::expr::PredicateKind::False => write!(f, "False"),
586            crate::kernel::expr::PredicateKind::Not => {
587                write!(f, "¬({})", pool.display(args[0]))
588            }
589            crate::kernel::expr::PredicateKind::And | crate::kernel::expr::PredicateKind::Or => {
590                write!(f, "(")?;
591                for (i, &arg) in args.iter().enumerate() {
592                    if i > 0 {
593                        write!(f, " {} ", kind)?;
594                    }
595                    write!(f, "{}", pool.display(arg))?;
596                }
597                write!(f, ")")
598            }
599            _ => {
600                write!(
601                    f,
602                    "({} {} {})",
603                    pool.display(args[0]),
604                    kind,
605                    pool.display(args[1])
606                )
607            }
608        },
609        ExprData::Forall { var, body } => {
610            write!(f, "∀ {} . {}", pool.display(*var), pool.display(*body))
611        }
612        ExprData::Exists { var, body } => {
613            write!(f, "∃ {} . {}", pool.display(*var), pool.display(*body))
614        }
615        ExprData::BigO(arg) => {
616            write!(f, "O({})", pool.display(*arg))
617        }
618        ExprData::RootSum { poly, var, body } => {
619            write!(
620                f,
621                "RootSum({}, {} . {})",
622                pool.display(*poly),
623                pool.display(*var),
624                pool.display(*body)
625            )
626        }
627    }
628}
629
630// ---------------------------------------------------------------------------
631// Unit tests
632// ---------------------------------------------------------------------------
633
634#[cfg(test)]
635mod tests {
636    use super::*;
637    use crate::kernel::domain::Domain;
638
639    fn pool() -> ExprPool {
640        ExprPool::new()
641    }
642
643    #[test]
644    fn noncommutative_mul_orders_distinct() {
645        let p = pool();
646        let a = p.symbol_commutative("A", Domain::Real, false);
647        let b = p.symbol_commutative("B", Domain::Real, false);
648        assert_ne!(
649            p.mul(vec![a, b]),
650            p.mul(vec![b, a]),
651            "A*B and B*A must not hash-cons together for NC symbols"
652        );
653    }
654
655    #[test]
656    fn symbol_commutative_is_structural() {
657        let p = pool();
658        let xc = p.symbol_commutative("x", Domain::Real, true);
659        let xnc = p.symbol_commutative("x", Domain::Real, false);
660        assert_ne!(xc, xnc);
661    }
662
663    // --- construction and equality ---
664
665    #[test]
666    fn symbol_interning() {
667        let p = pool();
668        let x1 = p.symbol("x", Domain::Real);
669        let x2 = p.symbol("x", Domain::Real);
670        assert_eq!(x1, x2, "same symbol must return same ExprId");
671    }
672
673    #[test]
674    fn domain_is_structural() {
675        let p = pool();
676        let xr = p.symbol("x", Domain::Real);
677        let xc = p.symbol("x", Domain::Complex);
678        assert_ne!(xr, xc, "same name but different domain must be distinct");
679    }
680
681    #[test]
682    fn integer_interning() {
683        let p = pool();
684        let a = p.integer(42_i32);
685        let b = p.integer(42_i32);
686        let c = p.integer(99_i32);
687        assert_eq!(a, b);
688        assert_ne!(a, c);
689    }
690
691    #[test]
692    fn rational_canonical() {
693        let p = pool();
694        // 2/4 reduces to 1/2
695        let r1 = p.rational(2_i32, 4_i32);
696        let r2 = p.rational(1_i32, 2_i32);
697        assert_eq!(r1, r2, "rationals must be reduced to canonical form");
698    }
699
700    #[test]
701    fn float_precision_is_structural() {
702        let p = pool();
703        let f53 = p.float(1.0, 53);
704        let f64_ = p.float(1.0, 64);
705        assert_ne!(
706            f53, f64_,
707            "same value but different precision is a different expr"
708        );
709    }
710
711    // --- compound expressions and subexpression sharing ---
712
713    #[test]
714    fn subexpression_sharing() {
715        let p = pool();
716        let x = p.symbol("x", Domain::Real);
717        let two = p.integer(2_i32);
718
719        // Build x^2 twice; both must return the same ExprId.
720        let xsq1 = p.pow(x, two);
721        let xsq2 = p.pow(x, two);
722        assert_eq!(xsq1, xsq2);
723
724        // Pool should have exactly 3 nodes: x, 2, x^2.
725        assert_eq!(p.len(), 3);
726    }
727
728    #[test]
729    fn add_interning() {
730        let p = pool();
731        let x = p.symbol("x", Domain::Real);
732        let y = p.symbol("y", Domain::Real);
733        let s1 = p.add(vec![x, y]);
734        let s2 = p.add(vec![x, y]);
735        assert_eq!(s1, s2);
736    }
737
738    #[test]
739    fn arg_order_is_canonical() {
740        // PA-3: Add/Mul children are sorted at construction time so that
741        // commutativity holds structurally — a+b and b+a intern to the same ExprId.
742        let p = pool();
743        let x = p.symbol("x", Domain::Real);
744        let y = p.symbol("y", Domain::Real);
745        let s1 = p.add(vec![x, y]);
746        let s2 = p.add(vec![y, x]);
747        assert_eq!(s1, s2, "a+b and b+a must be the same expression after PA-3");
748        let m1 = p.mul(vec![x, y]);
749        let m2 = p.mul(vec![y, x]);
750        assert_eq!(m1, m2, "a*b and b*a must be the same expression after PA-3");
751    }
752
753    #[test]
754    fn func_interning() {
755        let p = pool();
756        let x = p.symbol("x", Domain::Real);
757        let s1 = p.func("sin", vec![x]);
758        let s2 = p.func("sin", vec![x]);
759        let c1 = p.func("cos", vec![x]);
760        assert_eq!(s1, s2);
761        assert_ne!(s1, c1);
762    }
763
764    // --- display ---
765
766    #[test]
767    fn display_symbol() {
768        let p = pool();
769        let x = p.symbol("x", Domain::Real);
770        assert_eq!(p.display(x).to_string(), "x");
771    }
772
773    #[test]
774    fn display_integer() {
775        let p = pool();
776        let n = p.integer(42_i32);
777        assert_eq!(p.display(n).to_string(), "42");
778    }
779
780    #[test]
781    fn display_pow() {
782        let p = pool();
783        let x = p.symbol("x", Domain::Real);
784        let two = p.integer(2_i32);
785        let xsq = p.pow(x, two);
786        assert_eq!(p.display(xsq).to_string(), "x^2");
787    }
788
789    #[test]
790    fn display_add() {
791        let p = pool();
792        let x = p.symbol("x", Domain::Real);
793        let y = p.symbol("y", Domain::Real);
794        let s = p.add(vec![x, y]);
795        assert_eq!(p.display(s).to_string(), "(x + y)");
796    }
797
798    #[test]
799    fn display_func() {
800        let p = pool();
801        let x = p.symbol("x", Domain::Real);
802        let s = p.func("sin", vec![x]);
803        assert_eq!(p.display(s).to_string(), "sin(x)");
804    }
805
806    #[test]
807    fn display_nested() {
808        let p = pool();
809        let x = p.symbol("x", Domain::Real);
810        let two = p.integer(2_i32);
811        let xsq = p.pow(x, two);
812        let one = p.integer(1_i32);
813        let expr = p.add(vec![xsq, one]);
814        assert_eq!(p.display(expr).to_string(), "(x^2 + 1)");
815    }
816
817    // --- send + sync: compile-time check ---
818
819    fn assert_send_sync<T: Send + Sync>() {}
820
821    #[test]
822    fn pool_is_send_sync() {
823        assert_send_sync::<ExprPool>();
824    }
825}