gradient-script 0.1.0

package/dist/dsl/AST.d.ts

@@ -0,0 +1,123 @@
+/**
+ * AST nodes for GradientScript DSL
+ * Supports function definitions with structured types
+ */
+import { Type } from './Types.js';
+/**
+ * Base AST node
+ */
+export interface ASTNode {
+    type?: Type;
+}
+/**
+ * Program (top-level)
+ */
+export interface Program extends ASTNode {
+    kind: 'program';
+    functions: FunctionDef[];
+}
+/**
+ * Function definition
+ */
+export interface FunctionDef extends ASTNode {
+    kind: 'function';
+    name: string;
+    parameters: Parameter[];
+    body: Statement[];
+    returnExpr: Expression;
+}
+/**
+ * Function parameter
+ */
+export interface Parameter {
+    name: string;
+    requiresGrad: boolean;
+    paramType?: StructTypeAnnotation;
+}
+/**
+ * Type annotation for structured types
+ */
+export interface StructTypeAnnotation {
+    components: string[];
+}
+/**
+ * Statement types
+ */
+export type Statement = Assignment;
+/**
+ * Assignment statement
+ */
+export interface Assignment extends ASTNode {
+    kind: 'assignment';
+    variable: string;
+    expression: Expression;
+}
+/**
+ * Expression types
+ */
+export type Expression = NumberLiteral | Variable | BinaryOp | UnaryOp | FunctionCall | ComponentAccess;
+/**
+ * Number literal
+ */
+export interface NumberLiteral extends ASTNode {
+    kind: 'number';
+    value: number;
+}
+/**
+ * Variable reference
+ */
+export interface Variable extends ASTNode {
+    kind: 'variable';
+    name: string;
+}
+/**
+ * Binary operation
+ */
+export interface BinaryOp extends ASTNode {
+    kind: 'binary';
+    operator: '+' | '-' | '*' | '/' | '^' | '**';
+    left: Expression;
+    right: Expression;
+}
+/**
+ * Unary operation
+ */
+export interface UnaryOp extends ASTNode {
+    kind: 'unary';
+    operator: '-' | '+';
+    operand: Expression;
+}
+/**
+ * Function call
+ */
+export interface FunctionCall extends ASTNode {
+    kind: 'call';
+    name: string;
+    args: Expression[];
+}
+/**
+ * Component access (e.g., u.x, v.y)
+ */
+export interface ComponentAccess extends ASTNode {
+    kind: 'component';
+    object: Expression;
+    component: string;
+}
+/**
+ * Visitor pattern for AST traversal
+ */
+export interface ASTVisitor<T> {
+    visitProgram(node: Program): T;
+    visitFunction(node: FunctionDef): T;
+    visitAssignment(node: Assignment): T;
+    visitNumber(node: NumberLiteral): T;
+    visitVariable(node: Variable): T;
+    visitBinary(node: BinaryOp): T;
+    visitUnary(node: UnaryOp): T;
+    visitCall(node: FunctionCall): T;
+    visitComponent(node: ComponentAccess): T;
+}
+/**
+ * Helper to visit any expression node
+ */
+export declare function visitExpression<T>(visitor: ASTVisitor<T>, expr: Expression): T;

package/dist/dsl/AST.js

@@ -0,0 +1,23 @@
+/**
+ * AST nodes for GradientScript DSL
+ * Supports function definitions with structured types
+ */
+/**
+ * Helper to visit any expression node
+ */
+export function visitExpression(visitor, expr) {
+    switch (expr.kind) {
+        case 'number':
+            return visitor.visitNumber(expr);
+        case 'variable':
+            return visitor.visitVariable(expr);
+        case 'binary':
+            return visitor.visitBinary(expr);
+        case 'unary':
+            return visitor.visitUnary(expr);
+        case 'call':
+            return visitor.visitCall(expr);
+        case 'component':
+            return visitor.visitComponent(expr);
+    }
+}

