@mondaydotcomorg/atp-compiler 0.19.22 → 0.19.26

package/src/transformer/array-transformer-batch-reconstruct.ts

@@ -0,0 +1,179 @@
+import * as t from '@babel/types';
+import _traverse from '@babel/traverse';
+const traverse = typeof (_traverse as any).default === 'function' ? (_traverse as any).default : _traverse;
+import { generateUniqueId } from '../runtime/context.js';
+import { BatchParallelDetector } from './batch-detector.js';
+import { findAllLLMCallExpressions } from './array-transformer-utils.js';
+
+/**
+ * Transform array method to batch LLM calls while preserving callback logic.
+ * Supports multiple LLM calls per callback.
+ *
+ * This creates a multi-step transformation:
+ * 1. Batch all LLM calls in parallel (one batch per unique type:operation)
+ * 2. Reconstruct objects using the batched results
+ */
+export function transformToBatchWithReconstruction(
+	path: any,
+	node: t.CallExpression,
+	methodName: string,
+	callback: t.Function,
+	batchDetector: BatchParallelDetector,
+	onTransform: () => void
+): boolean {
+	if (methodName !== 'map') {
+		return false;
+	}
+
+	const paramName = callback.params[0];
+	if (!t.isIdentifier(paramName)) {
+		return false;
+	}
+	const param = paramName.name;
+	const array = (node.callee as t.MemberExpression).object;
+
+	// Find ALL LLM calls
+	const llmCalls = findAllLLMCallExpressions(callback.body, batchDetector);
+	if (llmCalls.length === 0) {
+		return false;
+	}
+
+	const methodId = generateUniqueId(`${methodName}_batch_reconstruct`);
+
+	const originalIndexParam = callback.params[1];
+	const indexVar =
+		originalIndexParam && t.isIdentifier(originalIndexParam) ? originalIndexParam.name : '__idx';
+
+	// Create batch declarations - one per LLM call (in order of appearance)
+	// This ensures each call gets its own result array
+	const batchDeclarations: t.Statement[] = [];
+	const resultVarByCallIndex = new Map<number, string>();
+
+	for (let i = 0; i < llmCalls.length; i++) {
+		const call = llmCalls[i]!;
+		const resultsVar = `__batch_results_${i}_${methodId.replace(/[^a-zA-Z0-9]/g, '_')}`;
+		resultVarByCallIndex.set(i, resultsVar);
+
+		const payloadMapper = t.arrowFunctionExpression(
+			[t.identifier(param)],
+			t.objectExpression([
+				t.objectProperty(t.identifier('type'), t.stringLiteral(call.callInfo.type)),
+				t.objectProperty(
+					t.identifier('operation'),
+					t.stringLiteral(call.callInfo.operation)
+				),
+				t.objectProperty(t.identifier('payload'), t.cloneNode(call.payloadNode, true)),
+			])
+		);
+
+		const batchCall = t.awaitExpression(
+			t.callExpression(
+				t.memberExpression(t.identifier('__runtime'), t.identifier('batchParallel')),
+				[
+					t.callExpression(
+						t.memberExpression(t.cloneNode(array, true), t.identifier('map')),
+						[payloadMapper]
+					),
+					t.stringLiteral(`${methodId}_${i}`),
+				]
+			)
+		);
+
+		batchDeclarations.push(
+			t.variableDeclaration('const', [
+				t.variableDeclarator(t.identifier(resultsVar), batchCall),
+			])
+		);
+	}
+
+	// Clone the callback body for reconstruction
+	const clonedBody = t.cloneNode(callback.body, true);
+
+	let traversableNode: t.Statement;
+	if (t.isBlockStatement(clonedBody)) {
+		traversableNode = t.functionDeclaration(t.identifier('__temp'), [], clonedBody);
+	} else {
+		traversableNode = t.expressionStatement(clonedBody as t.Expression);
+	}
+
+	// Replace each await expression with the corresponding result access
+	// We match calls by comparing their structure to the original calls
+	let replacementCount = 0;
+
+	traverse(t.file(t.program([traversableNode])), {
+		AwaitExpression(awaitPath: any) {
+			const arg = awaitPath.node.argument;
+			if (!t.isCallExpression(arg)) return;
+
+			const info = batchDetector.extractCallInfo(arg);
+			if (!info) return;
+
+			const key = `${info.type}:${info.operation}`;
+
+			// Find first unused call with matching key
+			let matchedIndex = -1;
+			for (let i = 0; i < llmCalls.length; i++) {
+				const original = llmCalls[i]!;
+				if (original.key === key && resultVarByCallIndex.has(i)) {
+					matchedIndex = i;
+					break;
+				}
+			}
+
+			if (matchedIndex === -1) return;
+
+			const resultsVar = resultVarByCallIndex.get(matchedIndex);
+			if (!resultsVar) return;
+			// Remove from map so we don't reuse it
+			resultVarByCallIndex.delete(matchedIndex);
+
+			const resultAccess = t.memberExpression(
+				t.identifier(resultsVar),
+				t.identifier(indexVar),
+				true
+			);
+
+			awaitPath.replaceWith(resultAccess);
+			replacementCount++;
+		},
+		noScope: true,
+	});
+
+	if (replacementCount === 0) {
+		return false;
+	}
+
+	let reconstructBody: t.BlockStatement | t.Expression;
+	if (t.isBlockStatement(clonedBody)) {
+		reconstructBody = clonedBody;
+	} else {
+		reconstructBody = clonedBody as t.Expression;
+	}
+
+	const reconstructMapper = t.arrowFunctionExpression(
+		[t.identifier(param), t.identifier(indexVar)],
+		reconstructBody
+	);
+	reconstructMapper.async = false;
+
+	const reconstructCall = t.callExpression(
+		t.memberExpression(t.cloneNode(array, true), t.identifier('map')),
+		[reconstructMapper]
+	);
+
+	// Build the IIFE with all batch declarations followed by reconstruction
+	const iife = t.callExpression(
+		t.arrowFunctionExpression(
+			[],
+			t.blockStatement([...batchDeclarations, t.returnStatement(reconstructCall)]),
+			true // async
+		),
+		[]
+	);
+
+	const awaitIife = t.awaitExpression(iife);
+
+	path.replaceWith(awaitIife);
+	onTransform();
+	return true;
+}

