@mondaydotcomorg/atp-compiler 0.17.14

package/src/transformer/detector.ts

@@ -0,0 +1,171 @@
+import { parse } from '@babel/parser';
+import _traverse from '@babel/traverse';
+// @ts-ignore - CommonJS/ESM compatibility
+const traverse =
+	typeof (_traverse as any).default === 'function' ? (_traverse as any).default : _traverse;
+import * as t from '@babel/types';
+import type { DetectionResult, AsyncPattern } from '../types.js';
+import { containsAwait, isArrayMethod, isPausableCallExpression } from './utils.js';
+
+export class AsyncIterationDetector {
+	detect(code: string): DetectionResult {
+		const patterns: AsyncPattern[] = [];
+		let batchableParallel = false;
+
+		try {
+			const ast = parse(code, {
+				sourceType: 'module',
+				plugins: ['typescript'],
+				allowAwaitOutsideFunction: true,
+				allowReturnOutsideFunction: true,
+			});
+
+			traverse(ast, {
+				ForOfStatement: (path: any) => {
+					if (containsAwait(path.node.body)) {
+						patterns.push('for-of-await');
+					}
+				},
+
+				WhileStatement: (path: any) => {
+					if (containsAwait(path.node.body)) {
+						patterns.push('while-await');
+					}
+				},
+
+				CallExpression: (path: any) => {
+					const node = path.node;
+
+					if (isArrayMethod(node, 'map')) {
+						const callback = node.arguments[0];
+						if (callback && t.isFunction(callback) && callback.async) {
+							patterns.push('map-async');
+						}
+					}
+
+					if (isArrayMethod(node, 'forEach')) {
+						const callback = node.arguments[0];
+						if (callback && t.isFunction(callback) && callback.async) {
+							patterns.push('forEach-async');
+						}
+					}
+
+					if (isArrayMethod(node, 'filter')) {
+						const callback = node.arguments[0];
+						if (callback && t.isFunction(callback) && callback.async) {
+							patterns.push('filter-async');
+						}
+					}
+
+					if (isArrayMethod(node, 'reduce')) {
+						const callback = node.arguments[0];
+						if (callback && t.isFunction(callback) && callback.async) {
+							patterns.push('reduce-async');
+						}
+					}
+
+					if (isArrayMethod(node, 'find')) {
+						const callback = node.arguments[0];
+						if (callback && t.isFunction(callback) && callback.async) {
+							patterns.push('find-async');
+						}
+					}
+
+					if (isArrayMethod(node, 'some')) {
+						const callback = node.arguments[0];
+						if (callback && t.isFunction(callback) && callback.async) {
+							patterns.push('some-async');
+						}
+					}
+
+					if (isArrayMethod(node, 'every')) {
+						const callback = node.arguments[0];
+						if (callback && t.isFunction(callback) && callback.async) {
+							patterns.push('every-async');
+						}
+					}
+
+					if (isArrayMethod(node, 'flatMap')) {
+						const callback = node.arguments[0];
+						if (callback && t.isFunction(callback) && callback.async) {
+							patterns.push('flatMap-async');
+						}
+					}
+
+					if (this.isPromiseAll(node)) {
+						patterns.push('promise-all');
+						if (this.canBatchPromiseAll(node)) {
+							batchableParallel = true;
+						}
+					}
+
+					if (this.isPromiseAllSettled(node)) {
+						patterns.push('promise-allSettled');
+					}
+				},
+			});
+
+			return {
+				needsTransform: patterns.length > 0,
+				patterns: [...new Set(patterns)],
+				batchableParallel,
+			};
+		} catch (error) {
+			return {
+				needsTransform: false,
+				patterns: [],
+				batchableParallel: false,
+			};
+		}
+	}
+
+	private isPromiseAll(node: t.CallExpression): boolean {
+		const callee = node.callee;
+		return (
+			t.isMemberExpression(callee) &&
+			t.isIdentifier(callee.object, { name: 'Promise' }) &&
+			t.isIdentifier(callee.property, { name: 'all' })
+		);
+	}
+
+	private isPromiseAllSettled(node: t.CallExpression): boolean {
+		const callee = node.callee;
+		return (
+			t.isMemberExpression(callee) &&
+			t.isIdentifier(callee.object, { name: 'Promise' }) &&
+			t.isIdentifier(callee.property, { name: 'allSettled' })
+		);
+	}
+
+	private canBatchPromiseAll(node: t.CallExpression): boolean {
+		const arrayArg = node.arguments[0];
+
+		if (!t.isArrayExpression(arrayArg)) {
+			return false;
+		}
+
+		if (arrayArg.elements.length === 0) {
+			return false;
+		}
+
+		return arrayArg.elements.every((el) => {
+			if (!el || t.isSpreadElement(el)) {
+				return false;
+			}
+
+			return this.isDirectPausableCall(el);
+		});
+	}
+
+	private isDirectPausableCall(node: t.Node): boolean {
+		if (t.isAwaitExpression(node)) {
+			node = node.argument;
+		}
+
+		if (!t.isCallExpression(node)) {
+			return false;
+		}
+
+		return isPausableCallExpression(node);
+	}
+}