package/dist/dsl/BuiltIns.d.ts

@@ -0,0 +1,58 @@
+/**
+ * Built-in functions for GradientScript DSL
+ * Defines dot2d, cross2d, magnitude2d, etc.
+ */
+import { Type } from './Types.js';
+/**
+ * Information about function discontinuities
+ */
+export interface DiscontinuityInfo {
+    description: string;
+    condition: string;
+}
+/**
+ * Signature of a built-in function
+ */
+export interface BuiltInSignature {
+    name: string;
+    params: Type[];
+    returnType: Type;
+    implementation?: (args: any[]) => any;
+    discontinuities?: DiscontinuityInfo[];
+}
+/**
+ * Registry of all built-in functions
+ */
+export declare class BuiltInRegistry {
+    private functions;
+    constructor();
+    /**
+     * Register all built-in functions
+     */
+    private registerAll;
+    /**
+     * Register a built-in function
+     */
+    register(sig: BuiltInSignature): void;
+    /**
+     * Look up a built-in function by name and argument types
+     */
+    lookup(name: string, argTypes: Type[]): BuiltInSignature | undefined;
+    /**
+     * Check if argument types match parameter types
+     */
+    private matchesSignature;
+    /**
+     * Check if a function is built-in
+     */
+    isBuiltIn(name: string): boolean;
+    /**
+     * Get all overloads for a function name
+     */
+    getOverloads(name: string): BuiltInSignature[];
+    /**
+     * Get discontinuity information for a function
+     */
+    getDiscontinuities(name: string): DiscontinuityInfo[];
+}
+export declare const builtIns: BuiltInRegistry;

package/dist/dsl/BuiltIns.js

