Skip to content

Instantly share code, notes, and snippets.

@Lucifier129

Lucifier129/modular-ast-optimizer.ts

Created December 28, 2025 10:19
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save Lucifier129/09668af1f113d744dfa561a4b571d3f8 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save Lucifier129/09668af1f113d744dfa561a4b571d3f8 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
modular-ast-optimizer.ts
type Result<T> = { ok: true; value: T } | { ok: false; error: string };
const Ok = <T>(value: T): Result<T> => ({ ok: true, value });
const Fail = (error: string): Result<any> => ({ ok: false, error });
class Accessor<S, T> {
constructor(
public readonly get: (s: S) => Result<T>,
public readonly set: (t: T, s: S) => Result<S>
) { }
static id<S>(): Accessor<S, S> { return new Accessor(Ok, (t, _) => Ok(t)); }
compose<Next>(next: Accessor<T, Next>): Accessor<S, Next> {
return new Accessor(
(s) => { const r = this.get(s); return r.ok ? next.get(r.value) : r; },
(n, s) => {
const rt = this.get(s);
if (!rt.ok) return rt;
const ut = next.set(n, rt.value);
return ut.ok ? this.set(ut.value, s) : ut;
}
);
}
prop<K extends keyof T>(key: K): Accessor<S, T[K]> {
return this.compose(new Accessor((t) => Ok(t[key]), (v, t) => Ok({ ...t, [key]: v })));
}
index(i: number): Accessor<S, T extends Array<infer I> ? I : never> {
return this.compose(new Accessor(
(t) => {
const arr = t as any as any[];
return i >= 0 && i < arr.length ? Ok(arr[i]) : Fail(`Index[${i}]`);
},
(v, t) => {
const arr = [...(t as any as any[])];
arr[i] = v;
return Ok(arr as any);
}
));
}
match<Sub extends T>(predicate: (t: T) => t is Sub): Accessor<S, Sub> {
return this.compose(new Accessor((t) => predicate(t) ? Ok(t) : Fail("Mismatch"), (v, _) => Ok(v as any)));
}
find<Item = T extends Array<infer I> ? I : never>(predicate: (item: Item) => boolean): Accessor<S, Item> {
return this.compose(new Accessor(
(t) => {
const arr = t as any as Item[];
const item = arr.find(predicate);
return item !== undefined ? Ok(item) : Fail("Not found");
},
(v, t) => {
const arr = [...(t as any as any[])];
const idx = arr.findIndex(predicate);
if (idx !== -1) arr[idx] = v;
return Ok(arr as any);
}
));
}
}
type Use<Out, Root> = <Local>(
processor: Processor<Local, Out, Root>,
accessor: Accessor<Root, Local>
) => Out;
type Processor<Local, Out, Root> = (
use: Use<Out, Root>,
accessor: Accessor<Root, Local>,
root: Root
) => Out;
function createRunner<Root, Local, Out>(
entryProcessor: Processor<Local, Out, Root>,
initialAccessor: Accessor<Root, Local>
): (root: Root) => Out {
return (root: Root) => {
const use: Use<Out, Root> = (proc, acc) => proc(use, acc, root);
return use(entryProcessor, initialAccessor)
};
}
abstract class Component<Local, Out, Root = any> {
constructor(
protected readonly useRunner: Use<Out, Root>,
protected readonly accessor: Accessor<Root, Local>
) { }
static asProcessor<L, O, R>(
this: typeof Component<L, O, R>
): Processor<L, O, R> {
const Ctor = this
return (useRunner, accessor, root) => {
const instance = new (Ctor as any)(useRunner, accessor);
const res = accessor.get(root);
if (!res.ok) return instance.catch(res.error);
return instance.impl(res.value);
};
}
protected use<L, SubC extends typeof Component<L, Out, Root>>(
Ctor: SubC,
accessor: Accessor<Root, L>
): Out {
return this.useRunner(Ctor.asProcessor(), accessor);
}
abstract impl(local: Local): Out;
abstract catch(error: string): Out;
protected prop<K extends keyof Local>(k: K) { return this.accessor.prop(k); }
protected index(i: number) { return this.accessor.index(i as any); }
protected match<S extends Local>(p: (l: Local) => l is S) { return this.accessor.match(p); }
protected find(p: (i: any) => boolean) { return this.accessor.find(p); }
}
// ---------------------------------------------------------
// 1. Domain: AST Definition
// ---------------------------------------------------------
type Expr =
| { type: 'Lit'; val: number }
| { type: 'Var'; name: string }
| { type: 'Add'; left: Expr; right: Expr }
| { type: 'Mul'; left: Expr; right: Expr };
// Type Guards for Lens matching
const isAdd = (e: Expr): e is Extract<Expr, { type: 'Add' }> => e.type === 'Add';
const isMul = (e: Expr): e is Extract<Expr, { type: 'Mul' }> => e.type === 'Mul';
const isLit = (e: Expr): e is Extract<Expr, { type: 'Lit' }> => e.type === 'Lit';
// ---------------------------------------------------------
// 2. Base Component: Rewriter
// ---------------------------------------------------------
// A Rewriter transforms a sub-tree of type Local into a generic Expr
abstract class Rewriter<Local extends Expr, Root> extends Component<Local, Expr, Root> {
// If lens fails, return a safe fallback (e.g. a 0 literal or error node)
// In a real system, you might propagate the error or return the original via a mechanism not shown here.
catch(msg: string): Expr {
console.warn(`[Rewriter Error]: ${msg}`);
return { type: 'Lit', val: 0 };
}
}
// ---------------------------------------------------------
// 3. Specific Rule Components
// ---------------------------------------------------------
class AddOptimizer<Root> extends Rewriter<Extract<Expr, { type: 'Add' }>, Root> {
impl(node: Extract<Expr, { type: 'Add' }>): Expr {
// 1. Recursively optimize children using Lenses
// We use 'Optimizer' (defined later) to handle polymorphism
const left = this.use(Optimizer, this.prop('left'));
const right = this.use(Optimizer, this.prop('right'));
// 2. Apply Rewrite Rules (Pattern Matching)
// Rule: Constant Folding (Lit + Lit)
if (left.type === 'Lit' && right.type === 'Lit') {
return { type: 'Lit', val: left.val + right.val };
}
// Rule: Identity (0 + x -> x)
if (left.type === 'Lit' && left.val === 0) return right;
// Rule: Identity (x + 0 -> x)
if (right.type === 'Lit' && right.val === 0) return left;
// No rule applied, reconstruct node with potentially optimized children
return { type: 'Add', left, right };
}
}
class MulOptimizer<Root> extends Rewriter<Extract<Expr, { type: 'Mul' }>, Root> {
impl(node: Extract<Expr, { type: 'Mul' }>): Expr {
// 1. Recurse
const left = this.use(Optimizer, this.prop('left'));
const right = this.use(Optimizer, this.prop('right'));
// 2. Apply Rules
// Rule: Zero Property (x * 0 -> 0, 0 * x -> 0)
if ((left.type === 'Lit' && left.val === 0) || (right.type === 'Lit' && right.val === 0)) {
return { type: 'Lit', val: 0 };
}
// Rule: Identity (x * 1 -> x)
if (right.type === 'Lit' && right.val === 1) return left;
if (left.type === 'Lit' && left.val === 1) return right;
// Rule: Constant Folding
if (left.type === 'Lit' && right.type === 'Lit') {
return { type: 'Lit', val: left.val * right.val };
}
return { type: 'Mul', left, right };
}
}
// ---------------------------------------------------------
// 4. Dispatcher Component
// ---------------------------------------------------------
class Optimizer<Root> extends Rewriter<Expr, Root> {
impl(node: Expr): Expr {
// Dispatch based on node type using the Accessor's match capability
// This allows us to "zoom in" to the specific subtype type-safely.
if (node.type === 'Add') {
// Switch context to AddOptimizer with a narrowed lens
return this.use(AddOptimizer, this.match(isAdd));
}
if (node.type === 'Mul') {
// Switch context to MulOptimizer with a narrowed lens
return this.use(MulOptimizer, this.match(isMul));
}
// Leaves (Lit, Var) are returned as-is (Base case)
return node;
}
}
// ---------------------------------------------------------
// 5. Execution & Demo
// ---------------------------------------------------------
// Helpers to build AST
const Lit = (val: number): Expr => ({ type: 'Lit', val });
const Var = (name: string): Expr => ({ type: 'Var', name });
const Add = (left: Expr, right: Expr): Expr => ({ type: 'Add', left, right });
const Mul = (left: Expr, right: Expr): Expr => ({ type: 'Mul', left, right });
// Helper to print AST
const prettyPrint = (e: Expr): string => {
switch (e.type) {
case 'Lit': return `${e.val}`;
case 'Var': return e.name;
case 'Add': return `(${prettyPrint(e.left)} + ${prettyPrint(e.right)})`;
case 'Mul': return `(${prettyPrint(e.left)} * ${prettyPrint(e.right)})`;
}
};
// Scenario: (x * 1) + (3 + 5) + (y * 0)
// Expectation: x + 8 + 0 -> (x + 8)
const complexAst: Expr = Add(
Add(
Mul(Var("x"), Lit(1)), // x * 1 -> x
Add(Lit(3), Lit(5)) // 3 + 5 -> 8
),
Mul(Var("y"), Lit(0)) // y * 0 -> 0
);
const optimize = createRunner(
Optimizer.asProcessor<Expr, Expr, Expr>(),
Accessor.id<Expr>()
);
console.log("--- BEFORE OPTIMIZATION ---");
console.log(prettyPrint(complexAst));
const optimizedAst = optimize(complexAst);
console.log("\n--- AFTER OPTIMIZATION ---");
console.log(prettyPrint(optimizedAst));
// Another run to show deep recursion handling
// (5 + 0) * (1 * 2) -> 5 * 2 -> 10
const nestedAst: Expr = Mul(
Add(Lit(5), Lit(0)),
Mul(Lit(1), Lit(2))
);
console.log("\n--- NESTED FOLDING ---");
console.log(`${prettyPrint(nestedAst)} => ${prettyPrint(optimize(nestedAst))}`);
@Lucifier129
Copy link
Author