package/src/transformer/index.ts

@@ -0,0 +1,155 @@
+import { parse } from '@babel/parser';
+import _traverse from '@babel/traverse';
+const traverse = (_traverse as any).default || _traverse;
+import _generate from '@babel/generator';
+const generate = (_generate as any).default || _generate;
+import * as t from '@babel/types';
+import { AsyncIterationDetector } from './detector.js';
+import { LoopTransformer } from './loop-transformer.js';
+import { ArrayTransformer } from './array-transformer.js';
+import { PromiseTransformer } from './promise-transformer.js';
+import type { TransformResult, CompilerConfig, TransformMetadata } from '../types.js';
+import { DEFAULT_COMPILER_CONFIG } from '../types.js';
+import { TransformationError } from '../runtime/errors.js';
+import { resetIdCounter } from '../runtime/context.js';
+
+export class ATPCompiler {
+	private config: CompilerConfig;
+	private detector: AsyncIterationDetector;
+	private loopTransformer: LoopTransformer;
+	private arrayTransformer: ArrayTransformer;
+	private promiseTransformer: PromiseTransformer;
+
+	constructor(config: Partial<CompilerConfig> = {}) {
+		this.config = { ...DEFAULT_COMPILER_CONFIG, ...config };
+		this.detector = new AsyncIterationDetector();
+		this.loopTransformer = new LoopTransformer(this.config.batchSizeThreshold);
+		this.arrayTransformer = new ArrayTransformer(this.config.batchSizeThreshold);
+		this.promiseTransformer = new PromiseTransformer(this.config.enableBatchParallel);
+	}
+
+	detect(code: string) {
+		return this.detector.detect(code);
+	}
+
+	transform(code: string): TransformResult {
+		resetIdCounter();
+
+		const detection = this.detector.detect(code);
+
+		if (!detection.needsTransform) {
+			return {
+				code,
+				transformed: false,
+				patterns: [],
+				metadata: {
+					loopCount: 0,
+					arrayMethodCount: 0,
+					parallelCallCount: 0,
+					batchableCount: 0,
+				},
+			};
+		}
+
+		try {
+			const ast = parse(code, {
+				sourceType: 'module',
+				plugins: ['typescript'],
+				allowAwaitOutsideFunction: true,
+				allowReturnOutsideFunction: true,
+			});
+
+			this.loopTransformer.resetTransformCount();
+			this.arrayTransformer.resetTransformCount();
+			this.promiseTransformer.resetTransformCount();
+
+			traverse(ast, {
+				ForOfStatement: (path: any) => {
+					this.loopTransformer.transformForOfLoop(path);
+				},
+
+				WhileStatement: (path: any) => {
+					this.loopTransformer.transformWhileLoop(path);
+				},
+
+				ForStatement: (path: any) => {
+					this.loopTransformer.transformForLoop(path);
+				},
+
+				CallExpression: (path: any) => {
+					if (this.isArrayMethodCall(path.node)) {
+						this.arrayTransformer.transformArrayMethod(path);
+					} else if (this.isPromiseAllCall(path.node)) {
+						this.promiseTransformer.transformPromiseAll(path);
+					} else if (this.isPromiseAllSettledCall(path.node)) {
+						this.promiseTransformer.transformPromiseAllSettled(path);
+					}
+				},
+			});
+
+			const output = generate(ast, {
+				sourceMaps: false,
+				retainLines: true,
+				comments: true,
+			});
+
+			const metadata: TransformMetadata = {
+				loopCount: this.loopTransformer.getTransformCount(),
+				arrayMethodCount: this.arrayTransformer.getTransformCount(),
+				parallelCallCount: this.promiseTransformer.getTransformCount(),
+				batchableCount: detection.batchableParallel ? 1 : 0,
+			};
+
+			return {
+				code: output.code,
+				transformed: true,
+				patterns: detection.patterns,
+				metadata,
+			};
+		} catch (error) {
+			const message = error instanceof Error ? error.message : String(error);
+			throw new TransformationError(message, code, 'unknown');
+		}
+	}
+
+	private isArrayMethodCall(node: t.CallExpression): boolean {
+		if (!t.isMemberExpression(node.callee)) {
+			return false;
+		}
+
+		const property = node.callee.property;
+		if (!t.isIdentifier(property)) {
+			return false;
+		}
+
+		const arrayMethods = ['map', 'forEach', 'filter', 'reduce', 'find', 'some', 'every', 'flatMap'];
+
+		return arrayMethods.includes(property.name);
+	}
+
+	private isPromiseAllCall(node: t.CallExpression): boolean {
+		const callee = node.callee;
+		return (
+			t.isMemberExpression(callee) &&
+			t.isIdentifier(callee.object, { name: 'Promise' }) &&
+			t.isIdentifier(callee.property, { name: 'all' })
+		);
+	}
+
+	private isPromiseAllSettledCall(node: t.CallExpression): boolean {
+		const callee = node.callee;
+		return (
+			t.isMemberExpression(callee) &&
+			t.isIdentifier(callee.object, { name: 'Promise' }) &&
+			t.isIdentifier(callee.property, { name: 'allSettled' })
+		);
+	}
+}
+
+export * from './detector.js';
+export * from './batch-detector.js';
+export * from './batch-optimizer.js';
+export * from './loop-transformer.js';
+export * from './array-transformer.js';
+export * from './promise-transformer.js';
+export * from './utils.js';

