gradient-script 0.2.0 → 0.3.1

package/dist/dsl/egraph/EGraph.js

@@ -0,0 +1,292 @@
+/**
+ * E-Graph: Equality Graph for expression optimization
+ *
+ * An e-graph efficiently represents equivalence classes of expressions.
+ * It supports:
+ * - Adding expressions (returns e-class ID)
+ * - Merging e-classes (union)
+ * - Finding canonical e-class (find)
+ * - Rebuilding after merges (maintains congruence)
+ */
+import { enodeKey, enodeChildren, enodeWithChildren } from './ENode.js';
+/**
+ * E-Graph: The main data structure
+ */
+export class EGraph {
+    nextId = 0;
+    classes = new Map();
+    parent = new Map(); // Union-find parent
+    rank = new Map(); // Union-find rank
+    hashcons = new Map(); // E-node key -> e-class
+    nodeStore = new Map(); // Key -> actual node
+    pending = []; // Classes needing rebuild
+    /**
+     * Find the canonical e-class ID (with path compression)
+     */
+    find(id) {
+        let root = id;
+        while (this.parent.get(root) !== root) {
+            root = this.parent.get(root);
+        }
+        // Path compression
+        let current = id;
+        while (this.parent.get(current) !== root) {
+            const next = this.parent.get(current);
+            this.parent.set(current, root);
+            current = next;
+        }
+        return root;
+    }
+    /**
+     * Add an e-node to the e-graph, returning its e-class ID
+     * If the node already exists, returns the existing class
+     */
+    add(node) {
+        // Canonicalize children first
+        const canonNode = this.canonicalize(node);
+        const key = enodeKey(canonNode);
+        // Check if already exists
+        const existing = this.hashcons.get(key);
+        if (existing !== undefined) {
+            return this.find(existing);
+        }
+        // Create new e-class
+        const id = this.nextId++;
+        this.parent.set(id, id);
+        this.rank.set(id, 0);
+        const eclass = {
+            id,
+            nodes: new Set([key]),
+            parents: new Set()
+        };
+        this.classes.set(id, eclass);
+        this.hashcons.set(key, id);
+        this.nodeStore.set(key, canonNode);
+        // Register as parent of children
+        for (const childId of enodeChildren(canonNode)) {
+            const childClass = this.classes.get(this.find(childId));
+            if (childClass) {
+                childClass.parents.add(key);
+            }
+        }
+        return id;
+    }
+    /**
+     * Merge two e-classes, returning the new canonical ID
+     */
+    merge(id1, id2) {
+        const root1 = this.find(id1);
+        const root2 = this.find(id2);
+        if (root1 === root2) {
+            return root1;
+        }
+        // Union by rank
+        const rank1 = this.rank.get(root1);
+        const rank2 = this.rank.get(root2);
+        let newRoot;
+        let oldRoot;
+        if (rank1 < rank2) {
+            newRoot = root2;
+            oldRoot = root1;
+        }
+        else if (rank1 > rank2) {
+            newRoot = root1;
+            oldRoot = root2;
+        }
+        else {
+            newRoot = root1;
+            oldRoot = root2;
+            this.rank.set(newRoot, rank1 + 1);
+        }
+        this.parent.set(oldRoot, newRoot);
+        // Merge e-class data
+        const newClass = this.classes.get(newRoot);
+        const oldClass = this.classes.get(oldRoot);
+        for (const nodeKey of oldClass.nodes) {
+            newClass.nodes.add(nodeKey);
+        }
+        for (const parentKey of oldClass.parents) {
+            newClass.parents.add(parentKey);
+        }
+        // Mark for rebuild
+        this.pending.push(newRoot);
+        return newRoot;
+    }
+    /**
+     * Rebuild the e-graph to restore congruence invariants
+     * Must be called after a batch of merges
+     */
+    rebuild() {
+        while (this.pending.length > 0) {
+            const todo = [...this.pending];
+            this.pending = [];
+            for (const classId of todo) {
+                this.repair(this.find(classId));
+            }
+        }
+    }
+    /**
+     * Repair an e-class after merges
+     */
+    repair(classId) {
+        const eclass = this.classes.get(classId);
+        if (!eclass)
+            return;
+        // Collect parent nodes that need re-canonicalization
+        const oldParents = new Set(eclass.parents);
+        eclass.parents.clear();
+        for (const parentKey of oldParents) {
+            const parentNode = this.nodeStore.get(parentKey);
+            if (!parentNode)
+                continue;
+            // Remove old hashcons entry
+            this.hashcons.delete(parentKey);
+            // Re-canonicalize and re-add
+            const canonNode = this.canonicalize(parentNode);
+            const newKey = enodeKey(canonNode);
+            const existingClass = this.hashcons.get(newKey);
+            if (existingClass !== undefined) {
+                // Node already exists in another class - merge
+                const parentClassId = this.findClassForNode(parentKey);
+                if (parentClassId !== undefined) {
+                    this.merge(parentClassId, existingClass);
+                }
+            }
+            else {
+                // Update hashcons with new key
+                const parentClassId = this.findClassForNode(parentKey);
+                if (parentClassId !== undefined) {
+                    this.hashcons.set(newKey, parentClassId);
+                    this.nodeStore.set(newKey, canonNode);
+                    const parentClass = this.classes.get(this.find(parentClassId));
+                    if (parentClass) {
+                        parentClass.nodes.delete(parentKey);
+                        parentClass.nodes.add(newKey);
+                    }
+                }
+            }
+        }
+        // Re-register parents for this class
+        for (const nodeKey of eclass.nodes) {
+            const node = this.nodeStore.get(nodeKey);
+            if (!node)
+                continue;
+            for (const childId of enodeChildren(node)) {
+                const childClass = this.classes.get(this.find(childId));
+                if (childClass) {
+                    childClass.parents.add(nodeKey);
+                }
+            }
+        }
+    }
+    /**
+     * Find which e-class contains a node (by key)
+     */
+    findClassForNode(nodeKey) {
+        for (const [id, eclass] of this.classes) {
+            if (eclass.nodes.has(nodeKey)) {
+                return this.find(id);
+            }
+        }
+        return undefined;
+    }
+    /**
+     * Canonicalize an e-node (update children to canonical IDs)
+     */
+    canonicalize(node) {
+        const children = enodeChildren(node);
+        if (children.length === 0) {
+            return node;
+        }
+        const canonChildren = children.map(id => this.find(id));
+        return enodeWithChildren(node, canonChildren);
+    }
+    /**
+     * Get all e-class IDs
+     */
+    getClassIds() {
+        const canonical = new Set();
+        for (const id of this.classes.keys()) {
+            canonical.add(this.find(id));
+        }
+        return [...canonical];
+    }
+    /**
+     * Get an e-class by ID
+     */
+    getClass(id) {
+        return this.classes.get(this.find(id));
+    }
+    /**
+     * Get all e-nodes in an e-class
+     */
+    getNodes(classId) {
+        const eclass = this.classes.get(this.find(classId));
+        if (!eclass)
+            return [];
+        const nodes = [];
+        for (const key of eclass.nodes) {
+            const node = this.nodeStore.get(key);
+            if (node) {
+                nodes.push(this.canonicalize(node));
+            }
+        }
+        return nodes;
+    }
+    /**
+     * Get the number of e-classes
+     */
+    get size() {
+        return this.getClassIds().length;
+    }
+    /**
+     * Get a node by its key
+     */
+    getNodeByKey(key) {
+        return this.nodeStore.get(key);
+    }
+    /**
+     * Lookup e-class by node (if it exists)
+     */
+    lookup(node) {
+        const canonNode = this.canonicalize(node);
+        const key = enodeKey(canonNode);
+        const id = this.hashcons.get(key);
+        return id !== undefined ? this.find(id) : undefined;
+    }
+    /**
+     * Debug: print e-graph state
+     */
+    dump() {
+        const lines = ['E-Graph:'];
+        for (const classId of this.getClassIds()) {
+            const eclass = this.classes.get(classId);
+            if (!eclass)
+                continue;
+            const nodeStrs = [...eclass.nodes].map(key => {
+                const node = this.nodeStore.get(key);
+                return node ? this.nodeToString(node) : key;
+            });
+            lines.push(`  [${classId}]: ${nodeStrs.join(' = ')}`);
+        }
+        return lines.join('\n');
+    }
+    /**
+     * Convert e-node to readable string
+     */
+    nodeToString(node) {
+        switch (node.tag) {
+            case 'num': return `${node.value}`;
+            case 'var': return node.name;
+            case 'add': return `(+ e${node.children[0]} e${node.children[1]})`;
+            case 'mul': return `(* e${node.children[0]} e${node.children[1]})`;
+            case 'sub': return `(- e${node.children[0]} e${node.children[1]})`;
+            case 'div': return `(/ e${node.children[0]} e${node.children[1]})`;
+            case 'pow': return `(^ e${node.children[0]} e${node.children[1]})`;
+            case 'neg': return `(neg e${node.child})`;
+            case 'inv': return `(inv e${node.child})`;
+            case 'call': return `(${node.name} ${node.children.map(c => `e${c}`).join(' ')})`;
+            case 'component': return `(. e${node.object} ${node.field})`;
+        }
+    }
+}