Lucifier129 commented Dec 28, 2025 

type Result<T> = { ok: true; value: T } | { ok: false; error: string };
const Ok = <T>(value: T): Result<T> => ({ ok: true, value });
const Fail = (error: string): Result<any> => ({ ok: false, error });

class Accessor<S, T> {
  constructor(
    public readonly get: (s: S) => Result<T>,
    public readonly set: (t: T, s: S) => Result<S>
  ) {}

  static id<S>(): Accessor<S, S> {
    return new Accessor(Ok, (t, _) => Ok(t));
  }

  compose<Next>(next: Accessor<T, Next>): Accessor<S, Next> {
    return new Accessor(
      (s) => {
        const r = this.get(s);
        return r.ok ? next.get(r.value) : r;
      },
      (n, s) => {
        const rt = this.get(s);
        if (!rt.ok) return rt;
        const ut = next.set(n, rt.value);
        return ut.ok ? this.set(ut.value, s) : ut;
      }
    );
  }

  prop<K extends keyof T>(key: K): Accessor<S, T[K]> {
    return this.compose(
      new Accessor(
        (t) => Ok(t[key]),
        (v, t) => Ok({ ...t, [key]: v })
      )
    );
  }

  index(i: number): Accessor<S, T extends Array<infer I> ? I : never> {
    return this.compose(
      new Accessor(
        (t) => {
          const arr = t as any as any[];
          return i >= 0 && i < arr.length ? Ok(arr[i]) : Fail(`Index[${i}]`);
        },
        (v, t) => {
          const arr = [...(t as any as any[])];
          arr[i] = v;
          return Ok(arr as any);
        }
      )
    );
  }