package/src/transformer/loop-transformer.ts

@@ -0,0 +1,285 @@
+import * as t from '@babel/types';
+import { generateUniqueId } from '../runtime/context.js';
+import { containsAwait } from './utils.js';
+import { BatchOptimizer } from './batch-optimizer.js';
+import { BatchParallelDetector } from './batch-detector.js';
+
+export class LoopTransformer {
+	private transformCount = 0;
+	private batchOptimizer: BatchOptimizer;
+	private batchDetector: BatchParallelDetector;
+	private batchSizeThreshold: number;
+
+	constructor(batchSizeThreshold: number = 10) {
+		this.batchOptimizer = new BatchOptimizer();
+		this.batchDetector = new BatchParallelDetector();
+		this.batchSizeThreshold = batchSizeThreshold;
+	}
+
+	transformForOfLoop(path: any): boolean {
+		const node = path.node as t.ForOfStatement;
+
+		if (!containsAwait(node.body)) {
+			return false;
+		}
+
+		const batchResult = this.batchOptimizer.canBatchForOfLoop(node);
+
+		if (batchResult.canBatch) {
+			const decision = this.batchOptimizer.makeSmartBatchDecision(
+				'for...of',
+				batchResult,
+				node.right,
+				this.batchSizeThreshold
+			);
+
+			if (decision.shouldBatch) {
+				return this.transformForOfToBatch(path, node);
+			}
+		}
+
+		return this.transformForOfToSequential(path, node);
+	}
+
+	/**
+	 * Transform simple for...of to batch parallel
+	 */
+	private transformForOfToBatch(path: any, node: t.ForOfStatement): boolean {
+		const loopId = generateUniqueId('for_of_batch');
+		const left = node.left;
+		const right = node.right;
+
+		let paramName: string;
+		if (t.isVariableDeclaration(left)) {
+			const id = left.declarations[0]?.id;
+			paramName = t.isIdentifier(id) ? id.name : 'item';
+		} else if (t.isIdentifier(left)) {
+			paramName = left.name;
+		} else {
+			paramName = 'item';
+		}
+
+		const bodyStatements = t.isBlockStatement(node.body) ? node.body.body : [node.body];
+		const callback = t.arrowFunctionExpression(
+			[t.identifier(paramName)],
+			t.blockStatement(bodyStatements),
+			true
+		);
+
+		const llmCall = this.findLLMCallInBody(node.body);
+		if (!llmCall) {
+			return this.transformForOfToSequential(path, node);
+		}
+
+		const callInfo = this.batchDetector.extractCallInfo(llmCall);
+		if (!callInfo) {
+			return this.transformForOfToSequential(path, node);
+		}
+
+		const payloadNode = this.batchDetector.extractPayloadNode(llmCall);
+		if (!payloadNode) {
+			return this.transformForOfToSequential(path, node);
+		}
+
+		const batchCallsArray = t.callExpression(t.memberExpression(right, t.identifier('map')), [
+			t.arrowFunctionExpression(
+				[t.identifier(paramName)],
+				t.objectExpression([
+					t.objectProperty(t.identifier('type'), t.stringLiteral(callInfo.type)),
+					t.objectProperty(t.identifier('operation'), t.stringLiteral(callInfo.operation)),
+					t.objectProperty(t.identifier('payload'), payloadNode),
+				])
+			),
+		]);
+
+		const batchCall = t.awaitExpression(
+			t.callExpression(
+				t.memberExpression(t.identifier('__runtime'), t.identifier('batchParallel')),
+				[batchCallsArray, t.stringLiteral(loopId)]
+			)
+		);
+
+		path.replaceWith(t.expressionStatement(batchCall));
+		this.transformCount++;
+		return true;
+	}
+
+	/**
+	 * Transform for...of to sequential with checkpoints (fallback)
+	 */
+	private transformForOfToSequential(path: any, node: t.ForOfStatement): boolean {
+		const loopId = generateUniqueId('for_of');
+		const left = node.left;
+		const right = node.right;
+
+		let paramName: string;
+		if (t.isVariableDeclaration(left)) {
+			const id = left.declarations[0]?.id;
+			paramName = t.isIdentifier(id) ? id.name : 'item';
+		} else if (t.isIdentifier(left)) {
+			paramName = left.name;
+		} else {
+			paramName = 'item';
+		}
+
+		const bodyStatements = t.isBlockStatement(node.body) ? node.body.body : [node.body];
+
+		const callbackFn = t.arrowFunctionExpression(
+			[t.identifier(paramName), t.identifier('__index')],
+			t.blockStatement(bodyStatements),
+			true
+		);
+
+		const runtimeCall = t.awaitExpression(
+			t.callExpression(
+				t.memberExpression(t.identifier('__runtime'), t.identifier('resumableForOf')),
+				[right, callbackFn, t.stringLiteral(loopId)]
+			)
+		);
+
+		path.replaceWith(t.expressionStatement(runtimeCall));
+		this.transformCount++;
+		return true;
+	}
+
+	/**
+	 * Find LLM call in loop body
+	 */
+	private findLLMCallInBody(body: t.Statement | t.BlockStatement): t.CallExpression | null {
+		let found: t.CallExpression | null = null;
+
+		const visit = (node: t.Node) => {
+			if (found) return;
+
+			if (t.isAwaitExpression(node) && t.isCallExpression(node.argument)) {
+				const call = node.argument;
+				if (t.isMemberExpression(call.callee)) {
+					found = call;
+					return;
+				}
+			}
+
+			Object.keys(node).forEach((key) => {
+				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 found;
+	}
+
+	transformWhileLoop(path: any): boolean {
+		const node = path.node as t.WhileStatement;
+
+		if (!containsAwait(node.body)) {
+			return false;
+		}
+
+		const loopId = generateUniqueId('while');
+
+		const conditionFn = t.arrowFunctionExpression([], node.test, false);
+
+		const bodyStatements = t.isBlockStatement(node.body) ? node.body.body : [node.body];
+
+		const bodyFn = t.arrowFunctionExpression(
+			[t.identifier('__iteration')],
+			t.blockStatement(bodyStatements),
+			true
+		);
+
+		const runtimeCall = t.awaitExpression(
+			t.callExpression(
+				t.memberExpression(t.identifier('__runtime'), t.identifier('resumableWhile')),
+				[conditionFn, bodyFn, t.stringLiteral(loopId)]
+			)
+		);
+
+		path.replaceWith(t.expressionStatement(runtimeCall));
+		this.transformCount++;
+		return true;
+	}
+
+	transformForLoop(path: any): boolean {
+		const node = path.node as t.ForStatement;
+
+		if (!containsAwait(node.body)) {
+			return false;
+		}
+
+		if (!node.init || !node.test || !node.update) {
+			return false;
+		}
+
+		const loopId = generateUniqueId('for');
+
+		let initValue: t.Expression = t.numericLiteral(0);
+		let loopVar = '__i';
+
+		if (t.isVariableDeclaration(node.init)) {
+			const decl = node.init.declarations[0];
+			if (decl && t.isIdentifier(decl.id) && decl.init) {
+				loopVar = decl.id.name;
+				initValue = decl.init;
+			}
+		}
+
+		const conditionFn = t.arrowFunctionExpression([t.identifier(loopVar)], node.test, false);
+
+		const bodyStatements = t.isBlockStatement(node.body) ? node.body.body : [node.body];
+
+		const bodyFn = t.arrowFunctionExpression(
+			[t.identifier(loopVar)],
+			t.blockStatement(bodyStatements),
+			true
+		);
+
+		let incrementFn: t.Expression;
+		if (t.isUpdateExpression(node.update)) {
+			if (node.update.operator === '++') {
+				incrementFn = t.arrowFunctionExpression(
+					[t.identifier(loopVar)],
+					t.binaryExpression('+', t.identifier(loopVar), t.numericLiteral(1)),
+					false
+				);
+			} else if (node.update.operator === '--') {
+				incrementFn = t.arrowFunctionExpression(
+					[t.identifier(loopVar)],
+					t.binaryExpression('-', t.identifier(loopVar), t.numericLiteral(1)),
+					false
+				);
+			} else {
+				return false;
+			}
+		} else {
+			return false;
+		}
+
+		const runtimeCall = t.awaitExpression(
+			t.callExpression(
+				t.memberExpression(t.identifier('__runtime'), t.identifier('resumableForLoop')),
+				[initValue, conditionFn, incrementFn, bodyFn, t.stringLiteral(loopId)]
+			)
+		);
+
+		path.replaceWith(t.expressionStatement(runtimeCall));
+		this.transformCount++;
+		return true;
+	}
+
+	getTransformCount(): number {
+		return this.transformCount;
+	}
+
+	resetTransformCount(): void {
+		this.transformCount = 0;
+	}
+}