package/dist/dsl/egraph/ENode.d.ts

@@ -0,0 +1,63 @@
+/**
+ * E-Node: Expression nodes in an e-graph
+ *
+ * E-nodes are hash-consed (deduplicated) and reference e-classes by ID.
+ * This allows the e-graph to represent equivalence classes efficiently.
+ */
+export type EClassId = number;
+/**
+ * E-node variants representing different expression types
+ */
+export type ENode = {
+    tag: 'num';
+    value: number;
+} | {
+    tag: 'var';
+    name: string;
+} | {
+    tag: 'add';
+    children: [EClassId, EClassId];
+} | {
+    tag: 'mul';
+    children: [EClassId, EClassId];
+} | {
+    tag: 'sub';
+    children: [EClassId, EClassId];
+} | {
+    tag: 'div';
+    children: [EClassId, EClassId];
+} | {
+    tag: 'pow';
+    children: [EClassId, EClassId];
+} | {
+    tag: 'neg';
+    child: EClassId;
+} | {
+    tag: 'inv';
+    child: EClassId;
+} | {
+    tag: 'call';
+    name: string;
+    children: EClassId[];
+} | {
+    tag: 'component';
+    object: EClassId;
+    field: string;
+};
+/**
+ * Create a canonical string key for an e-node (for hash-consing)
+ * This key is used to detect structurally identical nodes.
+ */
+export declare function enodeKey(node: ENode): string;
+/**
+ * Get all e-class IDs that this node references (its children)
+ */
+export declare function enodeChildren(node: ENode): EClassId[];
+/**
+ * Create a new e-node with updated children (after canonicalization)
+ */
+export declare function enodeWithChildren(node: ENode, newChildren: EClassId[]): ENode;
+/**
+ * Check if two e-nodes are structurally equal
+ */
+export declare function enodesEqual(a: ENode, b: ENode): boolean;