  match<Sub extends T>(predicate: (t: T) => t is Sub): Accessor<S, Sub> {
    return this.compose(
      new Accessor(
        (t) => (predicate(t) ? Ok(t) : Fail("Mismatch")),
        (v, _) => Ok(v as any)
      )
    );
  }

  find<Item = T extends Array<infer I> ? I : never>(
    predicate: (item: Item) => boolean
  ): Accessor<S, Item> {
    return this.compose(
      new Accessor(
        (t) => {
          const arr = t as any as Item[];
          const item = arr.find(predicate);
          return item !== undefined ? Ok(item) : Fail("Not found");
        },
        (v, t) => {
          const arr = [...(t as any as any[])];
          const idx = arr.findIndex(predicate);
          if (idx !== -1) arr[idx] = v;
          return Ok(arr as any);
        }
      )
    );
  }
}

type Use<Out> = <Input>(processor: Processor<Input, Out>, input: Input) => Out;

type Processor<Input, Out> = (use: Use<Out>, input: Input) => Out;

function createRunner<Input, Out>(
  processor: Processor<Input, Out>
): (input: Input) => Out {
  return (input) => {
    const use: Use<Out> = (subProcessor, subInput) => {
      return subProcessor(use, subInput);
    };

    return use(processor, input);
  };
}