package/src/transformer/array-transformer-utils.ts

@@ -1,23 +1,168 @@
 import * as t from '@babel/types';
 import { isArrayMethod } from './utils.js';
+import type { BatchParallelDetector } from './batch-detector.js';
+import type { BatchCallInfo } from '../types.js';
+
+export interface LLMCallInfo {
+	callNode: t.CallExpression;
+	callInfo: BatchCallInfo;
+	payloadNode: t.Expression;
+	// Unique key for grouping: "type:operation"
+	key: string;
+}
 
 /**
- * Find LLM call expression in AST node
+ * Collect all identifiers referenced in an expression
  */
-export function findLLMCallExpression(body: t.Node): t.CallExpression | null {
-	let found: t.CallExpression | null = null;
+function collectReferencedIdentifiers(node: t.Node): Set<string> {
+	const identifiers = new Set<string>();
+
+	const visit = (n: t.Node) => {
+		if (t.isIdentifier(n)) {
+			identifiers.add(n.name);
+		}
+
+		// Continue traversing
+		Object.keys(n).forEach((key) => {
+			// Skip 'type' and other metadata fields
+			if (key === 'type' || key === 'loc' || key === 'start' || key === 'end') return;
+			const value = (n as any)[key];
+			if (Array.isArray(value)) {
+				value.forEach((item) => {
+					if (item && typeof item === 'object' && item.type) {
+						visit(item);
+					}
+				});
+			} else if (value && typeof value === 'object' && value.type) {
+				visit(value);
+			}
+		});
+	};
+
+	visit(node);
+	return identifiers;
+}
+
+/**
+ * Collect all variable names declared inside a callback body.
+ * This includes: const/let/var declarations, function parameters, etc.
+ */
+function collectLocalVariables(body: t.Node): Set<string> {
+	const locals = new Set<string>();
 
 	const visit = (node: t.Node) => {
-		if (found) return;
+		// Variable declarations: const x = ..., let y = ..., var z = ...
+		if (t.isVariableDeclaration(node)) {
+			for (const decl of node.declarations) {
+				if (t.isIdentifier(decl.id)) {
+					locals.add(decl.id.name);
+				} else if (t.isObjectPattern(decl.id)) {
+					// Destructuring: const { a, b } = ...
+					for (const prop of decl.id.properties) {
+						if (t.isObjectProperty(prop) && t.isIdentifier(prop.value)) {
+							locals.add(prop.value.name);
+						} else if (t.isRestElement(prop) && t.isIdentifier(prop.argument)) {
+							locals.add(prop.argument.name);
+						}
+					}
+				} else if (t.isArrayPattern(decl.id)) {
+					// Destructuring: const [a, b] = ...
+					for (const elem of decl.id.elements) {
+						if (t.isIdentifier(elem)) {
+							locals.add(elem.name);
+						}
+					}
+				}
+			}
+		}
 
+		// Continue traversing (but don't descend into nested functions - their locals are their own scope)
+		if (t.isFunction(node)) {
+			return; // Don't traverse into nested functions
+		}
+
+		Object.keys(node).forEach((key) => {
+			if (key === 'type' || key === 'loc' || key === 'start' || key === 'end') return;
+			const value = (node as any)[key];
+			if (Array.isArray(value)) {
+				value.forEach((item) => {
+					if (item && typeof item === 'object' && item.type) {
+						visit(item);
+					}
+				});
+			} else if (value && typeof value === 'object' && value.type) {
+				visit(value);
+			}
+		});
+	};
+
+	visit(body);
+	return locals;
+}
+
+/**
+ * Check if LLM call payloads depend on variables that are local to the callback.
+ * This includes:
+ * - Variables computed from other expressions (const x = item.name.toUpperCase())
+ * - Results from previous LLM calls (const title = await atp.llm.call(...))
+ * - Any other locally-defined variable
+ *
+ * When batch transforming, the payload mapper only has access to the array item parameter
+ * and outer scope variables - NOT local variables defined inside the callback.
+ *
+ * @param body - The callback body
+ * @param itemParamName - The name of the array item parameter (e.g., "item")
+ * @param batchDetector - The batch detector for extracting LLM call info
+ * @returns true if there are dependencies that prevent batch transformation
+ */
+export function hasLLMCallDependencies(
+	body: t.Node,
+	batchDetector: BatchParallelDetector,
+	itemParamName?: string
+): boolean {
+	// Find all locally-defined variables in the callback body
+	const localVariables = collectLocalVariables(body);
+
+	// If no local variables, no dependencies possible
+	if (localVariables.size === 0) {
+		return false;
+	}
+
+	// Now check if any LLM call payload references these local variables
+	const allCalls = findAllAwaitedMemberCalls(body);
+
+	for (const call of allCalls) {
+		const payloadNode = batchDetector.extractPayloadNode(call);
+		if (payloadNode) {
+			const referencedIds = collectReferencedIdentifiers(payloadNode);
+			// Check if any referenced identifier is a local variable
+			// (excluding the item parameter which is passed to the payload mapper)
+			for (const id of referencedIds) {
+				if (localVariables.has(id) && id !== itemParamName) {
+					return true; // Found a dependency on a local variable
+				}
+			}
+		}
+	}
+
+	return false;
+}
+
+/**
+ * Find all awaited member expression calls in AST node
+ */
+function findAllAwaitedMemberCalls(body: t.Node): t.CallExpression[] {
+	const calls: t.CallExpression[] = [];
+
+	const visit = (node: t.Node) => {
 		if (t.isAwaitExpression(node) && t.isCallExpression(node.argument)) {
 			const call = node.argument;
 			if (t.isMemberExpression(call.callee)) {
-				found = call;
-				return;
+				calls.push(call);
 			}
 		}
 
+		// Continue traversing
 		Object.keys(node).forEach((key) => {
 			const value = (node as any)[key];
 			if (Array.isArray(value)) {
@@ -33,7 +178,42 @@ export function findLLMCallExpression(body: t.Node): t.CallExpression | null {
 	};
 
 	visit(body);
-	return found;
+	return calls;
+}
+
+/**
+ * Find ALL LLM call expressions in AST node with batch info
+ */
+export function findAllLLMCallExpressions(
+	body: t.Node,
+	batchDetector: BatchParallelDetector
+): LLMCallInfo[] {
+	const allCalls = findAllAwaitedMemberCalls(body);
+	const llmCalls: LLMCallInfo[] = [];
+
+	for (const call of allCalls) {
+		const callInfo = batchDetector.extractCallInfo(call);
+		const payloadNode = batchDetector.extractPayloadNode(call);
+
+		if (callInfo && payloadNode) {
+			llmCalls.push({
+				callNode: call,
+				callInfo,
+				payloadNode,
+				key: `${callInfo.type}:${callInfo.operation}`,
+			});
+		}
+	}
+
+	return llmCalls;
+}
+
+/**
+ * Find first LLM call expression in AST node
+ */
+export function findLLMCallExpression(body: t.Node): t.CallExpression | null {
+	const calls = findAllAwaitedMemberCalls(body);
+	return calls[0] ?? null;
 }
 
 /**

package/src/transformer/array-transformer.ts

@@ -1,9 +1,15 @@
 import * as t from '@babel/types';
 import { BatchOptimizer } from './batch-optimizer.js';
 import { BatchParallelDetector } from './batch-detector.js';
-import { getArrayMethodName, canUseBatchParallel } from './array-transformer-utils.js';
+import {
+	getArrayMethodName,
+	canUseBatchParallel,
+	findLLMCallExpression,
+	hasLLMCallDependencies,
+} from './array-transformer-utils.js';
 import { transformToBatchParallel } from './array-transformer-batch.js';
 import { transformToSequential } from './array-transformer-sequential.js';
+import { transformToBatchWithReconstruction } from './array-transformer-batch-reconstruct.js';
 
 export class ArrayTransformer {
 	private transformCount = 0;
@@ -31,6 +37,43 @@ export class ArrayTransformer {
 		}
 
 		const batchResult = this.batchOptimizer.canBatchArrayMethod(callback);
+		if (!batchResult.canBatch && methodName === 'map') {
+			const reason = batchResult.reason || '';
+			// Try batch-with-reconstruction for:
+			// 1. Object/array returns (would lose structure with simple batch)
+			// 2. Multiple pausable calls (can batch each call type separately)
+			const canTryBatchReconstruct =
+				reason.includes('object expression') ||
+				reason.includes('array expression') ||
+				reason.includes('Multiple pausable calls');
+
+			if (canTryBatchReconstruct) {
+				const llmCall = findLLMCallExpression(callback.body);
+				// Get the item parameter name (e.g., "item" in items.map(async (item) => ...))
+				const itemParam = callback.params[0];
+				const itemParamName = t.isIdentifier(itemParam) ? itemParam.name : undefined;
+				// Check for dependencies on local variables (e.g., computed values, previous LLM results)
+				// If dependencies exist, we can't batch because payload mapper only has access to item + outer scope
+				const hasDependencies = hasLLMCallDependencies(
+					callback.body,
+					this.batchDetector,
+					itemParamName
+				);
+				if (llmCall && !hasDependencies) {
+					const success = transformToBatchWithReconstruction(
+						path,
+						node,
+						methodName,
+						callback,
+						this.batchDetector,
+						() => this.transformCount++
+					);
+					if (success) {
+						return true;
+					}
+				}
+			}
+		}
 
 		if (batchResult.canBatch && canUseBatchParallel(methodName)) {
 			const array = (node.callee as t.MemberExpression).object;

package/src/transformer/index.ts

@@ -171,5 +171,6 @@ export * from './batch-detector.js';
 export * from './batch-optimizer.js';
 export * from './loop-transformer.js';
 export * from './array-transformer.js';
+export * from './array-transformer-batch-reconstruct.js';
 export * from './promise-transformer.js';
 export * from './utils.js';