package/dist/dsl/egraph/ENode.js

@@ -0,0 +1,94 @@
+/**
+ * E-Node: Expression nodes in an e-graph
+ *
+ * E-nodes are hash-consed (deduplicated) and reference e-classes by ID.
+ * This allows the e-graph to represent equivalence classes efficiently.
+ */
+/**
+ * Create a canonical string key for an e-node (for hash-consing)
+ * This key is used to detect structurally identical nodes.
+ */
+export function enodeKey(node) {
+    switch (node.tag) {
+        case 'num':
+            return `num:${node.value}`;
+        case 'var':
+            return `var:${node.name}`;
+        case 'add':
+            return `add:${node.children[0]},${node.children[1]}`;
+        case 'mul':
+            return `mul:${node.children[0]},${node.children[1]}`;
+        case 'sub':
+            return `sub:${node.children[0]},${node.children[1]}`;
+        case 'div':
+            return `div:${node.children[0]},${node.children[1]}`;
+        case 'pow':
+            return `pow:${node.children[0]},${node.children[1]}`;
+        case 'neg':
+            return `neg:${node.child}`;
+        case 'inv':
+            return `inv:${node.child}`;
+        case 'call':
+            return `call:${node.name}(${node.children.join(',')})`;
+        case 'component':
+            return `comp:${node.object}.${node.field}`;
+    }
+}
+/**
+ * Get all e-class IDs that this node references (its children)
+ */
+export function enodeChildren(node) {
+    switch (node.tag) {
+        case 'num':
+        case 'var':
+            return [];
+        case 'add':
+        case 'mul':
+        case 'sub':
+        case 'div':
+        case 'pow':
+            return [...node.children];
+        case 'neg':
+            return [node.child];
+        case 'inv':
+            return [node.child];
+        case 'call':
+            return [...node.children];
+        case 'component':
+            return [node.object];
+    }
+}
+/**
+ * Create a new e-node with updated children (after canonicalization)
+ */
+export function enodeWithChildren(node, newChildren) {
+    switch (node.tag) {
+        case 'num':
+        case 'var':
+            return node;
+        case 'add':
+            return { tag: 'add', children: [newChildren[0], newChildren[1]] };
+        case 'mul':
+            return { tag: 'mul', children: [newChildren[0], newChildren[1]] };
+        case 'sub':
+            return { tag: 'sub', children: [newChildren[0], newChildren[1]] };
+        case 'div':
+            return { tag: 'div', children: [newChildren[0], newChildren[1]] };
+        case 'pow':
+            return { tag: 'pow', children: [newChildren[0], newChildren[1]] };
+        case 'neg':
+            return { tag: 'neg', child: newChildren[0] };
+        case 'inv':
+            return { tag: 'inv', child: newChildren[0] };
+        case 'call':
+            return { tag: 'call', name: node.name, children: newChildren };
+        case 'component':
+            return { tag: 'component', object: newChildren[0], field: node.field };
+    }
+}
+/**
+ * Check if two e-nodes are structurally equal
+ */
+export function enodesEqual(a, b) {
+    return enodeKey(a) === enodeKey(b);
+}