abstract class Component<Input, Out> {
  constructor(protected readonly useRunner: Use<Out>) {}

  static asProcessor<Input, Output>(
    this: typeof Component<Input, Output>
  ): Processor<Input, Output> {
    const Ctor = this;
    return (useRunner, input) => {
      const instance = new (Ctor as any)(useRunner);
      return instance.impl(input);
    };
  }

  protected use<NextInput>(
    Ctor: typeof Component<NextInput, Out>,
    input: NextInput
  ): Out {
    return this.useRunner(Ctor.asProcessor(), input);
  }

  abstract impl(input: Input): Out;
}

// ---------------------------------------------------------
// 1. Domain: AST Definition
// ---------------------------------------------------------

type ValExpr = { type: "Lit"; val: number };
type VarExpr = { type: "Var"; name: string };
type AddExpr = { type: "Add"; left: Expr; right: Expr };
type MulExpr = { type: "Mul"; left: Expr; right: Expr };

type Expr = ValExpr | VarExpr | AddExpr | MulExpr;

// Type Guards for Lens matching
const isAdd = (e: Expr): e is Extract<Expr, { type: "Add" }> =>
  e.type === "Add";
const isMul = (e: Expr): e is Extract<Expr, { type: "Mul" }> =>
  e.type === "Mul";
const isLit = (e: Expr): e is Extract<Expr, { type: "Lit" }> =>
  e.type === "Lit";

// ---------------------------------------------------------
// 2. Base Component: Rewriter
// ---------------------------------------------------------

// A Rewriter transforms a sub-tree of type Local into a generic Expr
abstract class Rewriter<Input extends Expr> extends Component<Input, Expr> {}

// ---------------------------------------------------------
// 3. Specific Rule Components
// ---------------------------------------------------------

class AddOptimizer extends Rewriter<AddExpr> {
  impl(node: AddExpr): Expr {
    // 1. Recursively optimize children using Lenses
    //    We use 'Optimizer' (defined later) to handle polymorphism
    const left = this.use(Optimizer, node.left);
    const right = this.use(Optimizer, node.right);

    // 2. Apply Rewrite Rules (Pattern Matching)

    // Rule: Constant Folding (Lit + Lit)
    if (left.type === "Lit" && right.type === "Lit") {
      return { type: "Lit", val: left.val + right.val };
    }

