gradient-script 0.2.0 → 0.3.0

package/README.md

@@ -12,6 +12,8 @@
 
 GradientScript is a source-to-source compiler that automatically generates gradient functions from your mathematical code. Unlike numerical AD frameworks (JAX, PyTorch), it produces clean, human-readable gradient formulas you can inspect, optimize, and integrate directly into your codebase.
 
+It's perfect for LLM usage where the LLM can verify existing gradients or construct gradients you require with less risks of making errors.
+
 ## Why GradientScript?
 
 - **From real code to gradients**: Write natural math code, get symbolic derivatives
@@ -40,7 +42,7 @@ function distance(u: Vec2, v: Vec2): number {
 }
 ```
 
-Convert it to GradientScript by marking what you need gradients for:
+Convert it to GradientScript (realistically, let your LLM convert it giving it a reference here - and/or free usage of the CLI) by marking what you need gradients for:
 
 ```typescript
 // distance.gs

package/dist/cli.js

@@ -6,6 +6,68 @@ import { computeFunctionGradients } from './dsl/Differentiation.js';
 import { generateComplete } from './dsl/CodeGen.js';
 import { analyzeGuards, formatGuardWarnings } from './dsl/Guards.js';
 import { ParseError, formatParseError } from './dsl/Errors.js';
