Keep interned RPC refs alive across source files

rewrite-core/src/main/java/org/openrewrite/rpc/RewriteRpc.java

@@ -258,7 +258,7 @@ protected Boolean handle(Void noParams) {
         jsonRpc.rpc("Evict", new JsonRpcMethod<Evict>() {
             @Override
             protected Boolean handle(Evict request) {
-                // Inbound side has no per-file checkpoint, so refs are left for Reset.
+                // Refs are left for Reset, so the next file reuses them.
                 remoteObjects.remove(request.getId());
                 localObjects.remove(request.getId());
                 return true;
@@ -367,37 +367,16 @@ public void reset() {
     }
 
     /**
-     * Ref high-water marks (send-side count, receive-side max key) captured before a file is
-     * visited so {@link #evict} rolls back exactly that file's refs. Receive-side uses the max
-     * key, not the size, because remote ids may be zero-based.
+     * Drop a file's tree from both peers. Interned refs deliberately survive so the next file
+     * reuses them rather than re-sending the objects they stand for; only {@link #reset} clears
+     * them. Notified fire-and-forget — under source-outer iteration the file's transfer is
+     * already complete.
      */
-    public int[] refCheckpoint() {
-        int remoteRefsMax = -1;
-        for (Integer ref : remoteRefs.keySet()) {
-            if (ref > remoteRefsMax) {
-                remoteRefsMax = ref;
-            }
-        }
-        return new int[]{localRefs.size(), remoteRefsMax};
-    }
-
-    /**
-     * Drop a file's tree from both peers and roll their refs back to the pre-file checkpoint.
-     * Symmetric by design: dropping the send-side ref forces the next file to re-{@code ADD} the
-     * interned object instead of a {@code REF_USE} the rolled-back receiver would reject. Notified
-     * fire-and-forget — under source-outer iteration the file's transfer is already complete.
-     *
-     * @param localRefsCheckpoint  {@code refCheckpoint()[0]} captured before the file was visited
-     * @param remoteRefsCheckpoint {@code refCheckpoint()[1]} captured before the file was visited
-     */
-    public void evict(String id, int localRefsCheckpoint, int remoteRefsCheckpoint) {
+    public void evict(String id) {
         jsonRpc.notify(new JsonRpcRequest(null, "Evict", RawJson.of(new Evict(id))));
 
         remoteObjects.remove(id);
         localObjects.remove(id);
-
-        localRefs.values().removeIf(ref -> ref > localRefsCheckpoint);
-        remoteRefs.keySet().removeIf(ref -> ref > remoteRefsCheckpoint);
     }
 
     public <P> @Nullable Tree visit(SourceFile sourceFile, String visitorName, P p) {

rewrite-core/src/main/java/org/openrewrite/scheduling/RecipeRunCycle.java

@@ -125,7 +125,6 @@ public LSS scanSources(LSS sourceSet) {
             return sourceSetEditor.apply(sourceSet, sourceFile -> {
                 BatchState scanBatch = new BatchState();
                 Set<RewriteRpc> touched = newSetFromMap(new IdentityHashMap<>());
-                Map<RewriteRpc, int[]> refCheckpoints = new IdentityHashMap<>();
 
                 SourceFile result = allRecipeStack.reduce(sourceSet, recipe, ctx, (source, recipeStack) -> {
                     Recipe recipe = leaf(recipeStack);
@@ -140,11 +139,9 @@ public LSS scanSources(LSS sourceSet) {
                         RewriteRpc currentRpc = recipe instanceof RpcRecipe ? ((RpcRecipe) recipe).getRpc() : null;
                         String scanVisitorName = recipe instanceof RpcRecipe ? ((RpcRecipe) recipe).getScanVisitor() : null;
 
-                        if (scanVisitorName != null) {
-                            captureRpc(currentRpc, touched, refCheckpoints);
-                        }
-
                         if (currentRpc != null && scanVisitorName != null) {
+                            touched.add(currentRpc);
+
                             // Flush if switching to a different RPC instance
                             if (scanBatch.rpc != null && scanBatch.rpc != currentRpc) {
                                 flushScanBatch(scanBatch, source);
@@ -204,7 +201,7 @@ public LSS scanSources(LSS sourceSet) {
                     flushScanBatch(scanBatch, result);
                 }
 
-                evictSourceFile(sourceFile, touched, refCheckpoints);
+                evictSourceFile(sourceFile, touched);
                 return result;
             });
         }
@@ -230,29 +227,16 @@ private void flushScanBatch(BatchState batch, SourceFile source) {
     }
 
     /**
-     * Record a peer this file touched, snapshotting its ref high-water on first sight so
-     * {@link #evictSourceFile} can roll back exactly the refs this file introduced.
+     * Drop this source file's tree from every RPC peer that visited it, bounding each peer's tree
+     * cache to ~one file at a time. Interned refs survive so the next file reuses them.
      */
-    private static void captureRpc(@Nullable RewriteRpc rpc, Set<RewriteRpc> touched,
-                                   Map<RewriteRpc, int[]> refCheckpoints) {
-        if (rpc != null && touched.add(rpc)) {
-            refCheckpoints.put(rpc, rpc.refCheckpoint());
-        }
-    }
-
-    /**
-     * Drop this source file's tree from every RPC peer that visited it, rolling each peer's
-     * ref maps back to the pre-file checkpoint. Bounds RPC-server memory to ~one file at a time.
-     */
-    private static void evictSourceFile(@Nullable SourceFile sourceFile, Set<RewriteRpc> touched,
-                                        Map<RewriteRpc, int[]> refCheckpoints) {
+    private static void evictSourceFile(@Nullable SourceFile sourceFile, Set<RewriteRpc> touched) {
         if (sourceFile == null || touched.isEmpty()) {
             return;
         }
         String id = sourceFile.getId().toString();
         for (RewriteRpc rpc : touched) {
-            int[] cp = refCheckpoints.get(rpc);
-            rpc.evict(id, cp[0], cp[1]);
+            rpc.evict(id);
         }
     }
 
@@ -355,7 +339,6 @@ void clear() {
         recipeRunStats.recordSourceVisited(sourceFile);
         BatchState batch = new BatchState();
         Set<RewriteRpc> touched = newSetFromMap(new IdentityHashMap<>());
-        Map<RewriteRpc, int[]> refCheckpoints = new IdentityHashMap<>();
 
         SourceFile result = allRecipeStack.reduce(sourceSet, recipe, ctx, (source, recipeStack) -> {
             Recipe recipe = leaf(recipeStack);
@@ -364,7 +347,9 @@ void clear() {
             }
 
             RewriteRpc currentRpc = recipe instanceof RpcRecipe ? ((RpcRecipe) recipe).getRpc() : null;
-            captureRpc(currentRpc, touched, refCheckpoints);
+            if (currentRpc != null) {
+                touched.add(currentRpc);
+            }
 
             // Flush batch if switching to a different RPC or non-RPC recipe
             if (batch.rpc != null && batch.rpc != currentRpc) {
@@ -488,7 +473,7 @@ void clear() {
         }
 
         // Recipe errors are handled inside the reduce, so this runs on every normal return.
-        evictSourceFile(sourceFile, touched, refCheckpoints);
+        evictSourceFile(sourceFile, touched);
         return result;
     }
 

rewrite-core/src/test/java/org/openrewrite/rpc/RewriteRpcTest.java

@@ -45,8 +45,10 @@
 import java.io.PipedInputStream;
 import java.io.PipedOutputStream;
 import java.nio.file.Path;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.CountDownLatch;
 
@@ -199,8 +201,8 @@ void sendFailureCleansUpRemoteObjects() {
     }
 
     /**
-     * {@link RewriteRpc#evict} drops the tree from both peers and rolls the client's ref maps
-     * back to the pre-file checkpoint.
+     * {@link RewriteRpc#evict} drops the tree from both peers but leaves every interned ref in
+     * place, so the next source file reuses them instead of re-sending what they stand for.
      */
     @SneakyThrows
     @Test
@@ -212,9 +214,6 @@ void evictDropsTreeFromBothPeers() {
         String id = original.getId().toString();
         String sourceFileType = PlainText.class.getName();
 
-        // High-water before the client fetches anything, so evict rolls back exactly this exchange.
-        int[] checkpoint = client.refCheckpoint();
-
         // Server holds the tree; client fetches it → both peers cache it.
         server.localObjects.put(id, original);
         client.getObject(id, sourceFileType);
@@ -223,12 +222,14 @@ void evictDropsTreeFromBothPeers() {
         assertThat(server.localObjects).containsKey(id);
         assertThat(server.remoteObjects).containsKey(id);
 
-        client.evict(id, checkpoint[0], checkpoint[1]);
+        Set<Integer> refsBefore = new HashSet<>(client.remoteRefs.keySet());
+
+        client.evict(id);
 
-        // Client cleared synchronously, including refs rolled back to the checkpoint.
+        // Trees cleared synchronously; refs survive, which is what makes cross-file interning work.
         assertThat(client.localObjects).doesNotContainKey(id);
         assertThat(client.remoteObjects).doesNotContainKey(id);
-        assertThat(client.remoteRefs.keySet()).allMatch(ref -> ref <= checkpoint[1]);
+        assertThat(client.remoteRefs.keySet()).containsExactlyInAnyOrderElementsOf(refsBefore);
 
         // The Evict notification is fire-and-forget; wait for the server to apply it.
         long deadline = System.currentTimeMillis() + 5_000;

rewrite-csharp/csharp/OpenRewrite/CSharp/Rpc/RewriteRpcServer.cs

@@ -86,12 +86,6 @@ public class RewriteRpcServer
     /// </summary>
     private readonly ConcurrentDictionary<int, object> _remoteRefs = new();
 
-    /// <summary>
-    /// Ref high-water per source file (send-side _localRefs count, receive-side max _remoteRefs
-    /// key), captured before first visit so <see cref="Evict"/> rolls back exactly its refs.
-    /// </summary>
-    private readonly ConcurrentDictionary<string, (int LocalRefs, int RemoteRefsMax)> _refCheckpoints = new();
-
     /// <summary>
     /// DependencyTypes pages its (potentially hundreds-of-MB) response: the full RpcObjectData list
     /// is built once, cached keyed by coordinate, and handed back one
@@ -1448,7 +1442,6 @@ public async Task<VisitResponse> Visit(VisitRequest request)
         }
 
         // Fetch tree from the remote (Java) process
-        CaptureRefCheckpoint(request.TreeId);
         var tree = await GetObjectFromRemoteAsync(request.TreeId, request.SourceFileType);
 
         if (phase != "scan" && phase != "edit")
@@ -1503,7 +1496,6 @@ public async Task<BatchVisitResponse> BatchVisit(BatchVisitRequest request)
         }
 
         var sw = Stopwatch.StartNew();
-        CaptureRefCheckpoint(request.TreeId);
         var tree = await GetObjectFromRemoteAsync(request.TreeId, request.SourceFileType);
         var fetchMs = sw.ElapsedMilliseconds;
 
@@ -1806,35 +1798,15 @@ private void ClearLocalState()
         _remoteObjects.Clear();
         _localRefs.Clear();
         _remoteRefs.Clear();
-        _refCheckpoints.Clear();
         _preparedRecipes.Clear();
         _recipeAccumulators.Clear();
         _executionContexts.Clear();
     }
 
     /// <summary>
-    /// Records the ref high-water before a source file is first visited (first visit wins), so
-    /// <see cref="Evict"/> can roll back exactly the refs that file introduced.
-    /// </summary>
-    private void CaptureRefCheckpoint(string treeId)
-    {
-        _refCheckpoints.GetOrAdd(treeId, _ =>
-        {
-            var remoteMax = -1;
-            foreach (var key in _remoteRefs.Keys)
-            {
-                if (key > remoteMax)
-                {
-                    remoteMax = key;
-                }
-            }
-            return (_localRefs.Count, remoteMax);
-        });
-    }
-
-    /// <summary>
-    /// Drops one source file's tree and rolls back the refs it introduced; recipe/accumulator/
-    /// context state (keyed separately) is preserved. Fire-and-forget, so it returns no response.
+    /// Drops one source file's tree. Interned refs are deliberately kept so the next file reuses
+    /// them; only Reset clears them. Recipe/accumulator/context state (keyed separately) is
+    /// preserved. Fire-and-forget, so it returns no response.
     /// </summary>
     [JsonRpcMethod("Evict", UseSingleObjectParameterDeserialization = true)]
     public void Evict(EvictRequest request)
@@ -1845,23 +1817,6 @@ public void Evict(EvictRequest request)
         }
         _localObjects.TryRemove(request.Id, out _);
         _remoteObjects.TryRemove(request.Id, out _);
-        if (_refCheckpoints.TryRemove(request.Id, out var cp))
-        {
-            foreach (var kv in _localRefs)
-            {
-                if (kv.Value > cp.LocalRefs)
-                {
-                    _localRefs.TryRemove(kv.Key, out _);
-                }
-            }
-            foreach (var key in _remoteRefs.Keys)
-            {
-                if (key > cp.RemoteRefsMax)
-                {
-                    _remoteRefs.TryRemove(key, out _);
-                }
-            }
-        }
     }
 
     /// <summary>

rewrite-csharp/csharp/OpenRewrite/Core/Rpc/RpcSendQueue.cs

@@ -365,6 +365,7 @@ private void DoChange(object? after, object? before, Action? onChange, IRpcCodec
             {
                 "JRightPadded" or "JLeftPadded" or "JContainer" or "JavaType" =>
                     $"org.openrewrite.java.tree.{name}",
+                "JavaSourceSet" => "org.openrewrite.java.marker.JavaSourceSet",
                 _ => $"org.openrewrite.java.tree.J${name}",
             },
             "OpenRewrite.CSharp" => $"org.openrewrite.csharp.tree.Cs${name}",

rewrite-csharp/csharp/OpenRewrite/Java/JavaSourceSet.cs

@@ -0,0 +1,85 @@
+/*
+ * Copyright 2026 the original author or authors.
+ * <p>
+ * Licensed under the Moderne Source Available License (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * <p>
+ * https://docs.moderne.io/licensing/moderne-source-available-license
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+using OpenRewrite.Core;
+using OpenRewrite.Core.Rpc;
+using OpenRewrite.Java.Rpc;
+
+namespace OpenRewrite.Java;
+
+/// <summary>
+/// Mirrors org.openrewrite.java.marker.JavaSourceSet. It rides on every source file in a Java
+/// source set, resource files included, so it reaches this peer through Xml as well as C#.
+/// <para>
+/// Field order is the protocol; see JavaSourceSet#rpcSend on the Java side. Without this codec the
+/// marker resolves to <see cref="UnknownMarker"/>, which consumes one message where Java sends
+/// many, and the queue then desynchronizes with no diagnostic.
+/// </para>
+/// </summary>
+public sealed class JavaSourceSet(
+    Guid id,
+    string name,
+    IList<JavaType.FullyQualified> classpath,
+    IDictionary<string, IList<JavaType.FullyQualified>> gavToTypes
+) : Marker, IRpcCodec<JavaSourceSet>
+{
+    public Guid Id { get; } = id;
+    public string Name { get; } = name;
+    public IList<JavaType.FullyQualified> Classpath { get; } = classpath;
+    public IDictionary<string, IList<JavaType.FullyQualified>> GavToTypes { get; } = gavToTypes;
+
+    public void RpcSend(JavaSourceSet after, RpcSendQueue q)
+    {
+        var typeSender = new JavaSender();
+        q.GetAndSend(after, s => s.Id);
+        q.GetAndSend(after, s => s.Name);
+        q.GetAndSendListAsRef(after, s => s.Classpath, TypeKey, t => typeSender.VisitType(t, q));
+
+        var gavs = after.GavToTypes.Keys.ToList();
+        q.GetAndSendList(after, _ => gavs, gav => gav, null);
+        foreach (var gav in gavs)
+        {
+            q.GetAndSendListAsRef(after, s => s.GavToTypes[gav], TypeKey,
+                t => typeSender.VisitType(t, q));
+        }
+    }
+
+    // Classpath entries are Class instances in practice; the key only feeds the sender's own diff.
+    private static object TypeKey(JavaType.FullyQualified type) =>
+        type is JavaType.Class cls ? cls.FullyQualifiedName : type;
+
+    public JavaSourceSet RpcReceive(JavaSourceSet before, RpcReceiveQueue q)
+    {
+        var typeReceiver = new JavaReceiver();
+        var id = q.ReceiveAndGet<Guid, string>(before.Id, Guid.Parse);
+        var name = q.Receive(before.Name)!;
+        var classpath = q.ReceiveList(before.Classpath,
+            t => (JavaType.FullyQualified)typeReceiver.VisitType(t, q)!)!;
+
+        // The uninitialized instance the queue hands back on an ADD has null collections, so the
+        // before state is never dereferenced without a guard.
+        var beforeGavs = before.GavToTypes;
+        var gavs = q.ReceiveList(beforeGavs?.Keys.ToList(), null);
+        var gavToTypes = new Dictionary<string, IList<JavaType.FullyQualified>>();
+        foreach (var gav in gavs ?? [])
+        {
+            IList<JavaType.FullyQualified>? beforeBucket = null;
+            beforeGavs?.TryGetValue(gav, out beforeBucket);
+            gavToTypes[gav] = q.ReceiveList(beforeBucket,
+                t => (JavaType.FullyQualified)typeReceiver.VisitType(t, q)!)!;
+        }
+        return new JavaSourceSet(id, name, classpath, gavToTypes);
+    }
+}