Support shadowed and memoized loop-scoped tool aliases in the AOT compiler

.changeset/compiler-fallback-loop-tools.md

@@ -5,13 +5,13 @@
 "@stackables/bridge-parser": patch
 ---
 
-Add memoized tool handles with compiler fallback.
+Add memoized tool handles with compiler support.
 
 Bridge `with` declarations now support `memoize` for tool handles, including
 loop-scoped tool handles inside array mappings. Memoized handles reuse the same
 result for repeated calls with identical inputs, and each declared handle keeps
 its own cache.
 
-The AOT compiler does not compile memoized tool handles yet. It now throws a
-dedicated incompatibility error for those bridges, and compiler `executeBridge`
-automatically falls back to the core ExecutionTree interpreter.
+The AOT compiler now compiles memoized tool handles too, including loop-scoped
+tool handles inside array mappings. Compiled execution preserves request-scoped
+caching semantics and reuses results for repeated calls with identical inputs.

.changeset/compiler-shadowed-loop-tools.md

@@ -0,0 +1,10 @@
+---
+"@stackables/bridge-compiler": patch
+---
+
+Compile shadowed loop-scoped tool handles in the AOT compiler.
+
+Bridges can now redeclare the same tool alias in nested array scopes without
+triggering `BridgeCompilerIncompatibleError` or falling back to the interpreter.
+The compiler now assigns distinct tool instances to repeated handle bindings so
+each nested scope emits and reads from the correct tool call.

packages/bridge-compiler/src/bridge-asserts.ts

@@ -10,35 +10,6 @@ export class BridgeCompilerIncompatibleError extends Error {
   }
 }
 
-export function assertBridgeCompilerCompatible(bridge: Bridge): void {
-  const operation = `${bridge.type}.${bridge.field}`;
-  const memoizedHandles = bridge.handles
-    .filter((handle) => handle.kind === "tool" && handle.memoize)
-    .map((handle) => handle.handle);
-
-  if (memoizedHandles.length > 0) {
-    throw new BridgeCompilerIncompatibleError(
-      operation,
-      `[bridge-compiler] ${operation}: memoized tool handles are not supported by AOT compilation yet (${memoizedHandles.join(", ")}).`,
-    );
-  }
-
-  const seenHandles = new Set<string>();
-  const shadowedHandles = new Set<string>();
-
-  for (const handle of bridge.handles) {
-    if (handle.kind !== "tool") continue;
-    if (seenHandles.has(handle.handle)) {
-      shadowedHandles.add(handle.handle);
-      continue;
-    }
-    seenHandles.add(handle.handle);
-  }
-
-  if (shadowedHandles.size > 0) {
-    throw new BridgeCompilerIncompatibleError(
-      operation,
-      `[bridge-compiler] ${operation}: shadowed loop-scoped tool handles are not supported by AOT compilation yet (${[...shadowedHandles].join(", ")}).`,
-    );
-  }
+export function assertBridgeCompilerCompatible(_bridge: Bridge): void {
+  // Intentionally empty: all currently supported bridge constructs compile.
 }

packages/bridge-compiler/src/codegen.ts

@@ -297,6 +297,10 @@ class CodegenContext {
   private toolDepVars = new Map<string, string>();
   /** Sparse fieldset filter for output wire pruning. */
   private requestedFields: string[] | undefined;
+  /** Per tool signature cursor used to assign distinct wire instances to repeated handle bindings. */
+  private toolInstanceCursors = new Map<string, number>();
+  /** Tool trunk keys declared with `memoize`. */
+  private memoizedToolKeys = new Set<string>();
 
   constructor(
     bridge: Bridge,
@@ -365,11 +369,14 @@ class CodegenContext {
               break;
             }
           }
-          const instance = this.findInstance(module, refType, fieldName);
+          const instance = this.findNextInstance(module, refType, fieldName);
           const tk = `${module}:${refType}:${fieldName}:${instance}`;
           const vn = `_t${++this.toolCounter}`;
           this.varMap.set(tk, vn);
           this.tools.set(tk, { trunkKey: tk, toolName: h.name, varName: vn });
+          if (h.memoize) {
+            this.memoizedToolKeys.add(tk);
+          }
           break;
         }
       }
@@ -435,25 +442,36 @@ class CodegenContext {
   }
 
   /** Find the instance number for a tool from the wires. */
