gradient-script 0.1.0 → 0.3.0

package/dist/dsl/CodeGen.js

@@ -2,16 +2,33 @@
  * Code generation for GradientScript DSL
  * Generates TypeScript/JavaScript code with gradient functions
  */
-import { simplifyGradients } from './Simplify.js';
-import { eliminateCommonSubexpressionsStructured } from './CSE.js';
+import { simplifyGradients, simplifyPostCSE } from './Simplify.js';
+import { ExpressionTransformer } from './ExpressionTransformer.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);
+}
+function shouldTrackForForwardReuse(expr) {
+    switch (expr.kind) {
+        case 'number':
+        case 'variable':
+            return false;
+        default:
+            return true;
+    }
+}
 /**
  * Code generator for expressions
  */
 export class ExpressionCodeGen {
     format;
-    constructor(format = 'typescript') {
+    csharpFloatType;
+    constructor(format = 'typescript', csharpFloatType = 'float') {
         this.format = format;
+        this.csharpFloatType = csharpFloatType;
     }
     /**
      * Generate code for an expression
@@ -47,7 +64,7 @@ export class ExpressionCodeGen {
         if (this.format === 'python' && (op === '^' || op === '**')) {
             op = '**'; // Python uses **
         }
-        else if ((this.format === 'typescript' || this.format === 'javascript') && (op === '^' || op === '**')) {
+        else if ((this.format === 'typescript' || this.format === 'javascript' || this.format === 'csharp') && (op === '^' || op === '**')) {
             // Optimize: x^2 -> x*x, x^3 -> x*x*x (faster than Math.pow)
             // Only for simple expressions (variables, component access)
             const isSimple = expr.left.kind === 'variable' ||
@@ -63,15 +80,23 @@ 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})`;
                     }
                 }
             }
-            // Fall back to Math.pow for complex expressions or larger exponents
-            return `Math.pow(${left}, ${right})`;
+            // Fall back to Math.pow / MathF.Pow for complex expressions or larger exponents
+            if (this.format === 'csharp') {
+                const mathClass = this.csharpFloatType === 'float' ? 'MathF' : 'Math';
+                return `${mathClass}.Pow(${left}, ${right})`;
+            }
+            else {
+                return `Math.pow(${left}, ${right})`;
+            }
         }
         return `${left} ${op} ${right}`;
     }
@@ -160,6 +185,10 @@ export class ExpressionCodeGen {
             else if (this.format === 'python') {
                 return `max(${min}, min(${max}, ${x}))`;
             }
+            else if (this.format === 'csharp') {
+                const mathClass = this.csharpFloatType === 'float' ? 'MathF' : 'Math';
+                return `${mathClass}.Max(${min}, ${mathClass}.Min(${max}, ${x}))`;
+            }
         }
         // Map function names for different formats
         const funcName = this.mapFunctionName(expr.name);
@@ -167,63 +196,91 @@ export class ExpressionCodeGen {
     }
     genComponent(expr) {
         const obj = this.generate(expr.object);
+        if (this.format === 'csharp') {
+            // C# uses PascalCase for properties
+            return `${obj}.${capitalize(expr.component)}`;
+        }
         return `${obj}.${expr.component}`;
     }
+    // Math functions that should be mapped across all formats
+    static MATH_FUNCTIONS = [
+        'sin', 'cos', 'tan', 'asin', 'acos', 'atan', 'atan2',
+        'exp', 'log', 'sqrt', 'abs', 'pow', 'min', 'max'
+    ];
+    // Python built-in functions that don't need the math. prefix
+    static PYTHON_BUILTINS = ['abs', 'pow', 'min', 'max'];
     mapFunctionName(name) {
-        if (this.format === 'typescript' || this.format === 'javascript') {
-            const mathFuncs = {
-                'sin': 'Math.sin',
-                'cos': 'Math.cos',
-                'tan': 'Math.tan',
-                'asin': 'Math.asin',
-                'acos': 'Math.acos',
-                'atan': 'Math.atan',
-                'atan2': 'Math.atan2',
-                'exp': 'Math.exp',
-                'log': 'Math.log',
-                'sqrt': 'Math.sqrt',
-                'abs': 'Math.abs',
-                'pow': 'Math.pow',
-                'min': 'Math.min',
-                'max': 'Math.max'
-            };
-            return mathFuncs[name] || name;
-        }
-        else if (this.format === 'python') {
-            const mathFuncs = {
-                'atan2': 'math.atan2',
-                'sin': 'math.sin',
-                'cos': 'math.cos',
-                'tan': 'math.tan',
-                'asin': 'math.asin',
-                'acos': 'math.acos',
-                'atan': 'math.atan',
-                'exp': 'math.exp',
-                'log': 'math.log',
-                'sqrt': 'math.sqrt',
-                'abs': 'abs',
-                'pow': 'pow',
-                'min': 'min',
-                'max': 'max'
-            };
-            return mathFuncs[name] || name;
+        // Check if this is a known math function
+        if (!ExpressionCodeGen.MATH_FUNCTIONS.includes(name)) {
+            return name;
+        }
+        // Define format-specific mappers
+        const mappers = {
+            typescript: (fn) => `Math.${fn}`,
+            javascript: (fn) => `Math.${fn}`,
+            python: (fn) => ExpressionCodeGen.PYTHON_BUILTINS.includes(fn) ? fn : `math.${fn}`,
+            csharp: (fn) => {
+                const mathClass = this.csharpFloatType === 'float' ? 'MathF' : 'Math';
+                const capitalized = fn.charAt(0).toUpperCase() + fn.slice(1);
+                return `${mathClass}.${capitalized}`;
+            }
+        };
+        const mapper = mappers[this.format];
+        return mapper ? mapper(name) : name;
+    }
+}
+/**
+ * Generate C# struct for gradient return type
+ */
+function generateCSharpGradientStruct(func, gradients, floatType) {
+    const lines = [];
+    const structName = `${capitalize(func.name)}GradResult`;
+    lines.push(`public struct ${structName}`);
+    lines.push('{');
+    lines.push(`    public ${floatType} Value;`);
+    for (const [paramName, gradient] of gradients.gradients.entries()) {
+        const propName = capitalize(`d${paramName}`);
+        if (isStructuredGradient(gradient)) {
+            // Generate a nested struct type for structured gradients
+            const components = Array.from(gradient.components.keys());
+            lines.push(`    public ${capitalize(paramName)}Grad ${propName};`);
+        }
+        else {
+            lines.push(`    public ${floatType} ${propName};`);
+        }
+    }
+    lines.push('}');
+    // Generate nested struct types for structured gradients
+    for (const [paramName, gradient] of gradients.gradients.entries()) {
+        if (isStructuredGradient(gradient)) {
+            lines.push('');
+            lines.push(`public struct ${capitalize(paramName)}Grad`);
+            lines.push('{');
+            for (const comp of gradient.components.keys()) {
+                lines.push(`    public ${floatType} ${capitalize(comp)};`);
+            }
+            lines.push('}');
         }
-        return name;
     }
+    return lines;
 }
 /**
  * Generate complete gradient function code
  */
 export function generateGradientFunction(func, gradients, env, options = {}) {
     const format = options.format || 'typescript';
+    const csharpFloatType = options.csharpFloatType || 'float';
     const includeComments = options.includeComments !== false;
     const shouldSimplify = options.simplify !== false; // Default to true
-    // Simplify gradients if requested
-    const gradientsToUse = shouldSimplify
-        ? { gradients: simplifyGradients(gradients.gradients) }
-        : gradients;
-    const codegen = new ExpressionCodeGen(format);
+    // Note: We don't simplify here yet - we'll do it after forward expression substitution
+    const gradientsToUse = gradients;
+    const codegen = new ExpressionCodeGen(format, csharpFloatType);
     const lines = [];
+    // For C#, we need to generate a struct for the return type first
+    if (format === 'csharp') {
+        lines.push(...generateCSharpGradientStruct(func, gradientsToUse, csharpFloatType));
+        lines.push('');
+    }
     // Function signature
     const paramNames = func.parameters.map(p => p.name).join(', ');
     if (format === 'typescript' || format === 'javascript') {
@@ -232,44 +289,246 @@ export function generateGradientFunction(func, gradients, env, options = {}) {
     else if (format === 'python') {
         lines.push(`def ${func.name}_grad(${paramNames}):`);
     }
+    else if (format === 'csharp') {
+        const params = func.parameters.map(p => {
+            if (p.paramType && p.paramType.components) {
+                return `${capitalize(p.name)}Struct ${p.name}`;
+            }
+            return `${csharpFloatType} ${p.name}`;
+        }).join(', ');
+        const returnType = `${capitalize(func.name)}GradResult`;
+        lines.push(`public static ${returnType} ${capitalize(func.name)}_Grad(${params})`);
+        lines.push('{');
+    }
     // Forward pass - compute intermediate variables
     // Track which expressions are already computed for CSE reuse
-    const forwardPassVars = new Map();
+    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 expr = codegen.generate(stmt.expression);
-            // Track this for CSE reuse (store expression -> variable name mapping)
-            forwardPassVars.set(expr, varName);
+            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} = ${expr};`);
+                lines.push(`  const ${varName} = ${generatedExpr};`);
             }
-            else {
-                lines.push(`  ${varName} = ${expr}`);
+            else if (format === 'python') {
+                lines.push(`  ${varName} = ${generatedExpr}`);
+            }
+            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
-    let valueExpr = func.returnExpr;
+    const valueExpr = func.returnExpr;
+    const valueKey = serializeExpression(valueExpr);
+    const existingVar = forwardExpressionMap.get(valueKey);
     const valueCode = codegen.generate(valueExpr);
-    const existingVar = forwardPassVars.get(valueCode);
     if (existingVar) {
         // Reuse existing variable
         if (format === 'typescript' || format === 'javascript') {
             lines.push(`  const value = ${existingVar};`);
         }
-        else {
+        else if (format === 'python') {
             lines.push(`  value = ${existingVar}`);
         }
+        else if (format === 'csharp') {
+            lines.push(`    ${csharpFloatType} value = ${existingVar};`);
+        }
     }
     else {
         // Compute new value
         if (format === 'typescript' || format === 'javascript') {
             lines.push(`  const value = ${valueCode};`);
         }
-        else {
+        else if (format === 'python') {
             lines.push(`  value = ${valueCode}`);
         }
+        else if (format === 'csharp') {
+            lines.push(`    ${csharpFloatType} value = ${valueCode};`);
+        }
+        if (shouldTrackForForwardReuse(valueExpr) && !forwardExpressionMap.has(valueKey)) {
+            forwardExpressionMap.set(valueKey, 'value');
+        }
+    }
+    // Apply forward expression substitution multiple times until no more changes
+    // This handles nested expressions like sqrt(dx*dx + dy*dy) where dx = pix - pjx
+    reuseForwardExpressionsInGradients(gradientsToUse.gradients, forwardExpressionMap);
+    reuseForwardExpressionsInGradients(gradientsToUse.gradients, forwardExpressionMap);
+    reuseForwardExpressionsInGradients(gradientsToUse.gradients, forwardExpressionMap);
+    // Simplify gradients after forward expression substitution
+    if (shouldSimplify) {
+        const simplified = simplifyGradients(gradientsToUse.gradients);
+        gradientsToUse.gradients.clear();
+        for (const [key, value] of simplified.entries()) {
+            gradientsToUse.gradients.set(key, value);
+        }
     }
     lines.push('');
     // Generate gradients
@@ -277,74 +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 {
-                    lines.push(`  ${varName} = ${code}`);
-                }
+            else {
+                // Scalar gradient - wrap as a single 'value' component
+                allGradientComponents.set(paramName, new Map([['value', gradient]]));
+            }
+        }
+        // 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;
             }
-            // Emit epsilon guard if needed
-            if (shouldEmitGuards && denominatorVars.size > 0) {
-                lines.push('');
-                if (includeComments) {
-                    lines.push(`  ${comment} Guard against division by zero`);
+            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
@@ -366,7 +622,7 @@ export function generateGradientFunction(func, gradients, env, options = {}) {
                 }
                 lines.push(`  };`);
             }
-            else {
+            else if (format === 'python') {
                 lines.push(`  ${gradName} = {`);
                 for (const comp of components) {
                     const [key, value] = comp.split(': ');
@@ -374,6 +630,15 @@ export function generateGradientFunction(func, gradients, env, options = {}) {
                 }
                 lines.push(`  }`);
             }
+            else if (format === 'csharp') {
+                lines.push(`    var ${gradName} = new ${capitalize(paramName)}Grad`);
+                lines.push(`    {`);
+                for (const comp of components) {
+                    const [key, value] = comp.split(': ');
+                    lines.push(`        ${capitalize(key)} = ${value},`);
+                }
+                lines.push(`    };`);
+            }
         }
         else {
             // Scalar gradient
@@ -381,9 +646,12 @@ export function generateGradientFunction(func, gradients, env, options = {}) {
             if (format === 'typescript' || format === 'javascript') {
                 lines.push(`  const ${gradName} = ${code};`);
             }
-            else {
+            else if (format === 'python') {
                 lines.push(`  ${gradName} = ${code}`);
             }
+            else if (format === 'csharp') {
+                lines.push(`    ${csharpFloatType} ${gradName} = ${code};`);
+            }
         }
     }
     lines.push('');
@@ -399,7 +667,7 @@ export function generateGradientFunction(func, gradients, env, options = {}) {
         lines.push(`  };`);
         lines.push('}');
     }
-    else {
+    else if (format === 'python') {
         lines.push(`  return {`);
         lines.push(`    "value": value,`);
         for (const gradName of gradNames) {
@@ -407,33 +675,160 @@ export function generateGradientFunction(func, gradients, env, options = {}) {
         }
         lines.push(`  }`);
     }
+    else if (format === 'csharp') {
+        lines.push(`    return new ${capitalize(func.name)}GradResult`);
+        lines.push(`    {`);
+        lines.push(`        Value = value,`);
+        for (const gradName of gradNames) {
+            lines.push(`        ${capitalize(gradName)} = ${gradName},`);
+        }
+        lines.push(`    };`);
+        lines.push('}');
+    }
     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 codegen = new ExpressionCodeGen(format);
+    const csharpFloatType = options.csharpFloatType || 'float';
+    const shouldOptimize = options.cse !== false; // Optimize by default
+    const codegen = new ExpressionCodeGen(format, csharpFloatType);
     const lines = [];
     // Function signature
     const paramNames = func.parameters.map(p => p.name).join(', ');
     if (format === 'typescript' || format === 'javascript') {
         lines.push(`function ${func.name}(${paramNames}) {`);
     }
-    else {
+    else if (format === 'python') {
         lines.push(`def ${func.name}(${paramNames}):`);
     }
-    // Body
+    else if (format === 'csharp') {
+        const floatType = csharpFloatType;
+        const params = func.parameters.map(p => {
+            if (p.paramType && p.paramType.components) {
+                return `${capitalize(p.name)}Struct ${p.name}`;
+            }
+            return `${floatType} ${p.name}`;
+        }).join(', ');
+        lines.push(`public static ${floatType} ${capitalize(func.name)}(${params})`);
+        lines.push('{');
+    }
+    // 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 {
-                lines.push(`  ${varName} = ${expr}`);
+            else if (format === 'python') {
+                lines.push(`  ${varName} = ${code}`);
+            }
+            else if (format === 'csharp') {
+                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};`);
+                }
             }
         }
     }