package/dist/dsl/egraph/Extractor.d.ts

@@ -0,0 +1,49 @@
+/**
+ * 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 { EGraph } from './EGraph.js';
+import { EClassId } from './ENode.js';
+import { Expression } from '../AST.js';
+/**
+ * Cost of different operations
+ * Higher cost = less preferred
+ */
+export interface CostModel {
+    num: number;
+    var: number;
+    add: number;
+    sub: number;
+    mul: number;
+    div: number;
+    pow: number;
+    neg: number;
+    inv: number;
+    call: number;
+    component: number;
+}
+/**
+ * Default cost model - division is expensive
+ */
+export declare const defaultCostModel: CostModel;
+/**
+ * Result of extraction with CSE
+ */
+export interface ExtractionResult {
+    /** Temporary variable definitions (name -> expression) */
+    temps: Map<string, Expression>;
+    /** The extracted expressions (class ID -> expression using temps) */
+    expressions: Map<EClassId, Expression>;
+    /** Total cost */
+    totalCost: number;
+}
+/**
+ * Extract the best expression from an e-class
+ */
+export declare function extractBest(egraph: EGraph, rootId: EClassId, costModel?: CostModel): Expression;
+/**
+ * Extract multiple expressions with CSE (shared subexpressions become temps)
+ */
+export declare function extractWithCSE(egraph: EGraph, roots: EClassId[], costModel?: CostModel, minSharedCost?: number): ExtractionResult;