feat: QOL improvements

Author: arnog

src/compute-engine/boxed-expression/abstract-boxed-expression.ts

@@ -64,6 +64,27 @@ export function _setSerializeJson(fn: SerializeJsonFn) {
   _serializeJson = fn;
 }
 
+type ExpandFn = (expr: Expression) => Expression;
+let _expandForIs: ExpandFn;
+/** @internal */
+export function _setExpandForIs(fn: ExpandFn) {
+  _expandForIs = fn;
+}
+
+const EXPANDABLE_OPS = new Set(['Multiply', 'Power', 'Negate', 'Divide']);
+
+/** Return true if at least one side has an operator where expansion could
+ *  produce a structurally different result. */
+function _couldBenefitFromExpand(
+  a: _BoxedExpression,
+  b: Expression | number | bigint | boolean | string
+): boolean {
+  if (EXPANDABLE_OPS.has(a.operator)) return true;
+  if (typeof b === 'object' && b !== null && 'operator' in b)
+    return EXPANDABLE_OPS.has(b.operator);
+  return false;
+}
+
 /**
  * _BoxedExpression
  *
@@ -471,13 +492,38 @@ export abstract class _BoxedExpression implements Expression {
     return [this, this.engine.One];
   }
 
+  toRational(): [number, number] | null {
+    return null;
+  }
+
+  factors(): ReadonlyArray<Expression> {
+    return [this];
+  }
+
   is(
     other: Expression | number | bigint | boolean | string,
     tolerance?: number
   ): boolean {
     // Fast path: exact structural/value check
     if (this.isSame(other)) return true;
 
+    // Try expansion — catches equivalences like (x+1)^2 vs x^2+2x+1
+    // even when the expression has free variables (where the numeric
+    // fallback below would bail out).
+    //
+    // Only expand when at least one side contains Multiply, Power, or
+    // Negate — those are the only operators where expansion can produce
+    // a structurally different result. This avoids needlessly
+    // reconstructing Add trees (expand recurses into Add operands).
+    if (_expandForIs && _couldBenefitFromExpand(this, other)) {
+      const expandedThis = _expandForIs(this);
+      if (expandedThis !== this && expandedThis.isSame(other)) return true;
+      if (other instanceof _BoxedExpression) {
+        const expandedOther = _expandForIs(other);
+        if (expandedThis.isSame(expandedOther)) return true;
+      }
+    }
+
     // Numeric fallback only when there are no free variables
     if (this.freeVariables.length > 0) return false;
 

src/compute-engine/boxed-expression/boxed-function.ts

@@ -52,6 +52,7 @@ import { canonicalMultiply, mul, div, Product } from './arithmetic-mul-div';
 import { add } from './arithmetic-add';
 import { pow } from './arithmetic-power';
 import { asSmallInteger } from './numerics';
+import { gcd } from '../numerics/numeric';
 import { _BoxedExpression } from './abstract-boxed-expression';
 import { DEFAULT_COMPLEXITY, sortOperands } from './order';
 import {
@@ -629,6 +630,42 @@ export class BoxedFunction
     return [this, this.engine.One];
   }
 
+  factors(): ReadonlyArray<Expression> {
+    const op = this.operator;
+    if (op === 'Multiply') {
+      const result: Expression[] = [];
+      for (const arg of this.ops) result.push(...arg.factors());
+      return result;
+    }
+    if (op === 'Negate') {
+      return [this.engine.number(-1), ...this.op1.factors()];
+    }
+    return [this];
+  }
+
+  toRational(): [number, number] | null {
+    const op = this.operator;
+    if (op === 'Divide' || op === 'Rational') {
+      const num = this.op1.re;
+      const den = this.op2.re;
+      if (
+        Number.isInteger(num) &&
+        Number.isInteger(den) &&
+        den !== 0
+      ) {
+        const g = gcd(Math.abs(num), Math.abs(den));
+        const sign = den < 0 ? -1 : 1;
+        return [(sign * num) / g, (sign * den) / g];
+      }
+      return null;
+    }
+    if (op === 'Negate') {
+      const r = this.op1.toRational();
+      return r ? [-r[0], r[1]] : null;
+    }
+    return null;
+  }
+
   //
   //
   // ALGEBRAIC OPERATIONS

src/compute-engine/boxed-expression/boxed-number.ts

@@ -455,6 +455,25 @@ export class BoxedNumber
     ];
   }
 
+  toRational(): [number, number] | null {
+    if (typeof this._value === 'number') {
+      if (!Number.isFinite(this._value)) return null;
+      if (Number.isInteger(this._value)) return [this._value, 1];
+      return null;
+    }
+    // NumericValue — check it's a pure rational (no radical, no imaginary)
+    if (this._value.im !== 0) return null;
+    const exact = this._value.asExact;
+    if (!exact) return null;
+    const ev = exact as ExactNumericValue;
+    if (ev.radical !== 1) return null;
+    const r = ev.rational;
+    const num = Number(r[0]);
+    const den = Number(r[1]);
+    if (!Number.isFinite(num) || !Number.isFinite(den)) return null;
+    return [num, den];
+  }
+
   subs(
     sub: Substitution,
     options?: { canonical?: CanonicalOptions }

src/compute-engine/boxed-expression/compare.ts

@@ -12,7 +12,7 @@ import { stochasticEqual } from './stochastic-equal';
 
 // Lazy reference to break circular dependency:
 // expand → arithmetic-add → boxed-tensor → abstract-boxed-expression → compare
-type ExpandFn = (expr: Expression | undefined) => Expression | null;
+type ExpandFn = (expr: Expression) => Expression;
 let _expand: ExpandFn;
 /** @internal */
 export function _setExpand(fn: ExpandFn) {
@@ -135,8 +135,8 @@ export function eq(
     }
 
     // Try structural equality after expand+simplify first
-    a = (_expand(a) ?? a).simplify();
-    b = (_expand(b) ?? b).simplify();
+    a = _expand(a).simplify();
+    b = _expand(b).simplify();
     if (same(a, b)) return true;
 
     // Fall back to stochastic evaluation at random sample points

src/compute-engine/boxed-expression/expand.ts

@@ -129,7 +129,6 @@ export function expandFunction(
   //
   if (h === 'Divide') {
     const num = expand(ops[0]);
-    if (num === null) return null;
     if (isFunction(num, 'Add'))
       return add(...num.ops.map((x) => x.div(ops[1])));
     return ce._fn('Divide', [num, ops[1]]);
@@ -160,14 +159,14 @@ export function expandFunction(
   //
   // Negate
   //
-  if (h === 'Negate') return expand(ops[0])?.neg() ?? null;
+  if (h === 'Negate') return expand(ops[0]).neg();
 
   //
   //
   // Add
   //
 
-  if (h === 'Add') return add(...ops.map((x) => expand(x) ?? x));
+  if (h === 'Add') return add(...ops.map((x) => expand(x)));
 
   //
   // Power
@@ -188,11 +187,11 @@ export function expandFunction(
  * If the exression is a relational operator, expand the operands.
  * Return null if the expression cannot be expanded.
  */
-export function expand(expr: Expression | undefined): Expression | null {
+export function expand(expr: Expression): Expression {
   // To expand an expression, we need to use its canonical form
-  expr = expr?.canonical;
+  expr = expr.canonical;
 
-  if (!expr || typeof expr.operator !== 'string') return null;
+  if (typeof expr.operator !== 'string') return expr;
 
   //
   // Expand relational operators
@@ -201,14 +200,16 @@ export function expand(expr: Expression | undefined): Expression | null {
     const ops = isFunction(expr) ? expr.ops : [];
     return expr.engine._fn(
       expr.operator,
-      ops.map((x) => expand(x) ?? x)
+      ops.map((x) => expand(x))
     );
   }
 
-  return expandFunction(
-    expr.engine,
-    expr.operator,
-    isFunction(expr) ? expr.ops : []
+  return (
+    expandFunction(
+      expr.engine,
+      expr.operator,
+      isFunction(expr) ? expr.ops : []
+    ) ?? expr
   );
 }
 
@@ -217,11 +218,11 @@ export function expand(expr: Expression | undefined): Expression | null {
  *
  * `expand()` only expands the top level of the expression.
  */
-export function expandAll(expr: Expression): Expression | null {
-  if (!expr.operator || !isFunction(expr)) return null;
+export function expandAll(expr: Expression): Expression {
+  if (!expr.operator || !isFunction(expr)) return expr;
 
-  const ops = expr.ops.map((x) => expandAll(x) ?? x);
+  const ops = expr.ops.map((x) => expandAll(x));
 
   const result = expr.engine.function(expr.operator, ops);
-  return expand(result) ?? result;
+  return expand(result);
 }

src/compute-engine/boxed-expression/init-lazy-refs.ts

@@ -6,14 +6,15 @@ import { expand } from './expand';
 import { _setExpand } from './compare';
 
 import { serializeJson, _setProduct } from './serialize';
-import { _setSerializeJson } from './abstract-boxed-expression';
+import { _setSerializeJson, _setExpandForIs } from './abstract-boxed-expression';
 
 import { Product } from './arithmetic-mul-div';
 
 import { compile } from '../compilation/compile-expression';
 import { _setCompile } from './stochastic-equal';
 
 _setExpand(expand);
+_setExpandForIs(expand);
 _setSerializeJson(serializeJson);
 _setProduct(Product);
 _setCompile(compile);

src/compute-engine/boxed-expression/polynomials.ts

@@ -221,7 +221,7 @@ export function getPolynomialCoefficients(
   const coeffs: Expression[] = new Array(degree + 1).fill(ce.Zero);
 
   // Expand the expression to get standard form
-  const expanded = expand(expr) ?? expr;
+  const expanded = expand(expr);
 
   // Helper to add a term's coefficient at a specific degree
   const addCoefficient = (coef: Expression, deg: number): boolean => {

src/compute-engine/boxed-expression/solve-linear-system.ts

@@ -174,10 +174,10 @@ function extractLinearCoefficients(
   if (isFunction(equation, 'Equal')) {
     const lhs = equation.op1;
     const rhs = equation.op2;
-    expr = expand(lhs.sub(rhs)) ?? lhs.sub(rhs);
+    expr = expand(lhs.sub(rhs));
   } else {
     // Assume equation = 0
-    expr = expand(equation) ?? equation;
+    expr = expand(equation);
   }
 
   // Check that all variables appear with degree at most 1 (linear)
@@ -680,9 +680,9 @@ export function solvePolynomialSystem(
   const normalized = equations.map((eq) => {
     if (isFunction(eq, 'Equal')) {
       const diff = eq.op1.sub(eq.op2);
-      return (expand(diff) ?? diff).simplify();
+      return expand(diff).simplify();
     }
-    return (expand(eq) ?? eq).simplify();
+    return expand(eq).simplify();
   });
 
   // Try product + sum pattern first
@@ -1182,12 +1182,12 @@ function extractLinearConstraint(
 
   if (op === 'Less' || op === 'LessEqual') {
     const diff1 = lhs.sub(rhs);
-    expr = (expand(diff1) ?? diff1).simplify();
+    expr = expand(diff1).simplify();
     strict = op === 'Less';
   } else {
     // Greater or GreaterEqual: flip to Less form
     const diff2 = rhs.sub(lhs);
-    expr = (expand(diff2) ?? diff2).simplify();
+    expr = expand(diff2).simplify();
     strict = op === 'Greater';
   }
 

src/compute-engine/boxed-expression/solve.ts

@@ -1151,10 +1151,10 @@ export function findUnivariateRoots(
   const ce = expr.engine;
 
   if (isFunction(expr, 'Equal')) {
-    const lhs = expand(expr.op1) ?? expr.op1;
-    const rhs = expand(expr.op2) ?? expr.op2;
+    const lhs = expand(expr.op1);
+    const rhs = expand(expr.op2);
     expr = lhs.sub(rhs).simplify();
-  } else expr = (expand(expr) ?? expr).simplify();
+  } else expr = expand(expr).simplify();
 
   // Save the expression BEFORE clearing denominators and other transformations.
   // This is crucial for validating roots: clearing denominators and harmonization

src/compute-engine/free-functions.ts

@@ -90,7 +90,7 @@ export function assign(
   getDefaultEngine().assign(arg1 as any, arg2 as any);
 }
 
-export function expand(expr: LatexString | ExpressionInput): Expression | null {
+export function expand(expr: LatexString | ExpressionInput): Expression {
   return expandExpr(toExpression(expr));
 }
 
@@ -107,7 +107,7 @@ export function solve(
 
 export function expandAll(
   expr: LatexString | ExpressionInput
-): Expression | null {
+): Expression {
   return expandAllExpr(toExpression(expr));
 }