Add defer functionality to generator executor (#64157)

Author: weswigham

packages/typescript/src/api/async/api.ts

@@ -221,8 +221,8 @@ export interface TranspileOutput {
 }
 
 // @sync-only-start
-// export { all } from "./generatorSupport.ts";
-// import {all, type ExecutedGeneratorsResults, type APIRequestGenerator} from "./generatorSupport.ts";
+// export { all, defer, type APIRequestGenerator, type AnyAPIRequestGenerator, type DeferredAPIRequestGenerator, type ExecutedGeneratorsResults } from "./generatorSupport.ts";
+// import {executeRequestGenerators, type ExecutedGeneratorsResults, type AnyAPIRequestGenerator} from "./generatorSupport.ts";
 // @sync-only-end
 
 export class API<FromLSP extends boolean = false> implements FormatDiagnosticsHost {
@@ -259,13 +259,8 @@ export class API<FromLSP extends boolean = false> implements FormatDiagnosticsHo
     }
     // @sync-skip-block-end
     // @sync-only-start
-    // batch<T extends readonly APIRequestGenerator[]>(...requestGenerators: T): ExecutedGeneratorsResults<T> {
-    //     const batches = all(...requestGenerators);
-    //     let state = batches.next();
-    //     while (!state.done) {
-    //         state = batches.next(this.client.batchRequests(state.value).responses);
-    //     }
-    //     return state.value;
+    // batch<T extends readonly AnyAPIRequestGenerator[]>(...requestGenerators: T): ExecutedGeneratorsResults<T> {
+    //     return executeRequestGenerators(requestGenerators, requests => this.client.batchRequests(requests).responses);
     // }
     // @sync-only-end
 

packages/typescript/src/api/sync/api.ts

@@ -237,11 +237,11 @@ export interface TranspileOutput {
     sourceMapText?: string;
 }
 
-export { all } from "./generatorSupport.ts";
+export { all, type AnyAPIRequestGenerator, type APIRequestGenerator, defer, type DeferredAPIRequestGenerator, type ExecutedGeneratorsResults } from "./generatorSupport.ts";
 import {
-    all,
-    type APIRequestGenerator,
+    type AnyAPIRequestGenerator,
     type ExecutedGeneratorsResults,
+    executeRequestGenerators,
 } from "./generatorSupport.ts";
 
 export class API<FromLSP extends boolean = false> implements FormatDiagnosticsHost {
@@ -287,13 +287,8 @@ export class API<FromLSP extends boolean = false> implements FormatDiagnosticsHo
         );
     }
 
