@mondaydotcomorg/atp-compiler
Version:
Production-ready compiler for transforming async iteration patterns into resumable operations with checkpoint-based state management
55 lines • 2.17 kB
JavaScript
import * as t from '@babel/types';
import { generateUniqueId } from '../runtime/context.js';
import { findLLMCallExpression } from './array-transformer-utils.js';
import { wrapBatchResultIfNeeded } from './array-transformer-wrappers.js';
/**
* Extract batch call info from callback
*/
export function extractBatchCallInfo(callback, batchDetector) {
const paramName = callback.params[0];
if (!t.isIdentifier(paramName)) {
return null;
}
const param = paramName.name;
const llmCall = findLLMCallExpression(callback.body);
if (!llmCall) {
return null;
}
const callInfo = batchDetector.extractCallInfo(llmCall);
if (!callInfo) {
return null;
}
const payloadNode = batchDetector.extractPayloadNode(llmCall);
if (!payloadNode) {
return null;
}
const mapperFunction = t.arrowFunctionExpression([t.identifier(param)], 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),
]));
return { mapperFunction };
}
/**
* Transform simple array method to batch parallel execution
*/
export function transformToBatchParallel(path, node, methodName, callback, batchDetector, onTransform, fallbackTransform) {
const methodId = generateUniqueId(`${methodName}_batch`);
const array = node.callee.object;
const callInfo = extractBatchCallInfo(callback, batchDetector);
if (!callInfo) {
return fallbackTransform();
}
const batchCallsArray = t.callExpression(t.memberExpression(array, t.identifier('map')), [
callInfo.mapperFunction,
]);
const batchCall = t.awaitExpression(t.callExpression(t.memberExpression(t.identifier('__runtime'), t.identifier('batchParallel')), [
batchCallsArray,
t.stringLiteral(methodId),
]));
const finalCall = wrapBatchResultIfNeeded(batchCall, methodName, array, methodId);
path.replaceWith(finalCall);
onTransform();
return true;
}
//# sourceMappingURL=array-transformer-batch.js.map