-  private findInstance(module: string, type: string, field: string): number {
+  private findNextInstance(module: string, type: string, field: string): number {
+    const sig = `${module}:${type}:${field}`;
+    const instances: number[] = [];
     for (const w of this.bridge.wires) {
       if (
         w.to.module === module &&
         w.to.type === type &&
         w.to.field === field &&
         w.to.instance != null
       )
-        return w.to.instance;
+        instances.push(w.to.instance);
       if (
         "from" in w &&
         w.from.module === module &&
         w.from.type === type &&
         w.from.field === field &&
         w.from.instance != null
       )
-        return w.from.instance;
-    }
-    return 1;
+        instances.push(w.from.instance);
+    }
+    const uniqueInstances = [...new Set(instances)].sort((a, b) => a - b);
+    const nextIndex = this.toolInstanceCursors.get(sig) ?? 0;
+    this.toolInstanceCursors.set(sig, nextIndex + 1);
+    if (uniqueInstances[nextIndex] != null) return uniqueInstances[nextIndex]!;
+    const lastInstance = uniqueInstances.at(-1) ?? 0;
+    // Some repeated handle bindings are never referenced in wires (for example,
+    // an unused shadowed tool alias in a nested loop). In that case we still
+    // need a distinct synthetic instance number so later bindings don't collide
+    // with earlier tool registrations.
+    return lastInstance + (nextIndex - uniqueInstances.length) + 1;
   }
 
   // ── Main compilation entry point ──────────────────────────────────────────