@@ -0,0 +1,181 @@
+/**
+ * Built-in functions for GradientScript DSL
+ * Defines dot2d, cross2d, magnitude2d, etc.
+ */
+import { Types } from './Types.js';
+/**
+ * Registry of all built-in functions
+ */
+export class BuiltInRegistry {
+    functions = new Map();
+    constructor() {
+        this.registerAll();
+    }
+    /**
+     * Register all built-in functions
+     */
+    registerAll() {
+        // 2D operations
+        this.register({
+            name: 'dot2d',
+            params: [Types.vec2(), Types.vec2()],
+            returnType: Types.scalar()
+        });
+        this.register({
+            name: 'cross2d',
+            params: [Types.vec2(), Types.vec2()],
+            returnType: Types.scalar()
+        });
+        this.register({
+            name: 'magnitude2d',
+            params: [Types.vec2()],
+            returnType: Types.scalar()
+        });
+        this.register({
+            name: 'normalize2d',
+            params: [Types.vec2()],
+            returnType: Types.vec2()
+        });
+        this.register({
+            name: 'distance2d',
+            params: [Types.vec2(), Types.vec2()],
+            returnType: Types.scalar()
+        });
+        // 3D operations
+        this.register({
+            name: 'dot3d',
+            params: [Types.vec3(), Types.vec3()],
+            returnType: Types.scalar()
+        });
+        this.register({
+            name: 'cross3d',
+            params: [Types.vec3(), Types.vec3()],
+            returnType: Types.vec3()
+        });
+        this.register({
+            name: 'magnitude3d',
+            params: [Types.vec3()],
+            returnType: Types.scalar()
+        });
+        this.register({
+            name: 'normalize3d',
+            params: [Types.vec3()],
+            returnType: Types.vec3()
+        });
+        // Math functions (scalar only)
+        const scalarMath = [
+            'sin', 'cos', 'tan',
+            'asin', 'acos', 'atan',
+            'exp', 'log', 'sqrt',
+            'abs'
+        ];
+        for (const name of scalarMath) {
+            this.register({
+                name,
+                params: [Types.scalar()],
+                returnType: Types.scalar()
+            });
+        }
+        // Binary math functions
+        this.register({
+            name: 'atan2',
+            params: [Types.scalar(), Types.scalar()],
+            returnType: Types.scalar(),
+            discontinuities: [{
+                    description: 'Branch cut discontinuity',
+                    condition: 'x < 0 and y ≈ 0 (near ±180°)'
+                }]
+        });
+        this.register({
+            name: 'pow',
+            params: [Types.scalar(), Types.scalar()],
+            returnType: Types.scalar()
+        });
+        this.register({
+            name: 'min',
+            params: [Types.scalar(), Types.scalar()],
+            returnType: Types.scalar()
+        });
+        this.register({
+            name: 'max',
+            params: [Types.scalar(), Types.scalar()],
+            returnType: Types.scalar(),
+            discontinuities: [{
+                    description: 'Non-smooth at equality',
+                    condition: 'a = b'
+                }]
+        });
+        this.register({
+            name: 'clamp',
+            params: [Types.scalar(), Types.scalar(), Types.scalar()],
+            returnType: Types.scalar(),
+            discontinuities: [{
+                    description: 'Non-smooth at boundaries',
+                    condition: 'x = min or x = max'
+                }]
+        });
+    }
+    /**
+     * Register a built-in function
+     */
+    register(sig) {
+        if (!this.functions.has(sig.name)) {
+            this.functions.set(sig.name, []);
+        }
+        this.functions.get(sig.name).push(sig);
+    }
+    /**
+     * Look up a built-in function by name and argument types
+     */
+    lookup(name, argTypes) {
+        const overloads = this.functions.get(name);
+        if (!overloads)
+            return undefined;
+        // Find matching overload
+        for (const sig of overloads) {
+            if (this.matchesSignature(argTypes, sig.params)) {
+                return sig;
+            }
+        }
+        return undefined;
+    }
+    /**
+     * Check if argument types match parameter types
+     */
+    matchesSignature(argTypes, paramTypes) {
+        if (argTypes.length !== paramTypes.length)
+            return false;
+        return argTypes.every((argType, i) => {
+            return Types.equals(argType, paramTypes[i]);
+        });
+    }
+    /**
+     * Check if a function is built-in
+     */
+    isBuiltIn(name) {
+        return this.functions.has(name);
+    }
+    /**
+     * Get all overloads for a function name
+     */
+    getOverloads(name) {
+        return this.functions.get(name) || [];
+    }
+    /**
+     * Get discontinuity information for a function
+     */
+    getDiscontinuities(name) {
+        const overloads = this.functions.get(name);
+        if (!overloads)
+            return [];
+        const discontinuities = [];
+        for (const sig of overloads) {
+            if (sig.discontinuities) {
+                discontinuities.push(...sig.discontinuities);
+            }
+        }
+        return discontinuities;
+    }
+}
+// Global built-in registry instance
+export const builtIns = new BuiltInRegistry();

package/dist/dsl/CSE.d.ts

@@ -0,0 +1,21 @@
+/**
+ * Common Subexpression Elimination (CSE)
+ * Identifies repeated expressions and factors them out
+ */
+import { Expression } from './AST.js';
+export interface CSEResult {
+    intermediates: Map<string, Expression>;
+    simplified: Expression;
+}
+export interface StructuredCSEResult {
+    intermediates: Map<string, Expression>;
+    components: Map<string, Expression>;
+}
+/**
+ * Perform CSE on an expression
+ */
+export declare function eliminateCommonSubexpressions(expr: Expression, minCount?: number): CSEResult;
+/**
+ * Perform CSE on structured gradients (for structured types like {x, y})
+ */
+export declare function eliminateCommonSubexpressionsStructured(components: Map<string, Expression>, minCount?: number): StructuredCSEResult;

package/dist/dsl/CSE.js

