gradient-script 0.2.0 → 0.3.0

package/dist/dsl/egraph/Extractor.js

@@ -0,0 +1,1068 @@
+/**
+ * Cost-Based Extraction from E-Graphs
+ *
+ * Extracts the lowest-cost expression from each e-class.
+ * Also detects common subexpressions (shared e-classes) for CSE.
+ */
+import { enodeChildren } from './ENode.js';
+// =============================================================================
+// Constant Folding Helpers
+// =============================================================================
+/**
+ * Create a binary expression, folding constants when possible
+ */
+function makeBinary(operator, left, right) {
+    // If both operands are numbers, fold the constant
+    if (left.kind === 'number' && right.kind === 'number') {
+        const l = left.value;
+        const r = right.value;
+        let result;
+        switch (operator) {
+            case '+':
+                result = l + r;
+                break;
+            case '-':
+                result = l - r;
+                break;
+            case '*':
+                result = l * r;
+                break;
+            case '/':
+                result = r !== 0 ? l / r : NaN;
+                break;
+            case '^':
+                result = Math.pow(l, r);
+                break;
+        }
+        // Only fold if result is a finite number
+        if (Number.isFinite(result)) {
+            return { kind: 'number', value: result };
+        }
+    }
+    return { kind: 'binary', operator, left, right };
+}
+/**
+ * Create a unary expression, folding constants when possible
+ */
+function makeUnary(operator, operand) {
+    // If operand is a number, fold the constant
+    if (operand.kind === 'number') {
+        return { kind: 'number', value: -operand.value };
+    }
+    return { kind: 'unary', operator, operand };
+}
+/**
+ * Check if an expression is trivial (should never be a temp)
+ * - Constants: 5, -2, etc.
+ * - Simple negations of constants: -(2)
+ * - Variables: a, b
+ */
+function isTrivialExpression(expr) {
+    if (expr.kind === 'number')
+        return true;
+    if (expr.kind === 'variable')
+        return true;
+    if (expr.kind === 'unary' && expr.operand.kind === 'number')
+        return true;
+    return false;
+}
+/**
+ * Default cost model - division is expensive
+ */
+export const defaultCostModel = {
+    num: 1,
+    var: 1,
+    add: 2,
+    sub: 2,
+    mul: 2,
+    div: 8, // Division is expensive - encourage factoring
+    pow: 4,
+    neg: 1,
+    inv: 5, // Inverse (1/x) - cheaper than div but significant
+    call: 3,
+    component: 1,
+};
+/**
+ * Extract the best expression from an e-class
+ */
+export function extractBest(egraph, rootId, costModel = defaultCostModel) {
+    const costs = computeCosts(egraph, costModel);
+    return extractFromClass(egraph, rootId, costs, costModel);
+}
+/**
+ * Extract multiple expressions with CSE (shared subexpressions become temps)
+ */
+export function extractWithCSE(egraph, roots, costModel = defaultCostModel, minSharedCost = 3 // Only extract temps if cost > this
+) {
+    // Compute costs for all e-classes
+    const costs = computeCosts(egraph, costModel);
+    // Count references to each e-class from roots
+    const refCounts = countReferences(egraph, roots, costs, costModel);
+    // Decide which classes should become temps (count >= 2 means used multiple times)
+    const tempsToExtract = new Map();
+    let tempCounter = 0;
+    for (const [classId, count] of refCounts) {
+        if (count >= 2) {
+            const classCost = costs.get(egraph.find(classId)) ?? Infinity;
+            if (classCost > minSharedCost) {
+                tempsToExtract.set(egraph.find(classId), `_tmp${tempCounter++}`);
+            }
+        }
+    }
+    // Extract temp definitions (without referencing other temps initially)
+    const temps = new Map();
+    for (const [classId, tempName] of tempsToExtract) {
+        const expr = extractFromClass(egraph, classId, costs, costModel);
+        // Skip trivial expressions that shouldn't be temps
+        // (constants, or simple negations of constants)
+        if (isTrivialExpression(expr)) {
+            tempsToExtract.delete(classId);
+            continue;
+        }
+        temps.set(tempName, expr);
+    }
+    // Extract root expressions, using temps where available
+    const expressions = new Map();
+    for (const rootId of roots) {
+        const expr = extractWithTemps(egraph, rootId, costs, costModel, tempsToExtract);
+        expressions.set(rootId, expr);
+    }
+    // Post-process: substitute temps into other temp definitions where possible
+    // Build a map from expression serialization to temp name
+    const exprToTemp = new Map();
+    for (const [tempName, expr] of temps) {
+        exprToTemp.set(serializeExpr(expr), tempName);
+    }
+    // Substitute temps into temp definitions
+    for (const [tempName, expr] of temps) {
+        temps.set(tempName, substituteTempRefs(expr, exprToTemp, tempName));
+    }
+    // Topologically sort temps by dependency (deps first)
+    const sortedTemps = topologicalSortTemps(temps);
+    temps.clear();
+    for (const [name, expr] of sortedTemps) {
+        temps.set(name, expr);
+    }
+    // Count actual usage of each temp in the final output
+    const tempUsageCounts = new Map();
+    function countTempUsage(expr) {
+        if (expr.kind === 'variable' && expr.name.startsWith('_tmp')) {
+            tempUsageCounts.set(expr.name, (tempUsageCounts.get(expr.name) ?? 0) + 1);
+        }
+        else if (expr.kind === 'binary') {
+            countTempUsage(expr.left);
+            countTempUsage(expr.right);
+        }
+        else if (expr.kind === 'unary') {
+            countTempUsage(expr.operand);
+        }
+        else if (expr.kind === 'call') {
+            expr.args.forEach(countTempUsage);
+        }
+        else if (expr.kind === 'component') {
+            countTempUsage(expr.object);
+        }
+    }
+    // Count in root expressions
+    for (const expr of expressions.values()) {
+        countTempUsage(expr);
+    }
+    // Also count in temp definitions (temps can reference other temps)
+    for (const expr of temps.values()) {
+        countTempUsage(expr);
+    }
+    // Identify temps to inline (used 0 or 1 times)
+    const tempsToInline = new Set();
+    for (const [tempName] of temps) {
+        const count = tempUsageCounts.get(tempName) ?? 0;
+        if (count <= 1) {
+            tempsToInline.add(tempName);
+        }
+    }
+    // If there are temps to inline, substitute them back
+    if (tempsToInline.size > 0) {
+        function inlineTemps(expr) {
+            if (expr.kind === 'variable' && tempsToInline.has(expr.name)) {
+                const tempExpr = temps.get(expr.name);
+                return tempExpr ? inlineTemps(tempExpr) : expr;
+            }
+            else if (expr.kind === 'binary') {
+                return {
+                    kind: 'binary',
+                    operator: expr.operator,
+                    left: inlineTemps(expr.left),
+                    right: inlineTemps(expr.right)
+                };
+            }
+            else if (expr.kind === 'unary') {
+                return {
+                    kind: 'unary',
+                    operator: expr.operator,
+                    operand: inlineTemps(expr.operand)
+                };
+            }
+            else if (expr.kind === 'call') {
+                return {
+                    kind: 'call',
+                    name: expr.name,
+                    args: expr.args.map(inlineTemps)
+                };
+            }
+            else if (expr.kind === 'component') {
+                return {
+                    kind: 'component',
+                    object: inlineTemps(expr.object),
+                    component: expr.component
+                };
+            }
+            return expr;
+        }
+        // Inline in remaining temps (temps not being inlined)
+        for (const [tempName, expr] of temps) {
+            if (!tempsToInline.has(tempName)) {
+                temps.set(tempName, inlineTemps(expr));
+            }
+        }
+        // Inline in root expressions
+        for (const [rootId, expr] of expressions) {
+            expressions.set(rootId, inlineTemps(expr));
+        }
+        // Remove inlined temps
+        for (const tempName of tempsToInline) {
+            temps.delete(tempName);
+        }
+        // Re-sort topologically after inlining (inlining may have changed dependencies)
+        const reSorted = topologicalSortTemps(temps);
+        temps.clear();
+        for (const [name, expr] of reSorted) {
+            temps.set(name, expr);
+        }
+    }
+    // Post-extraction CSE: find repeated patterns that emerge AFTER temp substitution
+    // e.g., "_tmp22 + _tmp23" appearing multiple times
+    postExtractionCSE(temps, expressions, minSharedCost, tempCounter, costModel);
+    // Detect and merge (a-b) / (b-a) patterns (these are negatives of each other)
+    mergeNegativePairs(temps, expressions);
+    // Normalize patterns like a + -1 * b to a - b (cleanup from e-graph extraction)
+    normalizeAddNegMul(temps, expressions);
+    // Calculate total cost
+    let totalCost = 0;
+    for (const [, expr] of temps) {
+        totalCost += expressionCost(expr, costModel);
+    }
+    for (const [, expr] of expressions) {
+        totalCost += expressionCost(expr, costModel);
+    }
+    return { temps, expressions, totalCost };
+}
+/**
+ * Compute the minimum cost for each e-class (bottom-up)
+ */
+function computeCosts(egraph, costModel) {
+    const costs = new Map();
+    const classIds = egraph.getClassIds();
+    // Initialize all costs to infinity
+    for (const id of classIds) {
+        costs.set(id, Infinity);
+    }
+    // Iterate until convergence
+    let changed = true;
+    let iterations = 0;
+    const maxIterations = 100;
+    while (changed && iterations < maxIterations) {
+        changed = false;
+        iterations++;
+        for (const classId of classIds) {
+            const canonId = egraph.find(classId);
+            const nodes = egraph.getNodes(canonId);
+            for (const node of nodes) {
+                const nodeCost = computeNodeCost(node, costs, costModel);
+                const currentCost = costs.get(canonId) ?? Infinity;
+                if (nodeCost < currentCost) {
+                    costs.set(canonId, nodeCost);
+                    changed = true;
+                }
+            }
+        }
+    }
+    return costs;
+}
+/**
+ * Compute cost of a single e-node
+ */
+function computeNodeCost(node, classCosts, costModel) {
+    const childCost = (id) => classCosts.get(id) ?? Infinity;
+    switch (node.tag) {
+        case 'num':
+            return costModel.num;
+        case 'var':
+            return costModel.var;
+        case 'add':
+            return costModel.add + childCost(node.children[0]) + childCost(node.children[1]);
+        case 'sub':
+            return costModel.sub + childCost(node.children[0]) + childCost(node.children[1]);
+        case 'mul':
+            return costModel.mul + childCost(node.children[0]) + childCost(node.children[1]);
+        case 'div':
+            return costModel.div + childCost(node.children[0]) + childCost(node.children[1]);
+        case 'pow':
+            return costModel.pow + childCost(node.children[0]) + childCost(node.children[1]);
+        case 'neg':
+            return costModel.neg + childCost(node.child);
+        case 'inv':
+            return costModel.inv + childCost(node.child);
+        case 'call':
+            return costModel.call + node.children.reduce((sum, id) => sum + childCost(id), 0);
+        case 'component':
+            return costModel.component + childCost(node.object);
+    }
+}
+/**
+ * Extract expression from an e-class using precomputed costs
+ */
+function extractFromClass(egraph, classId, costs, costModel) {
+    const canonId = egraph.find(classId);
+    const nodes = egraph.getNodes(canonId);
+    // Find lowest-cost node
+    let bestNode = null;
+    let bestCost = Infinity;
+    for (const node of nodes) {
+        const cost = computeNodeCost(node, costs, costModel);
+        if (cost < bestCost) {
+            bestCost = cost;
+            bestNode = node;
+        }
+    }
+    if (!bestNode) {
+        throw new Error(`No nodes in e-class ${canonId}`);
+    }
+    return nodeToExpression(bestNode, egraph, costs, costModel);
+}
+/**
+ * Extract expression, substituting temps where available
+ */
+function extractWithTemps(egraph, classId, costs, costModel, temps) {
+    const canonId = egraph.find(classId);
+    // Check if this class is a temp
+    const tempName = temps.get(canonId);
+    if (tempName) {
+        return { kind: 'variable', name: tempName };
+    }
+    const nodes = egraph.getNodes(canonId);
+    // Find lowest-cost node
+    let bestNode = null;
+    let bestCost = Infinity;
+    for (const node of nodes) {
+        const cost = computeNodeCost(node, costs, costModel);
+        if (cost < bestCost) {
+            bestCost = cost;
+            bestNode = node;
+        }
+    }
+    if (!bestNode) {
+        throw new Error(`No nodes in e-class ${canonId}`);
+    }
+    return nodeToExpressionWithTemps(bestNode, egraph, costs, costModel, temps);
+}
+/**
+ * Convert e-node to AST Expression (with constant folding)
+ */
+function nodeToExpression(node, egraph, costs, costModel) {
+    const extract = (id) => extractFromClass(egraph, id, costs, costModel);
+    switch (node.tag) {
+        case 'num':
+            return { kind: 'number', value: node.value };
+        case 'var':
+            return { kind: 'variable', name: node.name };
+        case 'add':
+            return makeBinary('+', extract(node.children[0]), extract(node.children[1]));
+        case 'sub':
+            return makeBinary('-', extract(node.children[0]), extract(node.children[1]));
+        case 'mul':
+            return makeBinary('*', extract(node.children[0]), extract(node.children[1]));
+        case 'div':
+            return makeBinary('/', extract(node.children[0]), extract(node.children[1]));
+        case 'pow':
+            return makeBinary('^', extract(node.children[0]), extract(node.children[1]));
+        case 'neg':
+            return makeUnary('-', extract(node.child));
+        case 'inv':
+            // inv(x) extracts as 1/x (with constant folding)
+            return makeBinary('/', { kind: 'number', value: 1 }, extract(node.child));
+        case 'call':
+            return {
+                kind: 'call',
+                name: node.name,
+                args: node.children.map(id => extract(id))
+            };
+        case 'component':
+            return {
+                kind: 'component',
+                object: extract(node.object),
+                component: node.field
+            };
+    }
+}
+/**
+ * Convert e-node to AST Expression, using temps
+ */
+function nodeToExpressionWithTemps(node, egraph, costs, costModel, temps) {
+    const extract = (id) => extractWithTemps(egraph, id, costs, costModel, temps);
+    switch (node.tag) {
+        case 'num':
+            return { kind: 'number', value: node.value };
+        case 'var':
+            return { kind: 'variable', name: node.name };
+        case 'add':
+            return makeBinary('+', extract(node.children[0]), extract(node.children[1]));
+        case 'sub':
+            return makeBinary('-', extract(node.children[0]), extract(node.children[1]));
+        case 'mul':
+            return makeBinary('*', extract(node.children[0]), extract(node.children[1]));
+        case 'div':
+            return makeBinary('/', extract(node.children[0]), extract(node.children[1]));
+        case 'pow':
+            return makeBinary('^', extract(node.children[0]), extract(node.children[1]));
+        case 'neg':
+            return makeUnary('-', extract(node.child));
+        case 'inv':
+            // inv(x) extracts as 1/x (with constant folding)
+            return makeBinary('/', { kind: 'number', value: 1 }, extract(node.child));
+        case 'call':
+            return {
+                kind: 'call',
+                name: node.name,
+                args: node.children.map(id => extract(id))
+            };
+        case 'component':
+            return {
+                kind: 'component',
+                object: extract(node.object),
+                component: node.field
+            };
+    }
+}
+/**
+ * Count references to each e-class from root expressions
+ */
+function countReferences(egraph, roots, costs, costModel) {
+    const counts = new Map();
+    function countInClass(classId, visited) {
+        const canonId = egraph.find(classId);
+        // Increment reference count
+        counts.set(canonId, (counts.get(canonId) ?? 0) + 1);
+        // Don't recurse if already visited in this path
+        if (visited.has(canonId)) {
+            return;
+        }
+        visited.add(canonId);
+        // Get best node and recurse into children
+        const nodes = egraph.getNodes(canonId);
+        let bestNode = null;
+        let bestCost = Infinity;
+        for (const node of nodes) {
+            const cost = computeNodeCost(node, costs, costModel);
+            if (cost < bestCost) {
+                bestCost = cost;
+                bestNode = node;
+            }
+        }
+        if (bestNode) {
+            for (const childId of enodeChildren(bestNode)) {
+                countInClass(childId, new Set(visited));
+            }
+        }
+    }
+    for (const rootId of roots) {
+        countInClass(rootId, new Set());
+    }
+    return counts;
+}
+/**
+ * Calculate cost of an AST expression
+ */
+function expressionCost(expr, costModel) {
+    switch (expr.kind) {
+        case 'number':
+            return costModel.num;
+        case 'variable':
+            return costModel.var;
+        case 'binary':
+            const opCost = expr.operator === '/' ? costModel.div :
+                expr.operator === '^' ? costModel.pow :
+                    expr.operator === '*' ? costModel.mul :
+                        costModel.add;
+            return opCost + expressionCost(expr.left, costModel) + expressionCost(expr.right, costModel);
+        case 'unary':
+            return costModel.neg + expressionCost(expr.operand, costModel);
+        case 'call':
+            return costModel.call + expr.args.reduce((sum, arg) => sum + expressionCost(arg, costModel), 0);
+        case 'component':
+            return costModel.component + expressionCost(expr.object, costModel);
+    }
+}
+/**
+ * Serialize an expression to a string for comparison
+ */
+function serializeExpr(expr) {
+    switch (expr.kind) {
+        case 'number':
+            return `N${expr.value}`;
+        case 'variable':
+            return `V${expr.name}`;
+        case 'binary':
+            return `(${serializeExpr(expr.left)}${expr.operator}${serializeExpr(expr.right)})`;
+        case 'unary':
+            return `U${expr.operator}${serializeExpr(expr.operand)}`;
+        case 'call':
+            return `C${expr.name}(${expr.args.map(serializeExpr).join(',')})`;
+        case 'component':
+            return `${serializeExpr(expr.object)}.${expr.component}`;
+    }
+}
+/**
+ * Substitute temp references into an expression (bottom-up)
+ * Looks for subexpressions that match other temps and replaces them
+ */
+function substituteTempRefs(expr, exprToTemp, currentTemp) {
+    // First, recursively substitute in children (bottom-up)
+    let result;
+    switch (expr.kind) {
+        case 'number':
+        case 'variable':
+            result = expr;
+            break;
+        case 'binary': {
+            const left = substituteTempRefs(expr.left, exprToTemp, currentTemp);
+            const right = substituteTempRefs(expr.right, exprToTemp, currentTemp);
+            result = (left === expr.left && right === expr.right)
+                ? expr
+                : { kind: 'binary', operator: expr.operator, left, right };
+            break;
+        }
+        case 'unary': {
+            const operand = substituteTempRefs(expr.operand, exprToTemp, currentTemp);
+            result = (operand === expr.operand)
+                ? expr
+                : { kind: 'unary', operator: expr.operator, operand };
+            break;
+        }
+        case 'call': {
+            const args = expr.args.map(arg => substituteTempRefs(arg, exprToTemp, currentTemp));
+            result = args.every((arg, i) => arg === expr.args[i])
+                ? expr
+                : { kind: 'call', name: expr.name, args };
+            break;
+        }
+        case 'component': {
+            const object = substituteTempRefs(expr.object, exprToTemp, currentTemp);
+            result = (object === expr.object)
+                ? expr
+                : { kind: 'component', object, component: expr.component };
+            break;
+        }
+    }
+    // Then check if the (possibly transformed) expression matches another temp
+    const serialized = serializeExpr(result);
+    const matchingTemp = exprToTemp.get(serialized);
+    if (matchingTemp && matchingTemp !== currentTemp) {
+        return { kind: 'variable', name: matchingTemp };
+    }
+    return result;
+}
+/**
+ * Topologically sort temps so dependencies come first
+ */
+function topologicalSortTemps(temps) {
+    // Find dependencies of each temp
+    const deps = new Map();
+    const tempNames = new Set(temps.keys());
+    function findDeps(expr, found) {
+        if (expr.kind === 'variable' && tempNames.has(expr.name)) {
+            found.add(expr.name);
+        }
+        else if (expr.kind === 'binary') {
+            findDeps(expr.left, found);
+            findDeps(expr.right, found);
+        }
+        else if (expr.kind === 'unary') {
+            findDeps(expr.operand, found);
+        }
+        else if (expr.kind === 'call') {
+            expr.args.forEach(arg => findDeps(arg, found));
+        }
+        else if (expr.kind === 'component') {
+            findDeps(expr.object, found);
+        }
+    }
+    for (const [name, expr] of temps) {
+        const d = new Set();
+        findDeps(expr, d);
+        deps.set(name, d);
+    }
+    // Topological sort using Kahn's algorithm
+    const result = [];
+    const remaining = new Set(temps.keys());
+    const processed = new Set();
+    while (remaining.size > 0) {
+        // Find a temp with no unprocessed dependencies
+        let found = false;
+        for (const name of remaining) {
+            const d = deps.get(name);
+            const hasUnprocessedDep = [...d].some(dep => !processed.has(dep));
+            if (!hasUnprocessedDep) {
+                result.push([name, temps.get(name)]);
+                remaining.delete(name);
+                processed.add(name);
+                found = true;
+                break;
+            }
+        }
+        if (!found) {
+            // Cycle detected - just add remaining in any order
+            for (const name of remaining) {
+                result.push([name, temps.get(name)]);
+            }
+            break;
+        }
+    }
+    return result;
+}
+/**
+ * Post-extraction CSE: find repeated patterns that emerge AFTER temp substitution
+ * e.g., "_tmp22 + _tmp23" appearing multiple times should become its own temp
+ */
+function postExtractionCSE(temps, expressions, minSharedCost, startingTempCounter, costModel) {
+    // Count occurrences of each subexpression
+    const exprCounts = new Map();
+    function countSubexprs(expr) {
+        // Don't count simple expressions
+        if (expr.kind === 'number' || expr.kind === 'variable')
+            return;
+        const serialized = serializeExpr(expr);
+        const cost = expressionCost(expr, costModel);
+        const existing = exprCounts.get(serialized);
+        if (existing) {
+            existing.count++;
+        }
+        else {
+            exprCounts.set(serialized, { count: 1, expr, cost });
+        }
+        // Recurse into children
+        if (expr.kind === 'binary') {
+            countSubexprs(expr.left);
+            countSubexprs(expr.right);
+        }
+        else if (expr.kind === 'unary') {
+            countSubexprs(expr.operand);
+        }
+        else if (expr.kind === 'call') {
+            expr.args.forEach(countSubexprs);
+        }
+        else if (expr.kind === 'component') {
+            countSubexprs(expr.object);
+        }
+    }
+    // Count in all temps and root expressions
+    for (const expr of temps.values()) {
+        countSubexprs(expr);
+    }
+    for (const expr of expressions.values()) {
+        countSubexprs(expr);
+    }
+    // Find subexpressions worth extracting (count >= 2 and cost > threshold)
+    const toExtract = [];
+    for (const [serialized, { count, expr, cost }] of exprCounts) {
+        if (count >= 2 && cost > minSharedCost) {
+            // Skip if it's just a temp reference
+            if (expr.kind === 'variable' && expr.name.startsWith('_tmp'))
+                continue;
+            // Skip trivial expressions (constants, negations of constants)
+            if (isTrivialExpression(expr))
+                continue;
+            toExtract.push({ serialized, expr, cost });
+        }
+    }
+    if (toExtract.length === 0)
+        return;
+    // Sort by cost ASCENDING (extract smaller/cheaper expressions first!)
+    // This is critical because larger patterns contain smaller ones.
+    // If we extract (a+b) first as _tmp100, then later patterns
+    // like (2 * (a+b)) will be serialized as (2 * _tmp100) and won't match.
+    toExtract.sort((a, b) => a.cost - b.cost);
+    // Build a map of existing temp RHS to prevent duplicates
+    const existingTempRHS = new Map();
+    for (const [tempName, expr] of temps) {
+        existingTempRHS.set(serializeExpr(expr), tempName);
+    }
+    // Create temps for repeated expressions
+    let tempCounter = startingTempCounter;
+    const serToTemp = new Map();
+    for (const { serialized, expr } of toExtract) {
+        // Skip if already defined as a temp
+        const existingTemp = existingTempRHS.get(serialized);
+        if (existingTemp) {
+            serToTemp.set(serialized, existingTemp);
+            continue;
+        }
+        // Find unique temp name
+        while (temps.has(`_tmp${tempCounter}`)) {
+            tempCounter++;
+        }
+        const tempName = `_tmp${tempCounter++}`;
+        serToTemp.set(serialized, tempName);
+        temps.set(tempName, expr);
+        existingTempRHS.set(serialized, tempName);
+    }
+    if (serToTemp.size === 0)
+        return;
+    // Substitute new temps into all expressions
+    function substitute(expr) {
+        if (expr.kind === 'number' || expr.kind === 'variable')
+            return expr;
+        const serialized = serializeExpr(expr);
+        const tempName = serToTemp.get(serialized);
+        if (tempName) {
+            return { kind: 'variable', name: tempName };
+        }
+        // Recurse
+        if (expr.kind === 'binary') {
+            const left = substitute(expr.left);
+            const right = substitute(expr.right);
+            return (left === expr.left && right === expr.right)
+                ? expr
+                : { kind: 'binary', operator: expr.operator, left, right };
+        }
+        else if (expr.kind === 'unary') {
+            const operand = substitute(expr.operand);
+            return (operand === expr.operand)
+                ? expr
+                : { kind: 'unary', operator: expr.operator, operand };
+        }
+        else if (expr.kind === 'call') {
+            const args = expr.args.map(substitute);
+            return args.every((arg, i) => arg === expr.args[i])
+                ? expr
+                : { kind: 'call', name: expr.name, args };
+        }
+        else if (expr.kind === 'component') {
+            const object = substitute(expr.object);
+            return (object === expr.object)
+                ? expr
+                : { kind: 'component', object, component: expr.component };
+        }
+        return expr;
+    }
+    // Substitute in ALL temps, including newly created ones
+    // But skip substituting a temp with itself (self-reference)
+    for (const [tempName, expr] of temps) {
+        // Create a substitute function that won't replace with the current temp
+        const subWithoutSelf = (e) => {
+            if (e.kind === 'number' || e.kind === 'variable')
+                return e;
+            const serialized = serializeExpr(e);
+            const targetTemp = serToTemp.get(serialized);
+            // Don't substitute if it would create self-reference
+            if (targetTemp && targetTemp !== tempName) {
+                return { kind: 'variable', name: targetTemp };
+            }
+            if (e.kind === 'binary') {
+                const left = subWithoutSelf(e.left);
+                const right = subWithoutSelf(e.right);
+                return (left === e.left && right === e.right)
+                    ? e
+                    : { kind: 'binary', operator: e.operator, left, right };
+            }
+            else if (e.kind === 'unary') {
+                const operand = subWithoutSelf(e.operand);
+                return (operand === e.operand)
+                    ? e
+                    : { kind: 'unary', operator: e.operator, operand };
+            }
+            else if (e.kind === 'call') {
+                const args = e.args.map(subWithoutSelf);
+                return args.every((arg, i) => arg === e.args[i])
+                    ? e
+                    : { kind: 'call', name: e.name, args };
+            }
+            else if (e.kind === 'component') {
+                const object = subWithoutSelf(e.object);
+                return (object === e.object)
+                    ? e
+                    : { kind: 'component', object, component: e.component };
+            }
+            return e;
+        };
+        temps.set(tempName, subWithoutSelf(expr));
+    }
+    // Substitute in root expressions
+    for (const [rootId, expr] of expressions) {
+        expressions.set(rootId, substitute(expr));
+    }
+    // Re-sort temps topologically
+    const sorted = topologicalSortTemps(temps);
+    temps.clear();
+    for (const [name, expr] of sorted) {
+        temps.set(name, expr);
+    }
+    // Inline temps that are now used only once (after all substitutions)
+    // This is critical because postExtractionCSE may have created temps
+    // that turned out to be used only once after substitution
+    const usageCounts = new Map();
+    function countUsage(expr) {
+        if (expr.kind === 'variable' && expr.name.startsWith('_tmp')) {
+            usageCounts.set(expr.name, (usageCounts.get(expr.name) ?? 0) + 1);
+        }
+        else if (expr.kind === 'binary') {
+            countUsage(expr.left);
+            countUsage(expr.right);
+        }
+        else if (expr.kind === 'unary') {
+            countUsage(expr.operand);
+        }
+        else if (expr.kind === 'call') {
+            expr.args.forEach(countUsage);
+        }
+        else if (expr.kind === 'component') {
+            countUsage(expr.object);
+        }
+    }
+    for (const expr of temps.values())
+        countUsage(expr);
+    for (const expr of expressions.values())
+        countUsage(expr);
+    // Find temps to inline (used 0 or 1 times)
+    const toInline = new Set();
+    for (const [name] of temps) {
+        const count = usageCounts.get(name) ?? 0;
+        if (count <= 1)
+            toInline.add(name);
+    }
+    if (toInline.size > 0) {
+        function inlineTemps(expr) {
+            if (expr.kind === 'variable' && toInline.has(expr.name)) {
+                const tempExpr = temps.get(expr.name);
+                return tempExpr ? inlineTemps(tempExpr) : expr;
+            }
+            else if (expr.kind === 'binary') {
+                const left = inlineTemps(expr.left);
+                const right = inlineTemps(expr.right);
+                return (left === expr.left && right === expr.right) ? expr
+                    : { kind: 'binary', operator: expr.operator, left, right };
+            }
+            else if (expr.kind === 'unary') {
+                const operand = inlineTemps(expr.operand);
+                return (operand === expr.operand) ? expr
+                    : { kind: 'unary', operator: expr.operator, operand };
+            }
+            else if (expr.kind === 'call') {
+                const args = expr.args.map(inlineTemps);
+                return args.every((a, i) => a === expr.args[i]) ? expr
+                    : { kind: 'call', name: expr.name, args };
+            }
+            else if (expr.kind === 'component') {
+                const object = inlineTemps(expr.object);
+                return (object === expr.object) ? expr
+                    : { kind: 'component', object, component: expr.component };
+            }
+            return expr;
+        }
+        // Inline in remaining temps
+        for (const [name, expr] of temps) {
+            if (!toInline.has(name)) {
+                temps.set(name, inlineTemps(expr));
+            }
+        }
+        // Inline in root expressions
+        for (const [rootId, expr] of expressions) {
+            expressions.set(rootId, inlineTemps(expr));
+        }
+        // Remove inlined temps
+        for (const name of toInline)
+            temps.delete(name);
+        // Re-sort topologically after inlining (inlining may have changed dependencies)
+        const finalSorted = topologicalSortTemps(temps);
+        temps.clear();
+        for (const [name, expr] of finalSorted) {
+            temps.set(name, expr);
+        }
+    }
+}
+/**
+ * Detect pairs of temps that are negatives of each other:
+ * e.g., _tmp1 = k * (a - b) and _tmp2 = k * (b - a)
+ * These can be merged: keep _tmp1, replace _tmp2 with -_tmp1
+ */
+function mergeNegativePairs(temps, expressions) {
+    // Build a map of "canonical subtraction form" -> temp name
+    // For k * (a - b), the canonical form is [k, a, b] sorted by serialization
+    const subPatterns = new Map();
+    for (const [tempName, expr] of temps) {
+        const pattern = extractSubtractionPattern(expr);
+        if (!pattern)
+            continue;
+        const { left, right, coefficient } = pattern;
+        // Canonical form: sort left and right alphabetically
+        const leftSer = serializeExpr(left);
+        const rightSer = serializeExpr(right);
+        const coeffSer = coefficient ? serializeExpr(coefficient) : '';
+        let canonKey;
+        let isNegated;
+        if (leftSer < rightSer) {
+            canonKey = `${coeffSer}:(${leftSer})-(${rightSer})`;
+            isNegated = false;
+        }
+        else {
+            canonKey = `${coeffSer}:(${rightSer})-(${leftSer})`;
+            isNegated = true;
+        }
+        const existing = subPatterns.get(canonKey);
+        if (existing) {
+            // Found a pair! One is the negative of the other
+            // Keep the non-negated one (or the first one if both are same)
+            if (existing.isNegated && !isNegated) {
+                // Current one is better, replace existing
+                replaceTempWithNegation(temps, expressions, existing.tempName, tempName);
+                subPatterns.set(canonKey, { tempName, isNegated, coefficient });
+            }
+            else if (!existing.isNegated && isNegated) {
+                // Existing is better, replace current
+                replaceTempWithNegation(temps, expressions, tempName, existing.tempName);
+            }
+            // If both have same negation status, do nothing
+        }
+        else {
+            subPatterns.set(canonKey, { tempName, isNegated, coefficient });
+        }
+    }
+}
+/**
+ * Extract subtraction pattern from expression:
+ * Returns { left, right, coefficient } for patterns like:
+ * - (a - b) -> { left: a, right: b, coefficient: null }
+ * - k * (a - b) -> { left: a, right: b, coefficient: k }
+ */
+function extractSubtractionPattern(expr) {
+    // Direct subtraction: (a - b)
+    if (expr.kind === 'binary' && expr.operator === '-') {
+        return { left: expr.left, right: expr.right, coefficient: null };
+    }
+    // Multiplication with subtraction: k * (a - b) or (a - b) * k
+    if (expr.kind === 'binary' && expr.operator === '*') {
+        if (expr.right.kind === 'binary' && expr.right.operator === '-') {
+            return { left: expr.right.left, right: expr.right.right, coefficient: expr.left };
+        }
+        if (expr.left.kind === 'binary' && expr.left.operator === '-') {
+            return { left: expr.left.left, right: expr.left.right, coefficient: expr.right };
+        }
+    }
+    return null;
+}
+/**
+ * Replace all uses of oldTemp with -newTemp, then delete oldTemp
+ */
+function replaceTempWithNegation(temps, expressions, oldTemp, newTemp) {
+    // Create negation expression
+    const negExpr = {
+        kind: 'unary',
+        operator: '-',
+        operand: { kind: 'variable', name: newTemp }
+    };
+    // Helper to replace references
+    function replaceRefs(expr) {
+        if (expr.kind === 'variable' && expr.name === oldTemp) {
+            return negExpr;
+        }
+        if (expr.kind === 'binary') {
+            const left = replaceRefs(expr.left);
+            const right = replaceRefs(expr.right);
+            return (left === expr.left && right === expr.right) ? expr
+                : { kind: 'binary', operator: expr.operator, left, right };
+        }
+        if (expr.kind === 'unary') {
+            const operand = replaceRefs(expr.operand);
+            return (operand === expr.operand) ? expr
+                : { kind: 'unary', operator: expr.operator, operand };
+        }
+        if (expr.kind === 'call') {
+            const args = expr.args.map(replaceRefs);
+            return args.every((a, i) => a === expr.args[i]) ? expr
+                : { kind: 'call', name: expr.name, args };
+        }
+        if (expr.kind === 'component') {
+            const object = replaceRefs(expr.object);
+            return (object === expr.object) ? expr
+                : { kind: 'component', object, component: expr.component };
+        }
+        return expr;
+    }
+    // Replace in all temps (except the one we're deleting)
+    for (const [name, expr] of temps) {
+        if (name !== oldTemp) {
+            temps.set(name, replaceRefs(expr));
+        }
+    }
+    // Replace in root expressions
+    for (const [rootId, expr] of expressions) {
+        expressions.set(rootId, replaceRefs(expr));
+    }
+    // Delete the old temp - its uses have been replaced with -newTemp
+    temps.delete(oldTemp);
+}
+/**
+ * Normalize patterns like a + -1 * b to a - b
+ * This cleans up cases where e-graph extraction picked the wrong form
+ */
+function normalizeAddNegMul(temps, expressions) {
+    function normalize(expr) {
+        if (expr.kind === 'number' || expr.kind === 'variable')
+            return expr;
+        // First normalize children
+        if (expr.kind === 'binary') {
+            const left = normalize(expr.left);
+            const right = normalize(expr.right);
+            // Check for a + (-1 * b) pattern: convert to a - b
+            if (expr.operator === '+' && right.kind === 'binary' && right.operator === '*') {
+                if (right.left.kind === 'number' && right.left.value === -1) {
+                    return { kind: 'binary', operator: '-', left, right: normalize(right.right) };
+                }
+                if (right.right.kind === 'number' && right.right.value === -1) {
+                    return { kind: 'binary', operator: '-', left, right: normalize(right.left) };
+                }
+            }
+            // Check for (-1 * b) + a pattern: convert to a - b
+            if (expr.operator === '+' && left.kind === 'binary' && left.operator === '*') {
+                if (left.left.kind === 'number' && left.left.value === -1) {
+                    return { kind: 'binary', operator: '-', left: right, right: normalize(left.right) };
+                }
+                if (left.right.kind === 'number' && left.right.value === -1) {
+                    return { kind: 'binary', operator: '-', left: right, right: normalize(left.left) };
+                }
+            }
+            return (left === expr.left && right === expr.right)
+                ? expr
+                : { kind: 'binary', operator: expr.operator, left, right };
+        }
+        if (expr.kind === 'unary') {
+            const operand = normalize(expr.operand);
+            return (operand === expr.operand) ? expr : { kind: 'unary', operator: expr.operator, operand };
+        }
+        if (expr.kind === 'call') {
+            const args = expr.args.map(normalize);
+            return args.every((a, i) => a === expr.args[i]) ? expr : { kind: 'call', name: expr.name, args };
+        }
+        if (expr.kind === 'component') {
+            const object = normalize(expr.object);
+            return (object === expr.object) ? expr : { kind: 'component', object, component: expr.component };
+        }
+        return expr;
+    }
+    // Normalize all temps
+    for (const [name, expr] of temps) {
+        temps.set(name, normalize(expr));
+    }
+    // Normalize root expressions
+    for (const [rootId, expr] of expressions) {
+        expressions.set(rootId, normalize(expr));
+    }
+}