@@ -729,6 +747,59 @@ class CodegenContext {
     lines.push(`      throw err;`);
     lines.push(`    }`);
     lines.push(`  }`);
+    if (this.memoizedToolKeys.size > 0) {
+      lines.push(`  const __toolMemoCache = new Map();`);
+      lines.push(`  function __stableMemoizeKey(value) {`);
+      lines.push(`    if (value === undefined) return "undefined";`);
+      lines.push('    if (typeof value === "bigint") return `${value}n`;');
+      lines.push(
+        `    if (value === null || typeof value !== "object") { const serialized = JSON.stringify(value); return serialized ?? String(value); }`,
+      );
+      lines.push(`    if (Array.isArray(value)) {`);
+      lines.push(
+        '      return `[${value.map((item) => __stableMemoizeKey(item)).join(",")}]`;',
+      );
+      lines.push(`    }`);
+      lines.push(
+        `    const entries = Object.entries(value).sort(([left], [right]) => (left < right ? -1 : left > right ? 1 : 0));`,
+      );
+      lines.push(
+        '    return `{${entries.map(([key, entryValue]) => `${JSON.stringify(key)}:${__stableMemoizeKey(entryValue)}`).join(",")}}`;',
+      );
+      lines.push(`  }`);
+      lines.push(
+        `  function __callMemoized(fn, input, toolName, memoizeKey) {`,
+      );
+      lines.push(`    let toolCache = __toolMemoCache.get(memoizeKey);`);
+      lines.push(`    if (!toolCache) {`);
+      lines.push(`      toolCache = new Map();`);
+      lines.push(`      __toolMemoCache.set(memoizeKey, toolCache);`);
+      lines.push(`    }`);
+      lines.push(`    const cacheKey = __stableMemoizeKey(input);`);
+      lines.push(`    const cached = toolCache.get(cacheKey);`);
+      lines.push(`    if (cached !== undefined) return cached;`);
+      lines.push(`    try {`);
+      lines.push(
+        `      const result = fn.bridge?.sync ? __callSync(fn, input, toolName) : __call(fn, input, toolName);`,
+      );
+      lines.push(`      if (result && typeof result.then === "function") {`);
+      lines.push(
+        `        const pending = Promise.resolve(result).catch((error) => {`,
+      );
+      lines.push(`          toolCache.delete(cacheKey);`);
+      lines.push(`          throw error;`);
+      lines.push(`        });`);
+      lines.push(`        toolCache.set(cacheKey, pending);`);
+      lines.push(`        return pending;`);
+      lines.push(`      }`);
+      lines.push(`      toolCache.set(cacheKey, result);`);
+      lines.push(`      return result;`);
+      lines.push(`    } catch (error) {`);
+      lines.push(`      toolCache.delete(cacheKey);`);
+      lines.push(`      throw error;`);
+      lines.push(`    }`);
+      lines.push(`  }`);
+    }
 
     // ── Dead tool detection ────────────────────────────────────────────
     // Detect which tools are reachable from the (possibly filtered) output
@@ -938,19 +1009,33 @@ class CodegenContext {
    * Generate a tool call expression that uses __callSync for sync tools at runtime,
    * falling back to `await __call` for async tools. Used at individual call sites.
    */
-  private syncAwareCall(fnName: string, inputObj: string): string {
+  private syncAwareCall(
+    fnName: string,
+    inputObj: string,
+    memoizeTrunkKey?: string,
+  ): string {
     const fn = `tools[${JSON.stringify(fnName)}]`;
     const name = JSON.stringify(fnName);
+    if (memoizeTrunkKey && this.memoizedToolKeys.has(memoizeTrunkKey)) {
+      return `await __callMemoized(${fn}, ${inputObj}, ${name}, ${JSON.stringify(memoizeTrunkKey)})`;
+    }
     return `(${fn}.bridge?.sync ? __callSync(${fn}, ${inputObj}, ${name}) : await __call(${fn}, ${inputObj}, ${name}))`;
   }
 
   /**
    * Same as syncAwareCall but without await — for use inside Promise.all() and
    * in sync array map bodies.  Returns a value for sync tools, a Promise for async.
    */
-  private syncAwareCallNoAwait(fnName: string, inputObj: string): string {
+  private syncAwareCallNoAwait(
+    fnName: string,
+    inputObj: string,
+    memoizeTrunkKey?: string,
+  ): string {
     const fn = `tools[${JSON.stringify(fnName)}]`;
     const name = JSON.stringify(fnName);
+    if (memoizeTrunkKey && this.memoizedToolKeys.has(memoizeTrunkKey)) {
+      return `__callMemoized(${fn}, ${inputObj}, ${name}, ${JSON.stringify(memoizeTrunkKey)})`;
+    }
     return `(${fn}.bridge?.sync ? __callSync(${fn}, ${inputObj}, ${name}) : __call(${fn}, ${inputObj}, ${name}))`;
   }
 
@@ -986,18 +1071,18 @@ class CodegenContext {
       );
       if (mode === "fire-and-forget") {
         lines.push(
-          `  try { ${this.syncAwareCall(tool.toolName, inputObj)}; } catch (_e) {}`,
+          `  try { ${this.syncAwareCall(tool.toolName, inputObj, tool.trunkKey)}; } catch (_e) {}`,
         );
         lines.push(`  const ${tool.varName} = undefined;`);
       } else if (mode === "catch-guarded") {
         // Catch-guarded: store result AND the actual error so unguarded wires can re-throw.
         lines.push(`  let ${tool.varName}, ${tool.varName}_err;`);
         lines.push(
-          `  try { ${tool.varName} = ${this.syncAwareCall(tool.toolName, inputObj)}; } catch (_e) { if (_e?.name === "BridgePanicError" || _e?.name === "BridgeAbortError") throw _e; ${tool.varName}_err = _e; }`,
+          `  try { ${tool.varName} = ${this.syncAwareCall(tool.toolName, inputObj, tool.trunkKey)}; } catch (_e) { if (_e?.name === "BridgePanicError" || _e?.name === "BridgeAbortError") throw _e; ${tool.varName}_err = _e; }`,
         );
       } else {
         lines.push(
-          `  const ${tool.varName} = ${this.syncAwareCall(tool.toolName, inputObj)};`,
+          `  const ${tool.varName} = ${this.syncAwareCall(tool.toolName, inputObj, tool.trunkKey)};`,
         );
       }
       return;
@@ -1070,7 +1155,7 @@ class CodegenContext {
       lines.push(`  let ${tool.varName};`);
       lines.push(`  try {`);
       lines.push(
-        `    ${tool.varName} = ${this.syncAwareCall(fnName, inputObj)};`,
+        `    ${tool.varName} = ${this.syncAwareCall(fnName, inputObj, tool.trunkKey)};`,
       );
       lines.push(`  } catch (_e) {`);
       if ("value" in onErrorWire) {
@@ -1087,18 +1172,18 @@ class CodegenContext {
       lines.push(`  }`);
     } else if (mode === "fire-and-forget") {
       lines.push(
-        `  try { ${this.syncAwareCall(fnName, inputObj)}; } catch (_e) {}`,
+        `  try { ${this.syncAwareCall(fnName, inputObj, tool.trunkKey)}; } catch (_e) {}`,
       );
       lines.push(`  const ${tool.varName} = undefined;`);
     } else if (mode === "catch-guarded") {
       // Catch-guarded: store result AND the actual error so unguarded wires can re-throw.
       lines.push(`  let ${tool.varName}, ${tool.varName}_err;`);
       lines.push(
-        `  try { ${tool.varName} = ${this.syncAwareCall(fnName, inputObj)}; } catch (_e) { if (_e?.name === "BridgePanicError" || _e?.name === "BridgeAbortError") throw _e; ${tool.varName}_err = _e; }`,
+        `  try { ${tool.varName} = ${this.syncAwareCall(fnName, inputObj, tool.trunkKey)}; } catch (_e) { if (_e?.name === "BridgePanicError" || _e?.name === "BridgeAbortError") throw _e; ${tool.varName}_err = _e; }`,
       );
     } else {
       lines.push(
-        `  const ${tool.varName} = ${this.syncAwareCall(fnName, inputObj)};`,
+        `  const ${tool.varName} = ${this.syncAwareCall(fnName, inputObj, tool.trunkKey)};`,
       );
     }
   }
@@ -1190,7 +1275,7 @@ class CodegenContext {
           4,
         );
         lines.push(
-          `  const ${tool.varName} = ${this.syncAwareCall(tool.toolName, inputObj)};`,
+          `  const ${tool.varName} = ${this.syncAwareCall(tool.toolName, inputObj, tool.trunkKey)};`,
         );
         return;
       }
@@ -2385,7 +2470,9 @@ class CodegenContext {
     // Non-internal tool in element scope — inline as an await __call
     const inputObj = this.buildElementToolInput(toolWires, elVar);
     const fnName = this.resolveToolDef(tool.toolName)?.fn ?? tool.toolName;
-    return `await __call(tools[${JSON.stringify(fnName)}], ${inputObj}, ${JSON.stringify(fnName)})`;
+    return this.memoizedToolKeys.has(trunkKey)
+      ? `await __callMemoized(tools[${JSON.stringify(fnName)}], ${inputObj}, ${JSON.stringify(fnName)}, ${JSON.stringify(trunkKey)})`
+      : `await __call(tools[${JSON.stringify(fnName)}], ${inputObj}, ${JSON.stringify(fnName)})`;
   }
 
   /**
@@ -2571,11 +2658,19 @@ class CodegenContext {
         const fnName = this.resolveToolDef(tool.toolName)?.fn ?? tool.toolName;
         if (syncOnly) {
           const fn = `tools[${JSON.stringify(fnName)}]`;
+          if (this.memoizedToolKeys.has(tk)) {
+            lines.push(
+              `const ${vn} = __callMemoized(${fn}, ${inputObj}, ${JSON.stringify(fnName)}, ${JSON.stringify(tk)});`,
+            );
+          } else {
+            lines.push(
+              `const ${vn} = __callSync(${fn}, ${inputObj}, ${JSON.stringify(fnName)});`,
+            );
+          }
+        } else {
           lines.push(
-            `const ${vn} = __callSync(${fn}, ${inputObj}, ${JSON.stringify(fnName)});`,
+            `const ${vn} = ${this.syncAwareCall(fnName, inputObj, tk)};`,
           );
-        } else {
-          lines.push(`const ${vn} = ${this.syncAwareCall(fnName, inputObj)};`);
         }
       }
     }
@@ -2919,7 +3014,9 @@ class CodegenContext {
           inputObj = this.buildObjectLiteral(toolWires, (w) => w.to.path, 4);
         }
 
-        let expr = `(await __call(tools[${JSON.stringify(fnName)}], ${inputObj}, ${JSON.stringify(fnName)}))`;
+        let expr = this.memoizedToolKeys.has(key)
+          ? `(await __callMemoized(tools[${JSON.stringify(fnName)}], ${inputObj}, ${JSON.stringify(fnName)}, ${JSON.stringify(key)}))`
+          : `(await __call(tools[${JSON.stringify(fnName)}], ${inputObj}, ${JSON.stringify(fnName)}))`;
         if (ref.path.length > 0) {
           expr =
             expr + ref.path.map((p) => `?.[${JSON.stringify(p)}]`).join("");
@@ -3349,7 +3446,7 @@ class CodegenContext {
         (w) => w.to.path,
         4,
       );
-      return this.syncAwareCallNoAwait(tool.toolName, inputObj);
+      return this.syncAwareCallNoAwait(tool.toolName, inputObj, tool.trunkKey);
     }
 
     const fnName = toolDef.fn ?? tool.toolName;
@@ -3384,7 +3481,7 @@ class CodegenContext {
     const inputParts = [...inputEntries.values()];
     const inputObj =
       inputParts.length > 0 ? `{\n${inputParts.join(",\n")},\n  }` : "{}";
-    return this.syncAwareCallNoAwait(fnName, inputObj);
+    return this.syncAwareCallNoAwait(fnName, inputObj, tool.trunkKey);
   }
 
   private topologicalLayers(toolWires: Map<string, Wire[]>): string[][] {