@@ -0,0 +1,194 @@
+/**
+ * Common Subexpression Elimination (CSE)
+ * Identifies repeated expressions and factors them out
+ */
+import { ExpressionTransformer } from './ExpressionTransformer.js';
+/**
+ * Serializes expressions to canonical string form for comparison
+ * This is a dedicated serializer that doesn't abuse the type system
+ */
+class ExpressionSerializer {
+    serialize(expr) {
+        switch (expr.kind) {
+            case 'number':
+                return `num(${expr.value})`;
+            case 'variable':
+                return `var(${expr.name})`;
+            case 'binary':
+                const left = this.serialize(expr.left);
+                const right = this.serialize(expr.right);
+                return `bin(${expr.operator},${left},${right})`;
+            case 'unary':
+                const operand = this.serialize(expr.operand);
+                return `un(${expr.operator},${operand})`;
+            case 'call':
+                const args = expr.args.map(arg => this.serialize(arg)).join(',');
+                return `call(${expr.name},${args})`;
+            case 'component':
+                const object = this.serialize(expr.object);
+                return `comp(${object},${expr.component})`;
+        }
+    }
+}
+/**
+ * Perform CSE on an expression
+ */
+export function eliminateCommonSubexpressions(expr, minCount = 2) {
+    const counter = new ExpressionCounter();
+    counter.count(expr);
+    const intermediates = new Map();
+    let varCounter = 0;
+    const subexprMap = new Map();
+    for (const [exprStr, count] of counter.counts.entries()) {
+        if (count >= minCount) {
+            const parsed = counter.expressions.get(exprStr);
+            if (parsed && shouldExtract(parsed)) {
+                const varName = `_tmp${varCounter++}`;
+                intermediates.set(varName, parsed);
+                subexprMap.set(exprStr, varName);
+            }
+        }
+    }
+    const simplified = substituteExpressions(expr, subexprMap, counter);
+    return { intermediates, simplified };
+}
+/**
+ * Perform CSE on structured gradients (for structured types like {x, y})
+ */
+export function eliminateCommonSubexpressionsStructured(components, minCount = 2) {
+    const counter = new ExpressionCounter();
+    for (const expr of components.values()) {
+        counter.count(expr);
+    }
+    const intermediates = new Map();
+    let varCounter = 0;
+    const subexprMap = new Map();
+    for (const [exprStr, count] of counter.counts.entries()) {
+        if (count >= minCount) {
+            const parsed = counter.expressions.get(exprStr);
+            if (parsed && shouldExtract(parsed)) {
+                const varName = `_tmp${varCounter++}`;
+                intermediates.set(varName, parsed);
+                subexprMap.set(exprStr, varName);
+            }
+        }
+    }
+    const simplifiedComponents = new Map();
+    for (const [comp, expr] of components.entries()) {
+        simplifiedComponents.set(comp, substituteExpressions(expr, subexprMap, counter));
+    }
+    return { intermediates, components: simplifiedComponents };
+}
+/**
+ * Check if an expression should be extracted
+ */
+function shouldExtract(expr) {
+    switch (expr.kind) {
+        case 'number':
+        case 'variable':
+            return false;
+        case 'component':
+            return expr.object.kind !== 'variable';
+        case 'unary':
+            return shouldExtract(expr.operand);
+        case 'binary':
+            return true;
+        case 'call':
+            return true;
+        default:
+            return false;
+    }
+}
+/**
+ * Counts occurrences of subexpressions during traversal
+ */
+class ExpressionCounter extends ExpressionTransformer {
+    counts = new Map();
+    expressions = new Map();
+    serializer = new ExpressionSerializer();
+    count(expr) {
+        this.transform(expr);
+    }
+    serialize(expr) {
+        return this.serializer.serialize(expr);
+    }
+    recordExpression(expr) {
+        const key = this.serialize(expr);
+        const currentCount = this.counts.get(key) || 0;
+        this.counts.set(key, currentCount + 1);
+        if (!this.expressions.has(key)) {
+            this.expressions.set(key, expr);
+        }
+    }
+    visitNumber(node) {
+        this.recordExpression(node);
+        return node;
+    }
+    visitVariable(node) {
+        this.recordExpression(node);
+        return node;
+    }
+    visitBinaryOp(node) {
+        this.recordExpression(node);
+        return super.visitBinaryOp(node);
+    }
+    visitUnaryOp(node) {
+        this.recordExpression(node);
+        return super.visitUnaryOp(node);
+    }
+    visitFunctionCall(node) {
+        this.recordExpression(node);
+        return super.visitFunctionCall(node);
+    }
+    visitComponentAccess(node) {
+        this.recordExpression(node);
+        return super.visitComponentAccess(node);
+    }
+}
+/**
+ * Substitute common subexpressions with variables
+ */
+function substituteExpressions(expr, subexprMap, counter) {
+    const key = counter.serialize(expr);
+    if (subexprMap.has(key)) {
+        return {
+            kind: 'variable',
+            name: subexprMap.get(key)
+        };
+    }
+    let result = expr;
+    for (const [exprStr, varName] of subexprMap.entries()) {
+        const exprToReplace = counter.expressions.get(exprStr);
+        if (exprToReplace && counter.serialize(result) !== exprStr) {
+            result = substituteInExpression(result, exprToReplace, { kind: 'variable', name: varName }, counter);
+        }
+    }
+    return result;
+}
+/**
+ * Transformer that substitutes a pattern with a replacement expression
+ */
+class SubstitutionTransformer extends ExpressionTransformer {
+    pattern;
+    replacement;
+    counter;
+    constructor(pattern, replacement, counter) {
+        super();
+        this.pattern = pattern;
+        this.replacement = replacement;
+        this.counter = counter;
+    }
+    transform(expr) {
+        if (this.counter.serialize(expr) === this.counter.serialize(this.pattern)) {
+            return this.replacement;
+        }
+        return super.transform(expr);
+    }
+}
+/**
+ * Helper to substitute an expression pattern with a replacement
+ */
+function substituteInExpression(expr, pattern, replacement, counter) {
+    const transformer = new SubstitutionTransformer(pattern, replacement, counter);
+    return transformer.transform(expr);
+}

package/dist/dsl/CodeGen.d.ts

@@ -0,0 +1,60 @@
+/**
+ * Code generation for GradientScript DSL
+ * Generates TypeScript/JavaScript code with gradient functions
+ */
+import { Expression, FunctionDef } from './AST.js';
+import { TypeEnv } from './Types.js';
+import { GradientResult } from './Differentiation.js';
+/**
+ * Code generation options
+ */
+export interface CodeGenOptions {
+    format?: 'typescript' | 'javascript' | 'python';
+    includeComments?: boolean;
+    simplify?: boolean;
+    cse?: boolean;
+    epsilon?: number;
+    emitGuards?: boolean;
+}
+/**
+ * Code generator for expressions
+ */
+export declare class ExpressionCodeGen {
+    private format;
+    constructor(format?: 'typescript' | 'javascript' | 'python');
+    /**
+     * Generate code for an expression
+     */
+    generate(expr: Expression): string;
+    private genNumber;
+    private genVariable;
+    private genBinary;
+    /**
+     * Generate expression with parentheses if needed based on precedence
+     */
+    private genWithPrecedence;
+    /**
+     * Determine if child expression needs parentheses
+     */
+    private needsParentheses;
+    /**
+     * Get operator precedence (higher number = higher precedence)
+     */
+    private getPrecedence;
+    private genUnary;
+    private genCall;
+    private genComponent;
+    private mapFunctionName;
+}
+/**
+ * Generate complete gradient function code
+ */
+export declare function generateGradientFunction(func: FunctionDef, gradients: GradientResult, env: TypeEnv, options?: CodeGenOptions): string;
+/**
+ * Generate the original forward function
+ */
+export declare function generateForwardFunction(func: FunctionDef, options?: CodeGenOptions): string;
+/**
+ * Generate complete output with both forward and gradient functions
+ */
+export declare function generateComplete(func: FunctionDef, gradients: GradientResult, env: TypeEnv, options?: CodeGenOptions): string;