@@ -443,9 +838,31 @@ export function generateForwardFunction(func, options = {}) {
         lines.push(`  return ${returnExpr};`);
         lines.push('}');
     }
-    else {
+    else if (format === 'python') {
         lines.push(`  return ${returnExpr}`);
     }
+    else if (format === 'csharp') {
+        lines.push(`    return ${returnExpr};`);
+        lines.push('}');
+    }
+    // For C#, prepend struct definitions
+    if (format === 'csharp') {
+        const structLines = [];
+        for (const param of func.parameters) {
+            if (param.paramType && param.paramType.components) {
+                structLines.push(`public struct ${capitalize(param.name)}Struct`);
+                structLines.push('{');
+                for (const comp of param.paramType.components) {
+                    structLines.push(`    public ${csharpFloatType} ${capitalize(comp)};`);
+                }
+                structLines.push('}');
+                structLines.push('');
+            }
+        }
+        if (structLines.length > 0) {
+            return structLines.join('\n') + lines.join('\n');
+        }
+    }
     return lines.join('\n');
 }
 /**
@@ -466,6 +883,40 @@ export function generateComplete(func, gradients, env, options = {}) {
     lines.push(generateGradientFunction(func, gradients, env, options));
     return lines.join('\n');
 }
+class ForwardExpressionSubstituter extends ExpressionTransformer {
+    forwardExpressions;
+    constructor(forwardExpressions) {
+        super();
+        this.forwardExpressions = forwardExpressions;
+    }
+    transform(expr) {
+        const key = serializeExpression(expr);
+        const varName = this.forwardExpressions.get(key);
+        if (varName) {
+            return {
+                kind: 'variable',
+                name: varName
+            };
+        }
+        return super.transform(expr);
+    }
+}
+function reuseForwardExpressionsInGradients(gradients, forwardExpressions) {
+    if (forwardExpressions.size === 0) {
+        return;
+    }
+    const substituter = new ForwardExpressionSubstituter(forwardExpressions);
+    for (const [paramName, gradient] of gradients.entries()) {
+        if (isStructuredGradient(gradient)) {
+            for (const [component, expr] of gradient.components.entries()) {
+                gradient.components.set(component, substituter.transform(expr));
+            }
+        }
+        else {
+            gradients.set(paramName, substituter.transform(gradient));
+        }
+    }
+}
 /**
  * Type guard for StructuredGradient
  */