+import { GradientChecker, formatGradCheckResult } from './dsl/GradientChecker.js';
+import { Types } from './dsl/Types.js';
+/**
+ * Generate random test points for gradient verification.
+ * Uses multiple test points to catch errors at different values.
+ */
+function generateTestPoints(func, env) {
+    const testPoints = [];
+    // Generate 3 different test points with varying scales
+    const scales = [1.0, 0.1, 10.0];
+    for (const scale of scales) {
+        const point = new Map();
+        for (const param of func.parameters) {
+            const paramType = env.getOrThrow(param.name);
+            if (Types.isScalar(paramType)) {
+                // Random scalar in range [-scale, scale], avoid zero
+                point.set(param.name, (Math.random() * 2 - 1) * scale + 0.1 * scale);
+            }
+            else {
+                // Structured type - get components
+                const struct = {};
+                for (const comp of paramType.components) {
+                    struct[comp] = (Math.random() * 2 - 1) * scale + 0.1 * scale;
+                }
+                point.set(param.name, struct);
+            }
+        }
+        testPoints.push(point);
+    }
+    return testPoints;
+}
+/**
+ * Verify gradients for a function using numerical differentiation.
+ * Returns true if all gradients pass, false otherwise.
+ */
+function verifyGradients(func, gradients, env) {
+    const checker = new GradientChecker(1e-5, 1e-4);
+    const testPoints = generateTestPoints(func, env);
+    let allPassed = true;
+    for (let i = 0; i < testPoints.length; i++) {
+        const result = checker.check(func, gradients, env, testPoints[i]);
+        if (!result.passed) {
+            if (allPassed) {
+                // First failure - print header (as comment for valid output)
+                console.error(`// Gradient verification FAILED for "${func.name}":`);
+            }
+            // Prefix each line with // so output remains valid code
+            const formattedResult = formatGradCheckResult(result, func.name)
+                .split('\n')
+                .map(line => '// ' + line)
+                .join('\n');
+            console.error(`//   Test point ${i + 1}: ${formattedResult}`);
+            allPassed = false;
+        }
+    }
+    if (allPassed) {
+        const result = checker.check(func, gradients, env, testPoints[0]);
+        // Prefix with // so output is valid code
+        console.error('// ' + formatGradCheckResult(result, func.name));
+    }
+    return allPassed;
+}
 function printUsage() {
     console.log(`
 GradientScript - Symbolic Differentiation for Structured Types
@@ -16,7 +78,7 @@ Usage:
 Options:
   --format <format>     Output format: typescript (default), javascript, python, csharp
   --no-simplify         Disable gradient simplification
-  --no-cse              Disable common subexpression elimination
+  --no-cse              Disable optimization (e-graph CSE)
   --no-comments         Omit comments in generated code
   --guards              Emit runtime guards for division by zero (experimental)
   --epsilon <value>     Epsilon value for guards (default: 1e-10)
@@ -45,6 +107,9 @@ For more information and examples:
 
   README (raw, LLM-friendly):
   https://raw.githubusercontent.com/mfagerlund/gradient-script/main/README.md
+
+  LLM Optimization Guide (for AI agents writing .gs files):
+  https://raw.githubusercontent.com/mfagerlund/gradient-script/main/docs/LLM-OPTIMIZATION-GUIDE.md
   `.trim());
 }
 function main() {
@@ -64,6 +129,7 @@ function main() {
         simplify: true,
         cse: true
     };
+    let skipVerify = false;
     for (let i = 1; i < args.length; i++) {
         const arg = args[i];
         if (arg === '--format') {
@@ -138,13 +204,20 @@ function main() {
             process.exit(1);
         }
         const outputs = [];
+        let hasVerificationFailure = false;
         program.functions.forEach((func, index) => {
             const env = inferFunction(func);
             const gradients = computeFunctionGradients(func, env);
+            // MANDATORY gradient verification
+            const verified = verifyGradients(func, gradients, env);
+            if (!verified) {
+                hasVerificationFailure = true;
+            }
             const guardAnalysis = analyzeGuards(func);
             if (guardAnalysis.hasIssues) {
-                console.error('Function "' + func.name + '" may have edge cases:');
-                console.error(formatGuardWarnings(guardAnalysis));
+                // Format warnings as comments so output remains valid code even if stderr is captured
+                console.error('// Function "' + func.name + '" may have edge cases:');
+                console.error(formatGuardWarnings(guardAnalysis, true));
             }
             const perFunctionOptions = { ...options };
             if (index > 0 && perFunctionOptions.includeComments !== false) {
@@ -153,6 +226,10 @@ function main() {
             const code = generateComplete(func, gradients, env, perFunctionOptions);
             outputs.push(code);
         });
+        if (hasVerificationFailure) {
+            console.error('// ERROR: Gradient verification failed. Output may contain incorrect gradients!');
+            process.exit(1);
+        }
         console.log(outputs.join('\n\n'));
     }
     catch (err) {

package/dist/dsl/CodeGen.d.ts

@@ -56,7 +56,7 @@ export declare class ExpressionCodeGen {
  */
 export declare function generateGradientFunction(func: FunctionDef, gradients: GradientResult, env: TypeEnv, options?: CodeGenOptions): string;
 /**
- * Generate the original forward function
+ * Generate the original forward function (with optional e-graph optimization)
  */
 export declare function generateForwardFunction(func: FunctionDef, options?: CodeGenOptions): string;
 /**

package/dist/dsl/CodeGen.js

@@ -2,11 +2,12 @@
  * Code generation for GradientScript DSL
  * Generates TypeScript/JavaScript code with gradient functions
  */
-import { simplifyGradients } from './Simplify.js';
+import { simplifyGradients, simplifyPostCSE } from './Simplify.js';
 import { ExpressionTransformer } from './ExpressionTransformer.js';
-import { eliminateCommonSubexpressionsStructured } from './CSE.js';
+import { optimizeWithEGraph } from './egraph/index.js';
 import { CodeGenError } from './Errors.js';
 import { serializeExpression } from './ExpressionUtils.js';
+import { inlineExpression } from './Inliner.js';
 function capitalize(str) {
     return str.charAt(0).toUpperCase() + str.slice(1);
 }
@@ -79,10 +80,12 @@ export class ExpressionCodeGen {
                         return left;
                     }
                     else if (exponent === 2) {
-                        return `${left} * ${left}`;
+                        // Wrap in parens because we're changing ^ (precedence 3) to * (precedence 2)
+                        // Without parens, a / b^2 would become a / b * b instead of a / (b * b)
+                        return `(${left} * ${left})`;
                     }
                     else if (exponent === 3) {
-                        return `${left} * ${left} * ${left}`;
+                        return `(${left} * ${left} * ${left})`;
                     }
                 }
             }
@@ -300,15 +303,162 @@ export function generateGradientFunction(func, gradients, env, options = {}) {
     // Forward pass - compute intermediate variables
     // Track which expressions are already computed for CSE reuse
     const forwardExpressionMap = new Map();
+    // Build substitution map for inlining (to match gradient expressions which are fully inlined)
+    const substitutionMap = new Map();
+    for (const stmt of func.body) {
+        if (stmt.kind === 'assignment') {
+            substitutionMap.set(stmt.variable, stmt.expression);
+        }
+    }
+    // Collect forward variable names and expressions for optimization
+    const forwardVars = new Set();
+    const forwardVarExprs = new Map();
+    for (const stmt of func.body) {
+        if (stmt.kind === 'assignment') {
+            forwardVars.add(stmt.variable);
+            forwardVarExprs.set(stmt.variable, stmt.expression);
+        }
+    }
+    // Optimize forward expressions with e-graph
+    let optimizedForwardExprs = forwardVarExprs;
+    let forwardCseTemps = new Map();
+    if (options.cse !== false && forwardVarExprs.size > 0) {
+        const forOptimizer = new Map();
+        forOptimizer.set('_forward', forwardVarExprs);
+        const forwardResult = optimizeWithEGraph(forOptimizer, { verbose: false });
+        // Rename forward temps to avoid conflicts with gradient temps (use _fwd prefix)
+        const rawTemps = forwardResult.intermediates;
+        const renameMap = new Map();
+        for (const oldName of rawTemps.keys()) {
+            const newName = oldName.replace('_tmp', '_fwd');
+            renameMap.set(oldName, newName);
+        }
+        // Apply renaming to temp definitions
+        for (const [oldName, expr] of rawTemps) {
+            const newName = renameMap.get(oldName);
+            forwardCseTemps.set(newName, renameTempRefs(expr, renameMap));
+        }
+        // Apply renaming to optimized expressions
+        optimizedForwardExprs = new Map();
+        for (const [varName, expr] of (forwardResult.gradients.get('_forward') || forwardVarExprs)) {
+            optimizedForwardExprs.set(varName, renameTempRefs(expr, renameMap));
+        }
+    }
+    // Helper to rename temp references in an expression
+    function renameTempRefs(expr, renameMap) {
+        if (expr.kind === 'variable') {
+            const newName = renameMap.get(expr.name);
+            return newName ? { kind: 'variable', name: newName } : expr;
+        }
+        else if (expr.kind === 'binary') {
+            return {
+                kind: 'binary',
+                operator: expr.operator,
+                left: renameTempRefs(expr.left, renameMap),
+                right: renameTempRefs(expr.right, renameMap)
+            };
+        }
+        else if (expr.kind === 'unary') {
+            return {
+                kind: 'unary',
+                operator: expr.operator,
+                operand: renameTempRefs(expr.operand, renameMap)
+            };
+        }
+        else if (expr.kind === 'call') {
+            return {
+                kind: 'call',
+                name: expr.name,
+                args: expr.args.map(a => renameTempRefs(a, renameMap))
+            };
+        }
+        else if (expr.kind === 'component') {
+            return {
+                kind: 'component',
+                object: renameTempRefs(expr.object, renameMap),
+                component: expr.component
+            };
+        }
+        return expr;
+    }
+    // Helper to find which forward vars an expression depends on
+    function findForwardVarDeps(expr) {
+        const deps = new Set();
+        function visit(e) {
+            if (e.kind === 'variable' && forwardVars.has(e.name)) {
+                deps.add(e.name);
+            }
+            else if (e.kind === 'binary') {
+                visit(e.left);
+                visit(e.right);
+            }
+            else if (e.kind === 'unary') {
+                visit(e.operand);
+            }
+            else if (e.kind === 'call') {
+                e.args.forEach(visit);
+            }
+            else if (e.kind === 'component') {
+                visit(e.object);
+            }
+        }
+        visit(expr);
+        return deps;
+    }
+    // Track which CSE temps need to be emitted after which forward var
+    const fwdTempAfterVar = new Map();
+    const fwdTempsBeforeAny = [];
+    for (const [tempName, tempExpr] of forwardCseTemps) {
+        const deps = findForwardVarDeps(tempExpr);
+        if (deps.size === 0) {
+            fwdTempsBeforeAny.push({ name: tempName, expr: tempExpr });
+        }
+        else {
+            let lastDep = '';
+            for (const stmt of func.body) {
+                if (stmt.kind === 'assignment' && deps.has(stmt.variable)) {
+                    lastDep = stmt.variable;
+                }
+            }
+            if (lastDep) {
+                if (!fwdTempAfterVar.has(lastDep)) {
+                    fwdTempAfterVar.set(lastDep, []);
+                }
+                fwdTempAfterVar.get(lastDep).push({ name: tempName, expr: tempExpr });
+            }
+        }
+    }
+    // Emit forward CSE temps that don't depend on forward vars
+    for (const { name: tempName, expr } of fwdTempsBeforeAny) {
+        const code = codegen.generate(expr);
+        if (format === 'typescript' || format === 'javascript') {
+            lines.push(`  const ${tempName} = ${code};`);
+        }
+        else if (format === 'python') {
+            lines.push(`  ${tempName} = ${code}`);
+        }
+        else if (format === 'csharp') {
+            lines.push(`    ${csharpFloatType} ${tempName} = ${code};`);
+        }
+    }
+    // Generate forward variable assignments with interleaved temps
     for (const stmt of func.body) {
         if (stmt.kind === 'assignment') {
             const varName = stmt.variable;
-            const generatedExpr = codegen.generate(stmt.expression);
+            const expr = optimizedForwardExprs.get(varName) || stmt.expression;
+            const generatedExpr = codegen.generate(expr);
             if (shouldTrackForForwardReuse(stmt.expression)) {
+                // Register the original expression
                 const exprKey = serializeExpression(stmt.expression);
                 if (!forwardExpressionMap.has(exprKey)) {
                     forwardExpressionMap.set(exprKey, varName);
                 }
+                // Also register the fully inlined form (this is what gradient expressions will have)
+                const inlinedExpr = inlineExpression(stmt.expression, substitutionMap);
+                const inlinedKey = serializeExpression(inlinedExpr);
+                if (inlinedKey !== exprKey && !forwardExpressionMap.has(inlinedKey)) {
+                    forwardExpressionMap.set(inlinedKey, varName);
+                }
             }
             if (format === 'typescript' || format === 'javascript') {
                 lines.push(`  const ${varName} = ${generatedExpr};`);
@@ -319,6 +469,20 @@ export function generateGradientFunction(func, gradients, env, options = {}) {
             else if (format === 'csharp') {
                 lines.push(`    ${csharpFloatType} ${varName} = ${generatedExpr};`);
             }
+            // Emit any forward CSE temps that depend on this var
+            const tempsForVar = fwdTempAfterVar.get(varName) || [];
+            for (const { name: tempName, expr: tempExpr } of tempsForVar) {
+                const tempCode = codegen.generate(tempExpr);
+                if (format === 'typescript' || format === 'javascript') {
+                    lines.push(`  const ${tempName} = ${tempCode};`);
+                }
+                else if (format === 'python') {
+                    lines.push(`  ${tempName} = ${tempCode}`);
+                }
+                else if (format === 'csharp') {
+                    lines.push(`    ${csharpFloatType} ${tempName} = ${tempCode};`);
+                }
+            }
         }
     }
     // Compute output value - reuse forward pass variables if possible
@@ -372,77 +536,71 @@ export function generateGradientFunction(func, gradients, env, options = {}) {
     if (includeComments) {
         lines.push(`  ${comment} Gradients`);
     }
-    // Apply CSE if requested
+    // Apply e-graph optimization (CSE + algebraic simplification)
     const shouldApplyCSE = options.cse !== false; // Default to true
-    const cseIntermediates = new Map();
+    let cseIntermediates = new Map();
     if (shouldApplyCSE) {
-        // Collect all gradient expressions for CSE analysis
+        // Collect all gradient components into a single map for global optimization
+        // Include both structured and scalar gradients (scalar uses 'value' as component)
+        const allGradientComponents = new Map();
         for (const [paramName, gradient] of gradientsToUse.gradients.entries()) {
             if (isStructuredGradient(gradient)) {
-                const cseResult = eliminateCommonSubexpressionsStructured(gradient.components);
-                // Merge intermediates
-                for (const [name, expr] of cseResult.intermediates.entries()) {
-                    cseIntermediates.set(name, expr);
-                }
-                // Update gradient components with CSE-simplified versions
-                gradient.components = cseResult.components;
-            }
-        }
-        // Generate intermediate variables from CSE
-        if (cseIntermediates.size > 0) {
-            // Check if we should emit guards (opt-in)
-            const shouldEmitGuards = options.emitGuards === true;
-            const epsilon = options.epsilon || 1e-10;
-            // Identify potential denominators (sum of squares patterns)
-            const denominatorVars = new Set();
-            for (const [varName, expr] of cseIntermediates.entries()) {
-                const code = codegen.generate(expr);
-                // Check if this looks like a denominator (contains + and squared terms)
-                if (code.includes('+') && (code.includes('* ') || code.includes('Math.pow'))) {
-                    denominatorVars.add(varName);
-                }
+                allGradientComponents.set(paramName, gradient.components);
             }
-            for (const [varName, expr] of cseIntermediates.entries()) {
-                const code = codegen.generate(expr);
-                if (format === 'typescript' || format === 'javascript') {
-                    lines.push(`  const ${varName} = ${code};`);
-                }
-                else if (format === 'python') {
-                    lines.push(`  ${varName} = ${code}`);
-                }
-                else if (format === 'csharp') {
-                    lines.push(`    ${csharpFloatType} ${varName} = ${code};`);
-                }
+            else {
+                // Scalar gradient - wrap as a single 'value' component
+                allGradientComponents.set(paramName, new Map([['value', gradient]]));
             }
-            // Emit epsilon guard if needed
-            if (shouldEmitGuards && denominatorVars.size > 0) {
-                lines.push('');
-                if (includeComments) {
-                    lines.push(`  ${comment} Guard against division by zero`);
+        }
+        // Run e-graph optimization globally across ALL gradient expressions
+        const globalCSE = optimizeWithEGraph(allGradientComponents, { verbose: false });
+        cseIntermediates = globalCSE.intermediates;
+        // Update gradient components with optimized versions
+        for (const [paramName, simplifiedComponents] of globalCSE.gradients.entries()) {
+            const gradient = gradientsToUse.gradients.get(paramName);
+            if (gradient && isStructuredGradient(gradient)) {
+                gradient.components = simplifiedComponents;
+            }
+            else {
+                // Scalar gradient - unwrap from 'value' component
+                const optimizedExpr = simplifiedComponents.get('value');
+                if (optimizedExpr) {
+                    gradientsToUse.gradients.set(paramName, optimizedExpr);
                 }
-                for (const denom of denominatorVars) {
-                    if (format === 'typescript' || format === 'javascript') {
-                        lines.push(`  if (Math.abs(${denom}) < ${epsilon}) {`);
-                        lines.push(`    ${comment} Return zero gradients for degenerate case`);
-                        // Emit zero gradient structure
-                        const zeroGrads = [];
-                        for (const [paramName, gradient] of gradientsToUse.gradients.entries()) {
-                            if (isStructuredGradient(gradient)) {
-                                const components = Array.from(gradient.components.keys());
-                                const zeroStruct = components.map(c => `${c}: 0`).join(', ');
-                                zeroGrads.push(`d${paramName}: { ${zeroStruct} }`);
-                            }
-                            else {
-                                zeroGrads.push(`d${paramName}: 0`);
-                            }
-                        }
-                        lines.push(`    return { value, ${zeroGrads.join(', ')} };`);
-                        lines.push(`  }`);
-                    }
+            }
+        }
+        // Post-CSE simplification: apply rules that were skipped to avoid CSE interference
+        // Specifically: a + a → 2 * a (now safe because temps have been extracted)
+        for (const [varName, expr] of cseIntermediates) {
+            cseIntermediates.set(varName, simplifyPostCSE(expr));
+        }
+        for (const [paramName, gradient] of gradientsToUse.gradients.entries()) {
+            if (isStructuredGradient(gradient)) {
+                for (const [comp, expr] of gradient.components.entries()) {
+                    gradient.components.set(comp, simplifyPostCSE(expr));
                 }
             }
-            lines.push('');
+            else {
+                // Scalar gradient - apply post-CSE simplification
+                gradientsToUse.gradients.set(paramName, simplifyPostCSE(gradient));
+            }
+        }
+    }
+    // Generate CSE intermediate variables
+    if (cseIntermediates.size > 0) {
+        for (const [varName, expr] of cseIntermediates.entries()) {
+            const code = codegen.generate(expr);
+            if (format === 'typescript' || format === 'javascript') {
+                lines.push(`  const ${varName} = ${code};`);
+            }
+            else if (format === 'python') {
+                lines.push(`  ${varName} = ${code}`);
+            }
+            else if (format === 'csharp') {
+                lines.push(`    ${csharpFloatType} ${varName} = ${code};`);
+            }
         }
+        lines.push('');
     }
     for (const [paramName, gradient] of gradientsToUse.gradients.entries()) {
         // Use shorter names: du, dv instead of grad_u, grad_v
@@ -530,11 +688,12 @@ export function generateGradientFunction(func, gradients, env, options = {}) {
     return lines.join('\n');
 }
 /**
- * Generate the original forward function
+ * Generate the original forward function (with optional e-graph optimization)
  */
 export function generateForwardFunction(func, options = {}) {
     const format = options.format || 'typescript';
     const csharpFloatType = options.csharpFloatType || 'float';
+    const shouldOptimize = options.cse !== false; // Optimize by default
     const codegen = new ExpressionCodeGen(format, csharpFloatType);
     const lines = [];
     // Function signature
@@ -549,28 +708,127 @@ export function generateForwardFunction(func, options = {}) {
         const floatType = csharpFloatType;
         const params = func.parameters.map(p => {
             if (p.paramType && p.paramType.components) {
-                // Structured parameter - create a struct type name
                 return `${capitalize(p.name)}Struct ${p.name}`;
             }
             return `${floatType} ${p.name}`;
         }).join(', ');
-        // Generate struct definitions for structured parameters first (we'll prepend them later)
         lines.push(`public static ${floatType} ${capitalize(func.name)}(${params})`);
         lines.push('{');
     }
-    // Body
+    // Collect all forward variable names (for dependency tracking)
+    const forwardVars = new Set();
+    for (const stmt of func.body) {
+        if (stmt.kind === 'assignment') {
+            forwardVars.add(stmt.variable);
+        }
+    }
+    // Collect expressions for optimization
+    const varExpressions = new Map();
+    for (const stmt of func.body) {
+        if (stmt.kind === 'assignment') {
+            varExpressions.set(stmt.variable, stmt.expression);
+        }
+    }
+    // Optimize with e-graph if enabled
+    let optimizedExprs = varExpressions;
+    let cseTemps = new Map();
+    if (shouldOptimize && varExpressions.size > 0) {
+        const forOptimizer = new Map();
+        forOptimizer.set('_forward', varExpressions);
+        const result = optimizeWithEGraph(forOptimizer, { verbose: false });
+        cseTemps = result.intermediates;
+        optimizedExprs = result.gradients.get('_forward') || varExpressions;
+    }
+    // Helper to find which forward vars an expression depends on
+    function findForwardVarDeps(expr) {
+        const deps = new Set();
+        function visit(e) {
+            if (e.kind === 'variable' && forwardVars.has(e.name)) {
+                deps.add(e.name);
+            }
+            else if (e.kind === 'binary') {
+                visit(e.left);
+                visit(e.right);
+            }
+            else if (e.kind === 'unary') {
+                visit(e.operand);
+            }
+            else if (e.kind === 'call') {
+                e.args.forEach(visit);
+            }
+            else if (e.kind === 'component') {
+                visit(e.object);
+            }
+        }
+        visit(expr);
+        return deps;
+    }
+    // Track which temps need to be emitted after which forward var
+    const tempAfterVar = new Map();
+    const tempsEmittedBeforeAny = [];
+    for (const [tempName, tempExpr] of cseTemps) {
+        const deps = findForwardVarDeps(tempExpr);
+        if (deps.size === 0) {
+            // Temp only depends on params - emit before any forward vars
+            tempsEmittedBeforeAny.push({ name: tempName, expr: tempExpr });
+        }
+        else {
+            // Find the last forward var this temp depends on
+            let lastDep = '';
+            for (const stmt of func.body) {
+                if (stmt.kind === 'assignment' && deps.has(stmt.variable)) {
+                    lastDep = stmt.variable;
+                }
+            }
+            if (lastDep) {
+                if (!tempAfterVar.has(lastDep)) {
+                    tempAfterVar.set(lastDep, []);
+                }
+                tempAfterVar.get(lastDep).push({ name: tempName, expr: tempExpr });
+            }
+        }
+    }
+    // Emit temps that don't depend on forward vars
+    for (const { name: tempName, expr } of tempsEmittedBeforeAny) {
+        const code = codegen.generate(expr);
+        if (format === 'typescript' || format === 'javascript') {
+            lines.push(`  const ${tempName} = ${code};`);
+        }
+        else if (format === 'python') {
+            lines.push(`  ${tempName} = ${code}`);
+        }
+        else if (format === 'csharp') {
+            lines.push(`    ${csharpFloatType} ${tempName} = ${code};`);
+        }
+    }
+    // Generate variable assignments with interleaved temps
     for (const stmt of func.body) {
         if (stmt.kind === 'assignment') {
             const varName = stmt.variable;
-            const expr = codegen.generate(stmt.expression);
+            const expr = optimizedExprs.get(varName) || stmt.expression;
+            const code = codegen.generate(expr);
             if (format === 'typescript' || format === 'javascript') {
-                lines.push(`  const ${varName} = ${expr};`);
+                lines.push(`  const ${varName} = ${code};`);
             }
             else if (format === 'python') {
-                lines.push(`  ${varName} = ${expr}`);
+                lines.push(`  ${varName} = ${code}`);
             }
             else if (format === 'csharp') {
-                lines.push(`    ${csharpFloatType} ${varName} = ${expr};`);
+                lines.push(`    ${csharpFloatType} ${varName} = ${code};`);
+            }
+            // Emit any temps that depend on this var
+            const tempsForVar = tempAfterVar.get(varName) || [];
+            for (const { name: tempName, expr: tempExpr } of tempsForVar) {
+                const tempCode = codegen.generate(tempExpr);
+                if (format === 'typescript' || format === 'javascript') {
+                    lines.push(`  const ${tempName} = ${tempCode};`);
+                }
+                else if (format === 'python') {
+                    lines.push(`  ${tempName} = ${tempCode}`);
+                }
+                else if (format === 'csharp') {
+                    lines.push(`    ${csharpFloatType} ${tempName} = ${tempCode};`);
+                }
             }
         }
     }

package/dist/dsl/ExpressionUtils.d.ts

@@ -54,10 +54,16 @@ export declare function containsVariable(expr: Expression, varName: string): boo
  */
 export declare function expressionDepth(expr: Expression): number;
 /**
- * Serializes an expression to canonical string representation.
- * Used for expression comparison and hashing (CSE, CodeGen forward reuse).
+ * Serializes an expression to structural string representation.
+ * Used for exact expression comparison - operand order matters.
  *
  * This ensures consistent string representation of expressions across different
  * parts of the codebase.
  */
 export declare function serializeExpression(expr: Expression): string;
+/**
+ * Serializes an expression to canonical form for CSE matching.
+ * Commutative operations (+ and *) have operands sorted lexicographically,
+ * so a*b and b*a produce the same canonical string.
+ */
+export declare function serializeCanonical(expr: Expression): string;