-    batch<T extends readonly APIRequestGenerator[]>(...requestGenerators: T): ExecutedGeneratorsResults<T> {
-        const batches = all(...requestGenerators);
-        let state = batches.next();
-        while (!state.done) {
-            state = batches.next(this.client.batchRequests(state.value).responses);
-        }
-        return state.value;
+    batch<T extends readonly AnyAPIRequestGenerator[]>(...requestGenerators: T): ExecutedGeneratorsResults<T> {
+        return executeRequestGenerators(requestGenerators, requests => this.client.batchRequests(requests).responses);
     }
 
     private get ensureInitialized(): {

packages/typescript/src/api/sync/generatorSupport.ts

@@ -4,7 +4,7 @@ import type {
     APIResponse,
 } from "../proto.ts";
 
-export function cacheGeneratorMethod<Sync extends (...args: any[]) => any, Gen extends (...args: any[]) => APIRequestGenerator>(
+export function cacheGeneratorMethod<Sync extends (...args: any[]) => any, Gen extends (...args: any[]) => AnyAPIRequestGenerator>(
     owner: object,
     name: PropertyKey,
     sync: Sync,
@@ -22,82 +22,144 @@ export function apiRequest<Method extends keyof APIMethodInfo>(
 export function* apiRequest(method: PropertyKey, params: unknown): Generator<{ method: PropertyKey; params: unknown; }, unknown, unknown> {
     return yield { method, params };
 }
-export type APIRequestGenerator<Return = any> = Generator<APIRequest | readonly APIRequest[], Return, any>;
+const deferredGeneratorMarker: unique symbol = Symbol();
+interface DeferredAPIRequest {
+    readonly deferred: APIRequestGenerator;
+}
+type APIRequestGeneratorYield = APIRequest | readonly APIRequest[] | DeferredAPIRequest;
+export type APIRequestGenerator<Return = any> = Generator<APIRequestGeneratorYield, Return, any>;
+export type DeferredAPIRequestGenerator = Generator<DeferredAPIRequest, void, unknown> & { readonly [deferredGeneratorMarker]: true; };
+export type AnyAPIRequestGenerator<Return = any> = APIRequestGenerator<Return> | DeferredAPIRequestGenerator;
 type GeneratorReturn<T> = T extends Generator<any, infer R, any> ? R : never;
-export type ExecutedGeneratorsResults<T extends readonly APIRequestGenerator[]> = {
-    [K in keyof T]: GeneratorReturn<T[K]>;
-};
+export type ExecutedGeneratorsResults<T extends readonly AnyAPIRequestGenerator[]> = number extends T["length"] ? GeneratorReturn<Exclude<T[number], DeferredAPIRequestGenerator>>[]
+    : T extends readonly [infer Head extends AnyAPIRequestGenerator, ...infer Tail extends readonly AnyAPIRequestGenerator[]] ? Head extends DeferredAPIRequestGenerator ? ExecutedGeneratorsResults<Tail> : [GeneratorReturn<Head>, ...ExecutedGeneratorsResults<Tail>]
+    : [];
 
 interface GeneratorResponse {
     result: unknown;
     error?: string | undefined;
 }
 
-export function all<const T extends readonly APIRequestGenerator[]>(
-    ...requestGenerators: T
-): Generator<APIRequest[], ExecutedGeneratorsResults<T>, readonly GeneratorResponse[]>;
-export function* all<T extends readonly APIRequestGenerator[]>(
-    ...requestGenerators: T
-): Generator<APIRequest[], ExecutedGeneratorsResults<T>, readonly GeneratorResponse[]> {
-    const results: any[] = [];
-    const requestObjects: (APIRequest | readonly APIRequest[] | undefined)[] = [];
-    const completedIndices = new Set<number>();
-    for (let i = 0; i < requestGenerators.length; i++) {
-        const state = requestGenerators[i].next();
+interface RequestRunnerOptions {
+    executeDeferred?: boolean;
+}
+
+function createRequestRunner<T extends readonly AnyAPIRequestGenerator[]>(requestGenerators: T, options: RequestRunnerOptions = {}) {
+    const registeredGenerators = new Set<AnyAPIRequestGenerator>();
+    const requestsByGenerator = new Map<AnyAPIRequestGenerator, APIRequestGeneratorYield>();
+    const resultsByGenerator = new Map<AnyAPIRequestGenerator, unknown>();
+    for (const generator of requestGenerators) {
+        addGenerator(generator);
+    }
+    const requestRounds = runRequestRounds();
+    return { requestRounds, getResults };
+
+    function advanceGenerator(generator: AnyAPIRequestGenerator, value?: unknown, error?: string): void {
+        let state = error === undefined
+            ? generator.next(value)
+            : generator.throw(new Error(error));
+        while (!state.done && isDeferredAPIRequest(state.value) && options.executeDeferred) {
+            addGenerator(state.value.deferred);
+            state = generator.next();
+        }
         if (state.done) {
-            results[i] = state.value;
-            completedIndices.add(i);
+            requestsByGenerator.delete(generator);
+            resultsByGenerator.set(generator, state.value);
         }
         else {
-            requestObjects[i] = state.value;
+            requestsByGenerator.set(generator, state.value);
         }
     }
-    while (completedIndices.size < requestGenerators.length) {
-        const requests: APIRequest[] = [];
-        const responseIndices: (number | readonly number[] | undefined)[] = [];
-        const responseIndexByDeduplicationKey = new Map<string, number>();
-        const addRequest = (request: APIRequest): number => {
-            const deduplicationKey = getRequestDeduplicationKey(request);
-            let responseIndex = deduplicationKey === undefined ? undefined : responseIndexByDeduplicationKey.get(deduplicationKey);
-            if (responseIndex === undefined) {
-                responseIndex = requests.length;
-                requests.push(request);
-                if (deduplicationKey !== undefined) responseIndexByDeduplicationKey.set(deduplicationKey, responseIndex);
-            }
-            return responseIndex;
-        };
-        for (let i = 0; i < requestGenerators.length; i++) {
-            if (completedIndices.has(i)) continue;
 
-            const request = requestObjects[i]!;
-            responseIndices[i] = isRequestGroup(request) ? request.map(addRequest) : addRequest(request);
-        }
+    function addGenerator(generator: AnyAPIRequestGenerator): void {
+        if (registeredGenerators.has(generator)) throw new Error("Cannot execute the same generator instance more than once");
+        registeredGenerators.add(generator);
+        advanceGenerator(generator);
+    }
 
-        const responses = yield requests;
-        for (let i = 0; i < requestGenerators.length; i++) {
-            if (completedIndices.has(i)) continue;
-
-            const requestGenerator = requestGenerators[i];
-            const responseIndex = responseIndices[i]!;
-            let state: IteratorResult<APIRequest | readonly APIRequest[], any>;
-            if (typeof responseIndex === "number") {
-                const result = responses[responseIndex];
-                state = result.error
-                    ? requestGenerator.throw(new Error(result.error))
-                    : requestGenerator.next(result.result);
+    function* runRequestRounds(): Generator<APIRequest[] | DeferredAPIRequest, void, readonly GeneratorResponse[]> {
+        while (requestsByGenerator.size) {
+            for (const generator of registeredGenerators) {
+                let request = requestsByGenerator.get(generator);
+                while (request && isDeferredAPIRequest(request)) {
+                    yield request;
+                    advanceGenerator(generator);
+                    request = requestsByGenerator.get(generator);
+                }
             }
-            else {
-                state = requestGenerator.next(responseIndex.map(index => responses[index]));
+            if (!requestsByGenerator.size) break;
+
+            const requests: APIRequest[] = [];
+            const responseIndexByDeduplicationKey = new Map<string, number>();
+            const addRequest = (request: APIRequest): number => {
+                const deduplicationKey = getRequestDeduplicationKey(request);
+                let responseIndex = deduplicationKey === undefined ? undefined : responseIndexByDeduplicationKey.get(deduplicationKey);
+                if (responseIndex === undefined) {
+                    responseIndex = requests.length;
+                    requests.push(request);
+                    if (deduplicationKey !== undefined) responseIndexByDeduplicationKey.set(deduplicationKey, responseIndex);
+                }
+                return responseIndex;
+            };
+            // TODO: Use Iterator.prototype.filter when target >= ES2025
+            const roundGenerators = [...registeredGenerators].filter(generator => requestsByGenerator.has(generator));
+            const responseIndices = new Map<AnyAPIRequestGenerator, number | readonly number[]>();
+            for (const generator of roundGenerators) {
+                const request = requestsByGenerator.get(generator) as APIRequest | readonly APIRequest[];
+                responseIndices.set(generator, isRequestGroup(request) ? request.map(addRequest) : addRequest(request));
             }
-            if (state.done) {
-                results[i] = state.value;
-                completedIndices.add(i);
-                requestObjects[i] = undefined;
+
+            const responses = yield requests;
+            for (const generator of roundGenerators) {
+                const responseIndex = responseIndices.get(generator)!;
+                if (typeof responseIndex === "number") {
+                    const result = responses[responseIndex];
+                    advanceGenerator(generator, result.result, result.error || undefined);
+                }
+                else {
+                    advanceGenerator(generator, responseIndex.map(index => responses[index]));
+                }
             }
-            else requestObjects[i] = state.value;
         }
     }
-    return results as any[] as ExecutedGeneratorsResults<T>;
+
+    function getResults(): ExecutedGeneratorsResults<T> {
+        return requestGenerators
+            .filter(generator => !isDeferredGenerator(generator))
+            .map(generator => resultsByGenerator.get(generator)) as ExecutedGeneratorsResults<T>;
+    }
+}
+
+export function all<const T extends readonly AnyAPIRequestGenerator[]>(
+    ...requestGenerators: T
+): APIRequestGenerator<ExecutedGeneratorsResults<T>>;
+export function* all<T extends readonly AnyAPIRequestGenerator[]>(
+    ...requestGenerators: T
+): APIRequestGenerator<ExecutedGeneratorsResults<T>> {
+    const { requestRounds, getResults } = createRequestRunner(requestGenerators);
+    yield* requestRounds;
+    return getResults();
+}
+
+export function executeRequestGenerators<T extends readonly AnyAPIRequestGenerator[]>(
+    requestGenerators: T,
+    executeRequests: (requests: APIRequest[]) => readonly GeneratorResponse[],
+): ExecutedGeneratorsResults<T> {
+    const { requestRounds, getResults } = createRequestRunner(requestGenerators, { executeDeferred: true });
+    let state = requestRounds.next();
+    while (!state.done) {
+        if (isDeferredAPIRequest(state.value)) throw new Error("Unexpected deferred request");
+        state = requestRounds.next(executeRequests(state.value));
+    }
+    return getResults();
+}
+
+function isDeferredAPIRequest(request: APIRequestGeneratorYield): request is DeferredAPIRequest {
+    return !Array.isArray(request) && "deferred" in request;
+}
+
+function isDeferredGenerator(generator: AnyAPIRequestGenerator): generator is DeferredAPIRequestGenerator {
+    return deferredGeneratorMarker in generator;
 }
 
 function isRequestGroup(request: APIRequest | readonly APIRequest[]): request is readonly APIRequest[] {
@@ -112,3 +174,11 @@ function getRequestDeduplicationKey(request: APIRequest): string | undefined {
             return undefined;
     }
 }
+
+export function defer(gen: APIRequestGenerator): DeferredAPIRequestGenerator {
+    const deferred = (function* (): Generator<DeferredAPIRequest, void, unknown> {
+        yield { deferred: gen };
+    })() as DeferredAPIRequestGenerator;
+    Object.defineProperty(deferred, deferredGeneratorMarker, { value: true });
+    return deferred;
+}