@openrewrite/rewrite 8.62.3 → 8.62.4

package/src/rpc/queue.ts

@@ -46,37 +46,58 @@ export interface RpcCodec<T> {
  * A registry for managing RPC codecs based on object types.
  */
 export class RpcCodecs {
-    private static codecs = new Map<string, RpcCodec<any>>();
+    private static nonTreeCodecs = new Map<string, RpcCodec<any>>();
+
+    /**
+     * The first key is on sourceFileType and the second on object type
+     */
+    private static treeCodecs = new Map<string, Map<string, RpcCodec<any>>>();
 
     /**
      * Registers an RPC codec for a given type.
      *
      * @param type - The string identifier of the object type.
      * @param codec - The codec implementation to be registered.
+     * @param sourceFileType The source file type of the source file containing (or will contain) this element.
      */
-    static registerCodec(type: string, codec: RpcCodec<any>): void {
-        this.codecs.set(type, codec);
+    static registerCodec(type: string, codec: RpcCodec<any>, sourceFileType?: string): void {
+        if (sourceFileType) {
+            let codecsForSourceFile = this.treeCodecs.get(sourceFileType);
+            if (!codecsForSourceFile) {
+                codecsForSourceFile = new Map<string, RpcCodec<any>>();
+                this.treeCodecs.set(sourceFileType, codecsForSourceFile);
+            }
+            codecsForSourceFile.set(type, codec);
+        } else {
+            this.nonTreeCodecs.set(type, codec);
+        }
     }
 
     /**
      * Retrieves the registered codec for a given type.
      *
      * @param type - The string identifier of the object type.
+     * @param sourceFileType The source file type of the source file containing (or will contain) this element.
      * @returns The corresponding `RpcCodec`, or `undefined` if not found.
      */
-    static forType(type: string): RpcCodec<any> | undefined {
-        return this.codecs.get(type);
+    static forType(type: string, sourceFileType?: string): RpcCodec<any> | undefined {
+        if (sourceFileType) {
+            const treeCodec = this.treeCodecs.get(sourceFileType)?.get(type);
+            return treeCodec || this.nonTreeCodecs.get(type);
+        }
+        return this.nonTreeCodecs.get(type);
     }
 
     /**
      * Determines the appropriate codec for an instance based on its `kind` property.
      *
      * @param before - The object instance to find a codec for.
+     * @param sourceFileType The source file type of the source file containing (or will contain) this element.
      * @returns The corresponding `RpcCodec`, or `undefined` if no matching codec is found.
      */
-    static forInstance(before: any): RpcCodec<any> | undefined {
+    static forInstance(before: any, sourceFileType?: string): RpcCodec<any> | undefined {
         if (before !== undefined && before !== null && typeof before === "object" && "kind" in before) {
-            return RpcCodecs.forType(before["kind"] as string);
+            return RpcCodecs.forType(before["kind"] as string, sourceFileType);
         }
     }
 }
@@ -86,7 +107,9 @@ export class RpcSendQueue {
 
     private before?: any;
 
-    constructor(private readonly refs: ReferenceMap, private readonly trace: boolean) {
+    constructor(private readonly refs: ReferenceMap,
+                private readonly sourceFileType: string | undefined,
+                private readonly trace: boolean) {
     }
 
     async generate(after: any, before: any): Promise<RpcObjectData[]> {
@@ -132,7 +155,7 @@ export class RpcSendQueue {
             } else if (after === undefined) {
                 this.put({state: RpcObjectState.DELETE});
             } else {
-                let afterCodec = onChange ? undefined : RpcCodecs.forInstance(after);
+                let afterCodec = onChange ? undefined : RpcCodecs.forInstance(after, this.sourceFileType);
                 this.put({state: RpcObjectState.CHANGE, value: onChange || afterCodec ? undefined : after});
                 await this.doChange(after, before, onChange, afterCodec);
             }
@@ -161,7 +184,7 @@ export class RpcSendQueue {
                         this.put({state: RpcObjectState.NO_CHANGE});
                     } else {
                         this.put({state: RpcObjectState.CHANGE});
-                        await this.doChange(anAfter, aBefore, onChangeRun, RpcCodecs.forInstance(anAfter));
+                        await this.doChange(anAfter, aBefore, onChangeRun, RpcCodecs.forInstance(anAfter, this.sourceFileType));
                     }
                 }
             }
@@ -199,7 +222,7 @@ export class RpcSendQueue {
             }
             ref = this.refs.create(after);
         }
-        let afterCodec = onChange ? undefined : RpcCodecs.forInstance(after);
+        let afterCodec = onChange ? undefined : RpcCodecs.forInstance(after, this.sourceFileType);
         this.put({
             state: RpcObjectState.ADD,
             valueType: this.getValueType(after),
@@ -236,6 +259,7 @@ export class RpcReceiveQueue {
     private batch: RpcObjectData[] = [];
 
     constructor(private readonly refs: Map<number, any>,
+                private readonly sourceFileType: string | undefined,
                 private readonly pull: () => Promise<RpcObjectData[]>,
                 private readonly logFile?: Writable) {
     }
@@ -301,7 +325,7 @@ export class RpcReceiveQueue {
                     let codec;
                     if (onChange) {
                         after = await onChange(before!);
-                    } else if ((codec = RpcCodecs.forInstance(before))) {
+                    } else if ((codec = RpcCodecs.forInstance(before, this.sourceFileType))) {
                         after = await codec.rpcReceive(before, this);
                     } else if (message.value !== undefined) {
                         after = message.valueType ? {kind: message.valueType, ...message.value} : message.value;

package/src/rpc/request/generate.ts

@@ -17,7 +17,11 @@ import * as rpc from "vscode-jsonrpc/node";
 import {Recipe, ScanningRecipe} from "../../recipe";
 import {Cursor, rootCursor} from "../../tree";
 import {ExecutionContext} from "../../execution";
-import {UUID} from "node:crypto";
+
+export interface GenerateResponse {
+    ids: string[]
+    sourceFileTypes: string[]
+}
 
 export class Generate {
     constructor(private readonly id: string, private readonly p: string) {
@@ -28,8 +32,13 @@ export class Generate {
                   preparedRecipes: Map<String, Recipe>,
                   recipeCursors: WeakMap<Recipe, Cursor>,
                   getObject: (id: string) => any): void {
-        connection.onRequest(new rpc.RequestType<Generate, UUID[], Error>("Generate"), async (request) => {
+        connection.onRequest(new rpc.RequestType<Generate, GenerateResponse, Error>("Generate"), async (request) => {
             const recipe = preparedRecipes.get(request.id);
+            const response = {
+                ids: [],
+                sourceFileTypes: []
+            } as GenerateResponse;
+
             if (recipe && recipe instanceof ScanningRecipe) {
                 let cursor = recipeCursors.get(recipe);
                 if (!cursor) {
@@ -39,12 +48,15 @@ export class Generate {
                 const ctx = getObject(request.p) as ExecutionContext;
                 const acc = recipe.accumulator(cursor, ctx);
                 const generated = await recipe.generate(acc, ctx)
-                return generated.map(g => {
+
+                for (const g of generated) {
                     localObjects.set(g.id.toString(), g);
-                    return g.id;
-                })
+                    response.ids.push(g.id.toString());
+                    response.sourceFileTypes.push(g.kind);
+                }
+
             }
-            return []
+            return response;
         });
     }
 }

package/src/rpc/request/get-object.ts

@@ -18,7 +18,8 @@ import {RpcObjectData, RpcObjectState, RpcSendQueue} from "../queue";
 import {ReferenceMap} from "../../reference";
 
 export class GetObject {
-    constructor(private readonly id: string, private readonly lastKnownId?: string) {
+    constructor(private readonly id: string,
+                private readonly sourceFileType?: string) {
     }
 
     static handle(
@@ -49,18 +50,9 @@ export class GetObject {
             let allData = pendingData.get(objId);
             if (!allData) {
                 const after = localObjects.get(objId);
-                
-                // Determine what the remote has cached
-                let before = undefined;
-                if (request.lastKnownId) {
-                    before = remoteObjects.get(request.lastKnownId);
-                    if (before === undefined) {
-                        // Remote had something cached, but we've evicted it - must send full object
-                        remoteObjects.delete(request.lastKnownId);
-                    }
-                }
+                const before = remoteObjects.get(objId);
 
-                allData = await new RpcSendQueue(localRefs, trace).generate(after, before);
+                allData = await new RpcSendQueue(localRefs, request.sourceFileType, trace).generate(after, before);
                 pendingData.set(objId, allData);
 
                 remoteObjects.set(objId, after);
@@ -75,4 +67,5 @@ export class GetObject {
 
             return batch;
         });
-    }}
+    }
+}

package/src/rpc/request/parse.ts

@@ -16,7 +16,7 @@
 import * as rpc from "vscode-jsonrpc/node";
 import {ExecutionContext} from "../../execution";
 import {UUID} from "node:crypto";
-import {Parser, ParserInput, Parsers} from "../../parser";
+import {ParserInput, Parsers} from "../../parser";
 import {randomId} from "../../uuid";
 import {produce} from "immer";
 import {SourceFile} from "../../tree";

package/src/rpc/request/print.ts

@@ -29,9 +29,9 @@ export class Print {
     }
 
     static handle(connection: rpc.MessageConnection,
-                  getObject: (id: string) => any): void {
+                  getObject: (id: string, sourceFileType?: string) => any): void {
         connection.onRequest(new rpc.RequestType<Print, string, Error>("Print"), async request => {
-            const tree: Tree = await getObject(request.treeId.toString());
+            const tree: Tree = await getObject(request.treeId.toString(), request.sourceFileType);
             const out = new PrintOutputCapture(PrintMarkerPrinter[request.markerPrinter]);
             if (isSourceFile(tree)) {
                 return await printer(tree).print(tree, out);

package/src/rpc/request/visit.ts

@@ -25,6 +25,7 @@ export interface VisitResponse {
 
 export class Visit {
     constructor(private readonly visitor: string,
+                private readonly sourceFileType: string,
                 private readonly visitorOptions: Map<string, any> | undefined,
                 private readonly treeId: string,
                 private readonly p: string,
@@ -35,15 +36,15 @@ export class Visit {
                   localObjects: Map<string, any>,
                   preparedRecipes: Map<String, Recipe>,
                   recipeCursors: WeakMap<Recipe, Cursor>,
-                  getObject: (id: string) => any,
-                  getCursor: (cursorIds: string[] | undefined) => Promise<Cursor>): void {
+                  getObject: (id: string, sourceFileType?: string) => any,
+                  getCursor: (cursorIds: string[] | undefined, sourceFileType?: string) => Promise<Cursor>): void {
         connection.onRequest(new rpc.RequestType<Visit, VisitResponse, Error>("Visit"), async (request) => {
-            const p = await getObject(request.p);
-            const before: Tree = await getObject(request.treeId);
+            const p = await getObject(request.p, undefined);
+            const before: Tree = await getObject(request.treeId, request.sourceFileType);
             localObjects.set(before.id.toString(), before);
 
             const visitor = await Visit.instantiateVisitor(request, preparedRecipes, recipeCursors, p);
-            const after = await visitor.visit(before, p, await getCursor(request.cursor));
+            const after = await visitor.visit(before, p, await getCursor(request.cursor, request.sourceFileType));
             if (!after) {
                 localObjects.delete(before.id.toString());
             } else if (after !== before) {

package/src/rpc/rewrite-rpc.ts

@@ -19,7 +19,7 @@ import {Cursor, isSourceFile, isTree, rootCursor, SourceFile, Tree} from "../tre
 import {Recipe, RecipeDescriptor, RecipeRegistry} from "../recipe";
 import {SnowflakeId} from "@akashrajpurohit/snowflake-id";
 import {
-    Generate,
+    Generate, GenerateResponse,
     GetObject,
     GetRecipes,
     Parse,
@@ -69,8 +69,8 @@ export class RewriteRpc {
         const recipeCursors: WeakMap<Recipe, Cursor> = new WeakMap()
 
         // Need this indirection, otherwise `this` will be undefined when executed in the handlers.
-        const getObject = (id: string) => this.getObject(id);
-        const getCursor = (cursorIds: string[] | undefined) => this.getCursor(cursorIds);
+        const getObject = (id: string, sourceFileType?: string) => this.getObject(id, sourceFileType);
+        const getCursor = (cursorIds: string[] | undefined, sourceFileType?: string) => this.getCursor(cursorIds, sourceFileType);
 
         const registry = options.registry || new RecipeRegistry();
 
@@ -104,14 +104,13 @@ export class RewriteRpc {
         return this;
     }
 
-    async getObject<P>(id: string): Promise<P> {
+    async getObject<P>(id: string, sourceFileType?: string): Promise<P> {
         const localObject = this.localObjects.get(id);
-        const lastKnownId = localObject ? id : undefined;
 
-        const q = new RpcReceiveQueue(this.remoteRefs, () => {
+        const q = new RpcReceiveQueue(this.remoteRefs, sourceFileType, () => {
             return this.connection.sendRequest(
                 new rpc.RequestType<GetObject, RpcObjectData[], Error>("GetObject"),
-                new GetObject(id, lastKnownId)
+                new GetObject(id, sourceFileType)
             );
         }, this.options.traceGetObjectInput);
 
@@ -128,11 +127,11 @@ export class RewriteRpc {
         return remoteObject;
     }
 
-    async getCursor(cursorIds: string[] | undefined): Promise<Cursor> {
+    async getCursor(cursorIds: string[] | undefined, sourceFileType?: string): Promise<Cursor> {
         let cursor = rootCursor();
         if (cursorIds) {
             for (let i = cursorIds.length - 1; i >= 0; i--) {
-                const cursorObject = await this.getObject(cursorIds[i]);
+                const cursorObject = await this.getObject(cursorIds[i], sourceFileType);
                 this.remoteObjects.set(cursorIds[i], cursorObject);
                 cursor = new Cursor(cursorObject, cursor);
             }
@@ -140,13 +139,13 @@ export class RewriteRpc {
         return cursor;
     }
 
-    async parse(inputs: ParserInput[], relativeTo?: string): Promise<SourceFile[]> {
+    async parse(inputs: ParserInput[], sourceFileType: string, relativeTo?: string): Promise<SourceFile[]> {
         const parsed: SourceFile[] = [];
         for (const g of await this.connection.sendRequest(
             new rpc.RequestType<Parse, string[], Error>("Parse"),
             new Parse(inputs, relativeTo)
         )) {
-            parsed.push(await this.getObject(g));
+            parsed.push(await this.getObject(g, sourceFileType));
         }
         return parsed;
     }
@@ -193,21 +192,26 @@ export class RewriteRpc {
         this.localObjects.set(tree.id.toString(), tree);
         const pId = this.localObject(p);
         const cursorIds = this.getCursorIds(cursor);
+
+        const sourceFileType = isSourceFile(tree) ? tree.kind :
+            cursor!.firstEnclosing(t => isSourceFile(t))!.kind;
+
         const response = await this.connection.sendRequest(
             new rpc.RequestType<Visit, VisitResponse, Error>("Visit"),
-            new Visit(visitorName, undefined, tree.id.toString(), pId, cursorIds)
+            new Visit(visitorName, sourceFileType, undefined, tree.id.toString(), pId, cursorIds)
         );
-        return response.modified ? this.getObject(tree.id.toString()) : tree;
+        return response.modified ? this.getObject(tree.id.toString(), sourceFileType) : tree;
     }
 
     async generate(remoteRecipeId: string, ctx: ExecutionContext): Promise<SourceFile[]> {
         const ctxId = this.localObject(ctx);
         const generated: SourceFile[] = [];
-        for (const g of await this.connection.sendRequest(
-            new rpc.RequestType<Generate, string[], Error>("Generate"),
+        const response = await this.connection.sendRequest(
+            new rpc.RequestType<Generate, GenerateResponse, Error>("Generate"),
             new Generate(remoteRecipeId, ctxId)
-        )) {
-            generated.push(await this.getObject(g));
+        );
+        for (let i = 0; i < response.ids.length; i++) {
+            generated.push(await this.getObject(response.ids[i], response.sourceFileTypes[i]));
         }
         return generated;
     }

package/src/test/rewrite-test.ts

@@ -212,7 +212,7 @@ export class RecipeSpec {
 }
 
 class ValidateWhitespaceVisitor extends JavaScriptVisitor<ExecutionContext> {
-    protected override async visitSpace(space: J.Space, p: ExecutionContext): Promise<J.Space> {
+    public override async visitSpace(space: J.Space, p: ExecutionContext): Promise<J.Space> {
         const ret = super.visitSpace(space, p);
         expect(space.whitespace).toMatch(/^\s*$/);
         return ret;

package/src/text/rpc.ts

@@ -13,10 +13,9 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-import {RpcCodec, RpcCodecs, RpcReceiveQueue, RpcSendQueue} from "../rpc";
+import {RpcCodecs, RpcReceiveQueue, RpcSendQueue} from "../rpc";
 import {PlainText} from "./tree";
 import {createDraft, Draft, finishDraft} from "immer";
-import {TreeKind} from "../tree";
 
 async function receiveSnippet(before: PlainText.Snippet, q: RpcReceiveQueue): Promise<PlainText.Snippet | undefined> {
     const draft: Draft<PlainText.Snippet> = createDraft(before);
@@ -26,40 +25,37 @@ async function receiveSnippet(before: PlainText.Snippet, q: RpcReceiveQueue): Pr
     return finishDraft(draft);
 }
 
-const textCodec: RpcCodec<PlainText> = {
-    async rpcReceive(before: PlainText, q: RpcReceiveQueue): Promise<PlainText> {
-        const draft: Draft<PlainText> = createDraft(before);
-        draft.id = await q.receive(before.id);
-        draft.markers = await q.receive(before.markers);
-        draft.sourcePath = await q.receive(before.sourcePath);
-        draft.charsetName = await q.receive(before.charsetName);
-        draft.charsetBomMarked = await q.receive(before.charsetBomMarked);
-        draft.checksum = await q.receive(before.checksum);
-        draft.fileAttributes = await q.receive(before.fileAttributes);
-        draft.text = await q.receive(before.text);
-        draft.snippets = (await q.receiveList(before.snippets, snippet => receiveSnippet(snippet, q)))!;
-        return finishDraft(draft);
-    },
+// Register codec for all Java AST node types
+for (const kind of Object.values(PlainText.Kind)) {
+    RpcCodecs.registerCodec(kind as string, {
+        async rpcReceive(before: PlainText, q: RpcReceiveQueue): Promise<PlainText> {
+            const draft: Draft<PlainText> = createDraft(before);
+            draft.id = await q.receive(before.id);
+            draft.markers = await q.receive(before.markers);
+            draft.sourcePath = await q.receive(before.sourcePath);
+            draft.charsetName = await q.receive(before.charsetName);
+            draft.charsetBomMarked = await q.receive(before.charsetBomMarked);
+            draft.checksum = await q.receive(before.checksum);
+            draft.fileAttributes = await q.receive(before.fileAttributes);
+            draft.text = await q.receive(before.text);
+            draft.snippets = (await q.receiveList(before.snippets, snippet => receiveSnippet(snippet, q)))!;
+            return finishDraft(draft);
+        },
 
-    async rpcSend(after: PlainText, q: RpcSendQueue): Promise<void> {
-        await q.getAndSend(after, p => p.id);
-        await q.getAndSend(after, p => p.markers);
-        await q.getAndSend(after, p => p.sourcePath);
-        await q.getAndSend(after, p => p.charsetName);
-        await q.getAndSend(after, p => p.charsetBomMarked);
-        await q.getAndSend(after, p => p.checksum);
-        await q.getAndSend(after, p => p.fileAttributes);
-        await q.getAndSend(after, p => p.text);
-        await q.getAndSendList(after, a => a.snippets, s => s.id, async (snippet) => {
-            await q.getAndSend(snippet, p => p.id);
-            await q.getAndSend(snippet, p => p.markers);
-            await q.getAndSend(snippet, p => p.text);
-        });
-    }
+        async rpcSend(after: PlainText, q: RpcSendQueue): Promise<void> {
+            await q.getAndSend(after, p => p.id);
+            await q.getAndSend(after, p => p.markers);
+            await q.getAndSend(after, p => p.sourcePath);
+            await q.getAndSend(after, p => p.charsetName);
+            await q.getAndSend(after, p => p.charsetBomMarked);
+            await q.getAndSend(after, p => p.checksum);
+            await q.getAndSend(after, p => p.fileAttributes);
+            await q.getAndSend(after, p => p.text);
+            await q.getAndSendList(after, a => a.snippets, s => s.id, async (snippet) => {
+                await q.getAndSend(snippet, p => p.id);
+                await q.getAndSend(snippet, p => p.markers);
+                await q.getAndSend(snippet, p => p.text);
+            });
+        }
+    }, PlainText.Kind.PlainText);
 }
-
-Object.values(PlainText.Kind).forEach(kind => {
-    if (!Object.values(TreeKind).includes(kind as any)) {
-        RpcCodecs.registerCodec(kind, textCodec);
-    }
-});