refactor(compile, evaluate, types, lambda): introduce FnNode and enhance control flow handling

Author: codehz

src/compile.ts

@@ -1,7 +1,7 @@
 import type { ASTNode } from "./ast-types";
 import { generate, transformIdentifiers, transformPlaceholders } from "./generate";
 import { serializeArgumentToAST } from "./proxy-variable";
-import type { BranchNode, CompiledData, CompiledExpression, ExprValue, JumpNode, PhiNode } from "./types";
+import type { BranchNode, CompiledData, CompiledExpression, ExprValue, FnNode, JumpNode, PhiNode } from "./types";
 import { getVariableId } from "./variable";
 
 const ALLOWED_GLOBALS = new Set([
@@ -46,6 +46,20 @@ const ALLOWED_GLOBALS = new Set([
  */
 export interface CompileOptions {}
 
+/**
+ * 编译上下文（用于 lambda 参数分配）
+ */
+interface CompileCtx {
+  /** 下一个可用的参数名索引（全局递增） */
+  nextParamIndex: number;
+  /** 当前正在编译的表达式列表（栈顶为当前 lambda 或顶层） */
+  expressionStack: CompiledExpression[][];
+  /** 下一个全局表达式索引 */
+  nextIndex: number;
+  /** 变量数量（用于索引偏移） */
+  variableCount: number;
+}
+
 /**
  * 将 Proxy Expression 编译为可序列化的 JSON 结构
  *
@@ -106,6 +120,9 @@ export function compile<TResult>(
     // 已经转换为 $[N] 的跳过
     if (name.startsWith("$[") && /^\$\[\d+\]$/.test(name)) return name;
 
+    // lambda 参数名 _N 跳过
+    if (/^_\d+$/.test(name)) return name;
+
     const index = variableToIndex.get(name);
     if (index !== undefined) return `$[${index}]`;
 
@@ -119,70 +136,191 @@ export function compile<TResult>(
   }
 
   // 生成编译后的表达式（短路求值总是启用）
-  const expressions: CompiledExpression[] = [];
-  let nextIndex = variableOrder.length;
+  const topLevelExprs: CompiledExpression[] = [];
+  const ctx: CompileCtx = {
+    nextParamIndex: 0,
+    expressionStack: [topLevelExprs],
+    nextIndex: variableOrder.length,
+    variableCount: variableOrder.length,
+  };
 
-  function compileAst(node: ASTNode): number {
-    if (node.type === "BinaryExpr" && (node.operator === "||" || node.operator === "&&" || node.operator === "??")) {
-      return compileShortCircuit(node);
-    }
-    if (node.type === "ConditionalExpr") {
-      return compileConditional(node);
+  compileAst(transformed, ctx);
+
+  return [variableOrder, ...topLevelExprs];
+}
+
+function currentExprs(ctx: CompileCtx): CompiledExpression[] {
+  return ctx.expressionStack[ctx.expressionStack.length - 1]!;
+}
+
+/**
+ * 提取 AST 中所有 ArrowFunctionExpr 节点，编译为 FnNode，
+ * 并将原始位置替换为 $[N] 标识符引用。
+ */
+function extractAndCompileArrowFunctions(node: ASTNode, ctx: CompileCtx): ASTNode {
+  switch (node.type) {
+    case "ArrowFunctionExpr": {
+      // 编译这个箭头函数为 FnNode，返回 $[N] 引用
+      const idx = compileArrowFunction(node, ctx);
+      return { type: "Identifier", name: `$[${idx}]` };
     }
-    const exprStr = generate(node);
-    expressions.push(exprStr);
-    return nextIndex++;
-  }
 
-  function compileShortCircuit(node: ASTNode & { type: "BinaryExpr" }): number {
-    const leftIdx = compileAst(node.left);
+    case "BinaryExpr":
+      return {
+        ...node,
+        left: extractAndCompileArrowFunctions(node.left, ctx),
+        right: extractAndCompileArrowFunctions(node.right, ctx),
+      };
+
+    case "UnaryExpr":
+      return {
+        ...node,
+        argument: extractAndCompileArrowFunctions(node.argument, ctx),
+      };
+
+    case "ConditionalExpr":
+      return {
+        ...node,
+        test: extractAndCompileArrowFunctions(node.test, ctx),
+        consequent: extractAndCompileArrowFunctions(node.consequent, ctx),
+        alternate: extractAndCompileArrowFunctions(node.alternate, ctx),
+      };
 
-    const branchConditions: Record<string, string> = {
-      "||": `$[${leftIdx}]`,
-      "&&": `!$[${leftIdx}]`,
-      "??": `$[${leftIdx}]!=null`,
-    };
+    case "MemberExpr":
+      return {
+        ...node,
+        object: extractAndCompileArrowFunctions(node.object, ctx),
+        property: node.computed ? extractAndCompileArrowFunctions(node.property, ctx) : node.property,
+      };
 
-    const branchIdx = expressions.length;
-    expressions.push(["br", branchConditions[node.operator], 0] as BranchNode);
-    nextIndex++;
+    case "CallExpr":
+      return {
+        ...node,
+        callee: extractAndCompileArrowFunctions(node.callee, ctx),
+        arguments: node.arguments.map((arg) => extractAndCompileArrowFunctions(arg, ctx)),
+      };
 
-    compileAst(node.right);
-    const skipCount = expressions.length - branchIdx - 1;
-    (expressions[branchIdx] as BranchNode)[2] = skipCount;
+    case "ArrayExpr":
+      return {
+        ...node,
+        elements: node.elements.map((el) => extractAndCompileArrowFunctions(el, ctx)),
+      };
 
-    const phiIdx = nextIndex++;
-    expressions.push(["phi"] as PhiNode);
+    case "ObjectExpr":
+      return {
+        ...node,
+        properties: node.properties.map((prop) => ({
+          ...prop,
+          key: prop.computed ? extractAndCompileArrowFunctions(prop.key, ctx) : prop.key,
+          value: extractAndCompileArrowFunctions(prop.value, ctx),
+        })),
+      };
 
-    return phiIdx;
+    default:
+      return node;
   }
+}
 
-  function compileConditional(node: ASTNode & { type: "ConditionalExpr" }): number {
-    const testIdx = compileAst(node.test);
+function compileAst(node: ASTNode, ctx: CompileCtx): number {
+  if (node.type === "BinaryExpr" && (node.operator === "||" || node.operator === "&&" || node.operator === "??")) {
+    return compileShortCircuit(node, ctx);
+  }
+  if (node.type === "ConditionalExpr") {
+    return compileConditional(node, ctx);
+  }
+  if (node.type === "ArrowFunctionExpr") {
+    return compileArrowFunction(node, ctx);
+  }
 
-    const branchIdx = expressions.length;
-    expressions.push(["br", `$[${testIdx}]`, 0] as BranchNode);
-    nextIndex++;
+  // 提取并编译嵌套的箭头函数，替换为 $[N] 引用
+  const processed = extractAndCompileArrowFunctions(node, ctx);
 
-    compileAst(node.alternate);
+  const exprStr = generate(processed);
+  currentExprs(ctx).push(exprStr);
+  return ctx.nextIndex++;
+}
 
-    const jmpIdx = expressions.length;
-    expressions.push(["jmp", 0] as JumpNode);
-    nextIndex++;
+function compileShortCircuit(node: ASTNode & { type: "BinaryExpr" }, ctx: CompileCtx): number {
+  const exprs = currentExprs(ctx);
+  const leftIdx = compileAst(node.left, ctx);
 
-    compileAst(node.consequent);
-    const thenEndIdx = expressions.length;
+  const branchConditions: Record<string, string> = {
+    "||": `$[${leftIdx}]`,
+    "&&": `!$[${leftIdx}]`,
+    "??": `$[${leftIdx}]!=null`,
+  };
 
-    (expressions[branchIdx] as BranchNode)[2] = jmpIdx - branchIdx;
-    (expressions[jmpIdx] as JumpNode)[1] = thenEndIdx - jmpIdx - 1;
+  const branchIdx = exprs.length;
+  exprs.push(["br", branchConditions[node.operator], 0] as BranchNode);
+  ctx.nextIndex++;
 
-    const phiIdx = nextIndex++;
-    expressions.push(["phi"] as PhiNode);
+  compileAst(node.right, ctx);
+  const skipCount = exprs.length - branchIdx - 1;
+  (exprs[branchIdx] as BranchNode)[2] = skipCount;
 
-    return phiIdx;
+  const phiIdx = ctx.nextIndex++;
+  exprs.push(["phi"] as PhiNode);
+
+  return phiIdx;
+}
+
+function compileConditional(node: ASTNode & { type: "ConditionalExpr" }, ctx: CompileCtx): number {
+  const exprs = currentExprs(ctx);
+  const testIdx = compileAst(node.test, ctx);
+
+  const branchIdx = exprs.length;
+  exprs.push(["br", `$[${testIdx}]`, 0] as BranchNode);
+  ctx.nextIndex++;
+
+  compileAst(node.alternate, ctx);
+
+  const jmpIdx = exprs.length;
+  exprs.push(["jmp", 0] as JumpNode);
+  ctx.nextIndex++;
+
+  compileAst(node.consequent, ctx);
+  const thenEndIdx = exprs.length;
+
+  (exprs[branchIdx] as BranchNode)[2] = jmpIdx - branchIdx;
+  (exprs[jmpIdx] as JumpNode)[1] = thenEndIdx - jmpIdx - 1;
+
+  const phiIdx = ctx.nextIndex++;
+  exprs.push(["phi"] as PhiNode);
+
+  return phiIdx;
+}
+
+function compileArrowFunction(node: ASTNode & { type: "ArrowFunctionExpr" }, ctx: CompileCtx): number {
+  const paramCount = node.params.length;
+
+  // 0. Claim this FnNode's global index first
+  const fnIndex = ctx.nextIndex++;
+
+  // 1. 为参数分配 _N 名称
+  const paramMapping = new Map<symbol, string>();
+  for (const param of node.params) {
+    if (param.type === "Placeholder") {
+      const paramName = `_${ctx.nextParamIndex++}`;
+      paramMapping.set(param.id, paramName);
+    }
   }
 
-  compileAst(transformed);
+  // 2. 将函数体中的 lambda 参数 Placeholder 转换为 _N 标识符
+  const transformedBody = transformPlaceholders(node.body, (id) => {
+    return paramMapping.get(id) ?? null;
+  });
+
+  // 3. 编译函数体到新的表达式列表
+  const lambdaStmts: CompiledExpression[] = [];
+  ctx.expressionStack.push(lambdaStmts);
+
+  compileAst(transformedBody, ctx);
+
+  ctx.expressionStack.pop();
+
+  // 4. 构造 FnNode
+  const fnNode: FnNode = ["fn", paramCount, ...lambdaStmts];
+  currentExprs(ctx).push(fnNode);
 
-  return [variableOrder, ...expressions];
+  return fnIndex;
 }