    // Rule: Identity (0 + x -> x)
    if (left.type === "Lit" && left.val === 0) return right;

    // Rule: Identity (x + 0 -> x)
    if (right.type === "Lit" && right.val === 0) return left;

    // No rule applied, reconstruct node with potentially optimized children
    return { type: "Add", left, right };
  }
}

class MulOptimizer extends Rewriter<MulExpr> {
  impl(node: MulExpr): Expr {
    // 1. Recurse
    const left = this.use(Optimizer, node.left);
    const right = this.use(Optimizer, node.right);

    // 2. Apply Rules

    // Rule: Zero Property (x * 0 -> 0, 0 * x -> 0)
    if (
      (left.type === "Lit" && left.val === 0) ||
      (right.type === "Lit" && right.val === 0)
    ) {
      return { type: "Lit", val: 0 };
    }

    // Rule: Identity (x * 1 -> x)
    if (right.type === "Lit" && right.val === 1) return left;
    if (left.type === "Lit" && left.val === 1) return right;

    // Rule: Constant Folding
    if (left.type === "Lit" && right.type === "Lit") {
      return { type: "Lit", val: left.val * right.val };
    }

    return { type: "Mul", left, right };
  }
}

// ---------------------------------------------------------
// 4. Dispatcher Component
// ---------------------------------------------------------

class Optimizer extends Rewriter<Expr> {
  impl(node: Expr): Expr {
    // Dispatch based on node type using the Accessor's match capability
    // This allows us to "zoom in" to the specific subtype type-safely.

    if (node.type === "Add") {
      // Switch context to AddOptimizer with a narrowed lens
      return this.use(AddOptimizer, node);
    }

    if (node.type === "Mul") {
      // Switch context to MulOptimizer with a narrowed lens
      return this.use(MulOptimizer, node);
    }

    // Leaves (Lit, Var) are returned as-is (Base case)
    return node;
  }
}

// ---------------------------------------------------------
// 5. Execution & Demo
// ---------------------------------------------------------

// Helpers to build AST
const Lit = (val: number): Expr => ({ type: "Lit", val });
const Var = (name: string): Expr => ({ type: "Var", name });
const Add = (left: Expr, right: Expr): Expr => ({ type: "Add", left, right });
const Mul = (left: Expr, right: Expr): Expr => ({ type: "Mul", left, right });

// Helper to print AST
const prettyPrint = (e: Expr): string => {
  switch (e.type) {
    case "Lit":
      return `${e.val}`;
    case "Var":
      return e.name;
    case "Add":
      return `(${prettyPrint(e.left)} + ${prettyPrint(e.right)})`;
    case "Mul":
      return `(${prettyPrint(e.left)} * ${prettyPrint(e.right)})`;
  }
};

// Scenario: (x * 1) + (3 + 5) + (y * 0)
// Expectation: x + 8 + 0 -> (x + 8)
const complexAst: Expr = Add(
  Add(
    Mul(Var("x"), Lit(1)), // x * 1 -> x
    Add(Lit(3), Lit(5)) // 3 + 5 -> 8
  ),
  Mul(Var("y"), Lit(0)) // y * 0 -> 0
);

const optimize = createRunner(Optimizer.asProcessor());

console.log("--- BEFORE OPTIMIZATION ---");
console.log(prettyPrint(complexAst));

const optimizedAst = optimize(complexAst);

console.log("\n--- AFTER OPTIMIZATION ---");
console.log(prettyPrint(optimizedAst));

// Another run to show deep recursion handling
// (5 + 0) * (1 * 2) -> 5 * 2 -> 10
const nestedAst: Expr = Mul(Add(Lit(5), Lit(0)), Mul(Lit(1), Lit(2)));

console.log("\n--- NESTED FOLDING ---");
console.log(
  `${prettyPrint(nestedAst)}  =>  ${prettyPrint(optimize(nestedAst))}`
);

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment