Skip to content

Commit 3f7be45

Browse files
Copilotaarne
Copilot
and
aarne
authored
Support shadowed and memoized loop-scoped tool aliases in the AOT compiler (#109)
* Initial plan * fix: support shadowed loop-scoped tools in AOT compiler Co-authored-by: aarne <82001+aarne@users.noreply.github.com> * test: cover shadowed loop-scoped compiler support edge cases Co-authored-by: aarne <82001+aarne@users.noreply.github.com> * feat: compile memoized tool handles in AOT Co-authored-by: aarne <82001+aarne@users.noreply.github.com> * test: cover compiled memoization edge cases Co-authored-by: aarne <82001+aarne@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: aarne <82001+aarne@users.noreply.github.com>
1 parent 79c8ea7 commit 3f7be45

8 files changed

Lines changed: 267 additions & 83 deletions

File tree

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,13 +5,13 @@
55
"@stackables/bridge-parser": patch
66
---
77

8-
Add memoized tool handles with compiler fallback.
8+
Add memoized tool handles with compiler support.
99

1010
Bridge `with` declarations now support `memoize` for tool handles, including
1111
loop-scoped tool handles inside array mappings. Memoized handles reuse the same
1212
result for repeated calls with identical inputs, and each declared handle keeps
1313
its own cache.
1414

15-
The AOT compiler does not compile memoized tool handles yet. It now throws a
16-
dedicated incompatibility error for those bridges, and compiler `executeBridge`
17-
automatically falls back to the core ExecutionTree interpreter.
15+
The AOT compiler now compiles memoized tool handles too, including loop-scoped
16+
tool handles inside array mappings. Compiled execution preserves request-scoped
17+
caching semantics and reuses results for repeated calls with identical inputs.

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
---
2+
"@stackables/bridge-compiler": patch
3+
---
4+
5+
Compile shadowed loop-scoped tool handles in the AOT compiler.
6+
7+
Bridges can now redeclare the same tool alias in nested array scopes without
8+
triggering `BridgeCompilerIncompatibleError` or falling back to the interpreter.
9+
The compiler now assigns distinct tool instances to repeated handle bindings so
10+
each nested scope emits and reads from the correct tool call.

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

Lines changed: 2 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -10,35 +10,6 @@ export class BridgeCompilerIncompatibleError extends Error {
1010
}
1111
}
1212

13-
export function assertBridgeCompilerCompatible(bridge: Bridge): void {
14-
const operation = `${bridge.type}.${bridge.field}`;
15-
const memoizedHandles = bridge.handles
16-
.filter((handle) => handle.kind === "tool" && handle.memoize)
17-
.map((handle) => handle.handle);
18-
19-
if (memoizedHandles.length > 0) {
20-
throw new BridgeCompilerIncompatibleError(
21-
operation,
22-
`[bridge-compiler] ${operation}: memoized tool handles are not supported by AOT compilation yet (${memoizedHandles.join(", ")}).`,
23-
);
24-
}
25-
26-
const seenHandles = new Set<string>();
27-
const shadowedHandles = new Set<string>();
28-
29-
for (const handle of bridge.handles) {
30-
if (handle.kind !== "tool") continue;
31-
if (seenHandles.has(handle.handle)) {
32-
shadowedHandles.add(handle.handle);
33-
continue;
34-
}
35-
seenHandles.add(handle.handle);
36-
}
37-
38-
if (shadowedHandles.size > 0) {
39-
throw new BridgeCompilerIncompatibleError(
40-
operation,
41-
`[bridge-compiler] ${operation}: shadowed loop-scoped tool handles are not supported by AOT compilation yet (${[...shadowedHandles].join(", ")}).`,
42-
);
43-
}
13+
export function assertBridgeCompilerCompatible(_bridge: Bridge): void {
14+
// Intentionally empty: all currently supported bridge constructs compile.
4415
}

‎packages/bridge-compiler/src/codegen.ts‎

Lines changed: 120 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -297,6 +297,10 @@ class CodegenContext {
297297
private toolDepVars = new Map<string, string>();
298298
/** Sparse fieldset filter for output wire pruning. */
299299
private requestedFields: string[] | undefined;
300+
/** Per tool signature cursor used to assign distinct wire instances to repeated handle bindings. */
301+
private toolInstanceCursors = new Map<string, number>();
302+
/** Tool trunk keys declared with `memoize`. */
303+
private memoizedToolKeys = new Set<string>();
300304

301305
constructor(
302306
bridge: Bridge,
@@ -365,11 +369,14 @@ class CodegenContext {
365369
break;
366370
}
367371
}
368-
const instance = this.findInstance(module, refType, fieldName);
372+
const instance = this.findNextInstance(module, refType, fieldName);
369373
const tk = `${module}:${refType}:${fieldName}:${instance}`;
370374
const vn = `_t${++this.toolCounter}`;
371375
this.varMap.set(tk, vn);
372376
this.tools.set(tk, { trunkKey: tk, toolName: h.name, varName: vn });
377+
if (h.memoize) {
378+
this.memoizedToolKeys.add(tk);
379+
}
373380
break;
374381
}
375382
}
@@ -435,25 +442,36 @@ class CodegenContext {
435442
}
436443

437444
/** Find the instance number for a tool from the wires. */
438-
private findInstance(module: string, type: string, field: string): number {
445+
private findNextInstance(module: string, type: string, field: string): number {
446+
const sig = `${module}:${type}:${field}`;
447+
const instances: number[] = [];
439448
for (const w of this.bridge.wires) {
440449
if (
441450
w.to.module === module &&
442451
w.to.type === type &&
443452
w.to.field === field &&
444453
w.to.instance != null
445454
)
446-
return w.to.instance;
455+
instances.push(w.to.instance);
447456
if (
448457
"from" in w &&
449458
w.from.module === module &&
450459
w.from.type === type &&
451460
w.from.field === field &&
452461
w.from.instance != null
453462
)
454-
return w.from.instance;
455-
}
456-
return 1;
463+
instances.push(w.from.instance);
464+
}
465+
const uniqueInstances = [...new Set(instances)].sort((a, b) => a - b);
466+
const nextIndex = this.toolInstanceCursors.get(sig) ?? 0;
467+
this.toolInstanceCursors.set(sig, nextIndex + 1);
468+
if (uniqueInstances[nextIndex] != null) return uniqueInstances[nextIndex]!;
469+
const lastInstance = uniqueInstances.at(-1) ?? 0;
470+
// Some repeated handle bindings are never referenced in wires (for example,
471+
// an unused shadowed tool alias in a nested loop). In that case we still
472+
// need a distinct synthetic instance number so later bindings don't collide
473+
// with earlier tool registrations.
474+
return lastInstance + (nextIndex - uniqueInstances.length) + 1;
457475
}
458476

459477
// ── Main compilation entry point ──────────────────────────────────────────
@@ -729,6 +747,59 @@ class CodegenContext {
729747
lines.push(` throw err;`);
730748
lines.push(` }`);
731749
lines.push(` }`);
750+
if (this.memoizedToolKeys.size > 0) {
751+
lines.push(` const __toolMemoCache = new Map();`);
752+
lines.push(` function __stableMemoizeKey(value) {`);
753+
lines.push(` if (value === undefined) return "undefined";`);
754+
lines.push(' if (typeof value === "bigint") return `${value}n`;');
755+
lines.push(
756+
` if (value === null || typeof value !== "object") { const serialized = JSON.stringify(value); return serialized ?? String(value); }`,
757+
);
758+
lines.push(` if (Array.isArray(value)) {`);
759+
lines.push(
760+
' return `[${value.map((item) => __stableMemoizeKey(item)).join(",")}]`;',
761+
);
762+
lines.push(` }`);
763+
lines.push(
764+
` const entries = Object.entries(value).sort(([left], [right]) => (left < right ? -1 : left > right ? 1 : 0));`,
765+
);
766+
lines.push(
767+
' return `{${entries.map(([key, entryValue]) => `${JSON.stringify(key)}:${__stableMemoizeKey(entryValue)}`).join(",")}}`;',
768+
);
769+
lines.push(` }`);
770+
lines.push(
771+
` function __callMemoized(fn, input, toolName, memoizeKey) {`,
772+
);
773+
lines.push(` let toolCache = __toolMemoCache.get(memoizeKey);`);
774+
lines.push(` if (!toolCache) {`);
775+
lines.push(` toolCache = new Map();`);
776+
lines.push(` __toolMemoCache.set(memoizeKey, toolCache);`);
777+
lines.push(` }`);
778+
lines.push(` const cacheKey = __stableMemoizeKey(input);`);
779+
lines.push(` const cached = toolCache.get(cacheKey);`);
780+
lines.push(` if (cached !== undefined) return cached;`);
781+
lines.push(` try {`);
782+
lines.push(
783+
` const result = fn.bridge?.sync ? __callSync(fn, input, toolName) : __call(fn, input, toolName);`,
784+
);
785+
lines.push(` if (result && typeof result.then === "function") {`);
786+
lines.push(
787+
` const pending = Promise.resolve(result).catch((error) => {`,
788+
);
789+
lines.push(` toolCache.delete(cacheKey);`);
790+
lines.push(` throw error;`);
791+
lines.push(` });`);
792+
lines.push(` toolCache.set(cacheKey, pending);`);
793+
lines.push(` return pending;`);
794+
lines.push(` }`);
795+
lines.push(` toolCache.set(cacheKey, result);`);
796+
lines.push(` return result;`);
797+
lines.push(` } catch (error) {`);
798+
lines.push(` toolCache.delete(cacheKey);`);
799+
lines.push(` throw error;`);
800+
lines.push(` }`);
801+
lines.push(` }`);
802+
}
732803

733804
// ── Dead tool detection ────────────────────────────────────────────
734805
// Detect which tools are reachable from the (possibly filtered) output
@@ -938,19 +1009,33 @@ class CodegenContext {
9381009
* Generate a tool call expression that uses __callSync for sync tools at runtime,
9391010
* falling back to `await __call` for async tools. Used at individual call sites.
9401011
*/
941-
private syncAwareCall(fnName: string, inputObj: string): string {
1012+
private syncAwareCall(
1013+
fnName: string,
1014+
inputObj: string,
1015+
memoizeTrunkKey?: string,
1016+
): string {
9421017
const fn = `tools[${JSON.stringify(fnName)}]`;
9431018
const name = JSON.stringify(fnName);
1019+
if (memoizeTrunkKey && this.memoizedToolKeys.has(memoizeTrunkKey)) {
1020+
return `await __callMemoized(${fn}, ${inputObj}, ${name}, ${JSON.stringify(memoizeTrunkKey)})`;
1021+
}
9441022
return `(${fn}.bridge?.sync ? __callSync(${fn}, ${inputObj}, ${name}) : await __call(${fn}, ${inputObj}, ${name}))`;
9451023
}
9461024

9471025
/**
9481026
* Same as syncAwareCall but without await — for use inside Promise.all() and
9491027
* in sync array map bodies. Returns a value for sync tools, a Promise for async.
9501028
*/
951-
private syncAwareCallNoAwait(fnName: string, inputObj: string): string {
1029+
private syncAwareCallNoAwait(
1030+
fnName: string,
1031+
inputObj: string,
1032+
memoizeTrunkKey?: string,
1033+
): string {
9521034
const fn = `tools[${JSON.stringify(fnName)}]`;
9531035
const name = JSON.stringify(fnName);
1036+
if (memoizeTrunkKey && this.memoizedToolKeys.has(memoizeTrunkKey)) {
1037+
return `__callMemoized(${fn}, ${inputObj}, ${name}, ${JSON.stringify(memoizeTrunkKey)})`;
1038+
}
9541039
return `(${fn}.bridge?.sync ? __callSync(${fn}, ${inputObj}, ${name}) : __call(${fn}, ${inputObj}, ${name}))`;
9551040
}
9561041

@@ -986,18 +1071,18 @@ class CodegenContext {
9861071
);
9871072
if (mode === "fire-and-forget") {
9881073
lines.push(
989-
` try { ${this.syncAwareCall(tool.toolName, inputObj)}; } catch (_e) {}`,
1074+
` try { ${this.syncAwareCall(tool.toolName, inputObj, tool.trunkKey)}; } catch (_e) {}`,
9901075
);
9911076
lines.push(` const ${tool.varName} = undefined;`);
9921077
} else if (mode === "catch-guarded") {
9931078
// Catch-guarded: store result AND the actual error so unguarded wires can re-throw.
9941079
lines.push(` let ${tool.varName}, ${tool.varName}_err;`);
9951080
lines.push(
996-
` try { ${tool.varName} = ${this.syncAwareCall(tool.toolName, inputObj)}; } catch (_e) { if (_e?.name === "BridgePanicError" || _e?.name === "BridgeAbortError") throw _e; ${tool.varName}_err = _e; }`,
1081+
` 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; }`,
9971082
);
9981083
} else {
9991084
lines.push(
1000-
` const ${tool.varName} = ${this.syncAwareCall(tool.toolName, inputObj)};`,
1085+
` const ${tool.varName} = ${this.syncAwareCall(tool.toolName, inputObj, tool.trunkKey)};`,
10011086
);
10021087
}
10031088
return;
@@ -1070,7 +1155,7 @@ class CodegenContext {
10701155
lines.push(` let ${tool.varName};`);
10711156
lines.push(` try {`);
10721157
lines.push(
1073-
` ${tool.varName} = ${this.syncAwareCall(fnName, inputObj)};`,
1158+
` ${tool.varName} = ${this.syncAwareCall(fnName, inputObj, tool.trunkKey)};`,
10741159
);
10751160
lines.push(` } catch (_e) {`);
10761161
if ("value" in onErrorWire) {
@@ -1087,18 +1172,18 @@ class CodegenContext {
10871172
lines.push(` }`);
10881173
} else if (mode === "fire-and-forget") {
10891174
lines.push(
1090-
` try { ${this.syncAwareCall(fnName, inputObj)}; } catch (_e) {}`,
1175+
` try { ${this.syncAwareCall(fnName, inputObj, tool.trunkKey)}; } catch (_e) {}`,
10911176
);
10921177
lines.push(` const ${tool.varName} = undefined;`);
10931178
} else if (mode === "catch-guarded") {
10941179
// Catch-guarded: store result AND the actual error so unguarded wires can re-throw.
10951180
lines.push(` let ${tool.varName}, ${tool.varName}_err;`);
10961181
lines.push(
1097-
` try { ${tool.varName} = ${this.syncAwareCall(fnName, inputObj)}; } catch (_e) { if (_e?.name === "BridgePanicError" || _e?.name === "BridgeAbortError") throw _e; ${tool.varName}_err = _e; }`,
1182+
` try { ${tool.varName} = ${this.syncAwareCall(fnName, inputObj, tool.trunkKey)}; } catch (_e) { if (_e?.name === "BridgePanicError" || _e?.name === "BridgeAbortError") throw _e; ${tool.varName}_err = _e; }`,
10981183
);
10991184
} else {
11001185
lines.push(
1101-
` const ${tool.varName} = ${this.syncAwareCall(fnName, inputObj)};`,
1186+
` const ${tool.varName} = ${this.syncAwareCall(fnName, inputObj, tool.trunkKey)};`,
11021187
);
11031188
}
11041189
}
@@ -1190,7 +1275,7 @@ class CodegenContext {
11901275
4,
11911276
);
11921277
lines.push(
1193-
` const ${tool.varName} = ${this.syncAwareCall(tool.toolName, inputObj)};`,
1278+
` const ${tool.varName} = ${this.syncAwareCall(tool.toolName, inputObj, tool.trunkKey)};`,
11941279
);
11951280
return;
11961281
}
@@ -2385,7 +2470,9 @@ class CodegenContext {
23852470
// Non-internal tool in element scope — inline as an await __call
23862471
const inputObj = this.buildElementToolInput(toolWires, elVar);
23872472
const fnName = this.resolveToolDef(tool.toolName)?.fn ?? tool.toolName;
2388-
return `await __call(tools[${JSON.stringify(fnName)}], ${inputObj}, ${JSON.stringify(fnName)})`;
2473+
return this.memoizedToolKeys.has(trunkKey)
2474+
? `await __callMemoized(tools[${JSON.stringify(fnName)}], ${inputObj}, ${JSON.stringify(fnName)}, ${JSON.stringify(trunkKey)})`
2475+
: `await __call(tools[${JSON.stringify(fnName)}], ${inputObj}, ${JSON.stringify(fnName)})`;
23892476
}
23902477

23912478
/**
@@ -2571,11 +2658,19 @@ class CodegenContext {
25712658
const fnName = this.resolveToolDef(tool.toolName)?.fn ?? tool.toolName;
25722659
if (syncOnly) {
25732660
const fn = `tools[${JSON.stringify(fnName)}]`;
2661+
if (this.memoizedToolKeys.has(tk)) {
2662+
lines.push(
2663+
`const ${vn} = __callMemoized(${fn}, ${inputObj}, ${JSON.stringify(fnName)}, ${JSON.stringify(tk)});`,
2664+
);
2665+
} else {
2666+
lines.push(
2667+
`const ${vn} = __callSync(${fn}, ${inputObj}, ${JSON.stringify(fnName)});`,
2668+
);
2669+
}
2670+
} else {
25742671
lines.push(
2575-
`const ${vn} = __callSync(${fn}, ${inputObj}, ${JSON.stringify(fnName)});`,
2672+
`const ${vn} = ${this.syncAwareCall(fnName, inputObj, tk)};`,
25762673
);
2577-
} else {
2578-
lines.push(`const ${vn} = ${this.syncAwareCall(fnName, inputObj)};`);
25792674
}
25802675
}
25812676
}
@@ -2919,7 +3014,9 @@ class CodegenContext {
29193014
inputObj = this.buildObjectLiteral(toolWires, (w) => w.to.path, 4);
29203015
}
29213016

2922-
let expr = `(await __call(tools[${JSON.stringify(fnName)}], ${inputObj}, ${JSON.stringify(fnName)}))`;
3017+
let expr = this.memoizedToolKeys.has(key)
3018+
? `(await __callMemoized(tools[${JSON.stringify(fnName)}], ${inputObj}, ${JSON.stringify(fnName)}, ${JSON.stringify(key)}))`
3019+
: `(await __call(tools[${JSON.stringify(fnName)}], ${inputObj}, ${JSON.stringify(fnName)}))`;
29233020
if (ref.path.length > 0) {
29243021
expr =
29253022
expr + ref.path.map((p) => `?.[${JSON.stringify(p)}]`).join("");
@@ -3349,7 +3446,7 @@ class CodegenContext {
33493446
(w) => w.to.path,
33503447
4,
33513448
);
3352-
return this.syncAwareCallNoAwait(tool.toolName, inputObj);
3449+
return this.syncAwareCallNoAwait(tool.toolName, inputObj, tool.trunkKey);
33533450
}
33543451

33553452
const fnName = toolDef.fn ?? tool.toolName;
@@ -3384,7 +3481,7 @@ class CodegenContext {
33843481
const inputParts = [...inputEntries.values()];
33853482
const inputObj =
33863483
inputParts.length > 0 ? `{\n${inputParts.join(",\n")},\n }` : "{}";
3387-
return this.syncAwareCallNoAwait(fnName, inputObj);
3484+
return this.syncAwareCallNoAwait(fnName, inputObj, tool.trunkKey);
33883485
}
33893486

33903487
private topologicalLayers(toolWires: Map<string, Wire[]>): string[][] {

0 commit comments

Comments
 (0)