RewriteRpc no longer uses identity hash codes

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

@@ -18,6 +18,8 @@
 import io.moderne.jsonrpc.JsonRpc;
 import io.moderne.jsonrpc.JsonRpcMethod;
 import io.moderne.jsonrpc.JsonRpcRequest;
+import io.moderne.jsonrpc.internal.SnowflakeId;
+import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.VisibleForTesting;
 import org.jspecify.annotations.Nullable;
 import org.openrewrite.*;
@@ -50,6 +52,8 @@ public class RewriteRpc {
 
     @VisibleForTesting
     final Map<String, Object> localObjects = new HashMap<>();
+    /* A reverse map of the objects back to their IDs */
+    final Map<Object, String> localObjectIds = new IdentityHashMap<>();
 
     private final Map<Integer, Object> remoteRefs = new IdentityHashMap<>();
 
@@ -117,26 +121,18 @@ public <P> VisitResponse scan(SourceFile sourceFile, String visitorName, P p,
         // Set the local state of this tree, so that when the remote
         // asks for it, we know what to send.
         localObjects.put(sourceFile.getId().toString(), sourceFile);
-        String pId = Integer.toString(System.identityHashCode(p));
-
-        Object p2 = p;
-        while (p2 instanceof DelegatingExecutionContext) {
-            p2 = ((DelegatingExecutionContext) p2).getDelegate();
-        }
-        if (p2 instanceof ExecutionContext) {
-            ((ExecutionContext) p2).putMessage("org.openrewrite.rpc.id", pId);
-        }
-        localObjects.put(pId, p2);
 
+        String pId = maybeUnwrapExecutionContext(p);
         List<String> cursorIds = getCursorIds(cursor);
 
         return send("Visit", new Visit(visitorName, null, sourceFile.getId().toString(), pId, cursorIds),
                 VisitResponse.class);
     }
 
     public Collection<? extends SourceFile> generate(String remoteRecipeId, ExecutionContext ctx) {
-        List<String> generated = send("Generate", new Generate(remoteRecipeId,
-                Integer.toString(System.identityHashCode(ctx))), GenerateResponse.class);
+        String ctxId = maybeUnwrapExecutionContext(ctx);
+        List<String> generated = send("Generate", new Generate(remoteRecipeId, ctxId),
+                GenerateResponse.class);
         if (!generated.isEmpty()) {
             return generated.stream()
                     .map(this::<SourceFile>getObject)
@@ -145,6 +141,27 @@ public Collection<? extends SourceFile> generate(String remoteRecipeId, Executio
         return emptyList();
     }
 
+    /**
+     * If the p object is a DelegatingExecutionContext, unwrap it to get the underlying
+     * ExecutionContext
+     *
+     * @param p   A visitor parameter, which may or may not be an ExecutionContext
+     * @param <P> The type of p
+     * @return The ID of p as represented in the local object cache.
+     */
+    private <P> String maybeUnwrapExecutionContext(P p) {
+        Object p2 = p;
+        while (p2 instanceof DelegatingExecutionContext) {
+            p2 = ((DelegatingExecutionContext) p2).getDelegate();
+        }
+        String pId = localObjectIds.computeIfAbsent(p2, p3 -> SnowflakeId.generateId());
+        if (p2 instanceof ExecutionContext) {
+            ((ExecutionContext) p2).putMessage("org.openrewrite.rpc.id", pId);
+        }
+        localObjects.put(pId, p2);
+        return pId;
+    }
+
     public List<RecipeDescriptor> getRecipes() {
         return send("GetRecipes", null, GetRecipesResponse.class);
     }
@@ -171,7 +188,7 @@ List<String> getCursorIds(@Nullable Cursor cursor) {
             cursorIds = cursor.getPathAsStream().map(c -> {
                 String id = c instanceof Tree ?
                         ((Tree) c).getId().toString()
-                        : Integer.toString(System.identityHashCode(c));
+                        : localObjectIds.computeIfAbsent(c, c2 -> SnowflakeId.generateId());
                 localObjects.put(id, c);
                 return id;
             }).collect(Collectors.toList());
@@ -195,6 +212,7 @@ public <T> T getObject(String id) {
         }
         // We are now in sync with the remote state of the object.
         remoteObjects.put(id, remoteObject);
+        localObjects.put(id, remoteObject);
 
         //noinspection unchecked
         return (T) remoteObject;

rewrite-core/src/main/java/org/openrewrite/rpc/request/Generate.java

@@ -19,13 +19,19 @@
 import lombok.RequiredArgsConstructor;
 import lombok.Value;
 import org.openrewrite.*;
+import org.openrewrite.scheduling.RecipeRunCycle;
+import org.openrewrite.scheduling.WatchableExecutionContext;
+import org.openrewrite.table.RecipeRunStats;
+import org.openrewrite.table.SourcesFileErrors;
+import org.openrewrite.table.SourcesFileResults;
 
 import java.util.Collection;
 import java.util.Map;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import static java.util.Collections.emptyList;
+import static org.openrewrite.ExecutionContext.CURRENT_RECIPE;
 
 @Value
 public class Generate implements RpcRequest {
@@ -45,8 +51,17 @@ public static class Handler extends JsonRpcMethod<Generate> {
 
         @Override
         protected Object handle(Generate request) throws Exception {
-            ExecutionContext ctx = (ExecutionContext) getObject.apply(request.getP());
             Recipe recipe = preparedRecipes.get(request.getId());
+
+            ExecutionContext ctx = (ExecutionContext) getObject.apply(request.getP());
+            if (ctx.getMessage(CURRENT_RECIPE) == null) {
+                WatchableExecutionContext wctx = new WatchableExecutionContext((ExecutionContext) ctx);
+                wctx.putCycle(new RecipeRunCycle<>(recipe, 0, new Cursor(null, Cursor.ROOT_VALUE), wctx,
+                        new RecipeRunStats(Recipe.noop()), new SourcesFileResults(Recipe.noop()),
+                        new SourcesFileErrors(Recipe.noop()), LargeSourceSet::edit));
+                ctx.putCurrentRecipe(recipe);
+            }
+
             if (recipe instanceof ScanningRecipe) {
                 //noinspection unchecked
                 ScanningRecipe<Object> scanningRecipe = (ScanningRecipe<Object>) recipe;

rewrite-core/src/main/java/org/openrewrite/rpc/request/Visit.java

@@ -127,12 +127,12 @@ private Object getVisitorP(Visit request) {
                 String visitorName = request.getVisitor();
 
                 if (visitorName.startsWith("scan:") || visitorName.startsWith("edit:")) {
-                    WatchableExecutionContext ctx = new WatchableExecutionContext((ExecutionContext) p);
                     Recipe recipe = preparedRecipes.get(visitorName.substring(
                             "edit:".length() /* 'scan:' has same length*/));
                     // This is really probably particular to the Java implementation,
                     // because we are carrying forward the legacy of cycles that are likely to be
                     // removed from OpenRewrite in the future.
+                    WatchableExecutionContext ctx = new WatchableExecutionContext((ExecutionContext) p);
                     ctx.putCycle(new RecipeRunCycle<>(recipe, 0, new Cursor(null, Cursor.ROOT_VALUE), ctx,
                             new RecipeRunStats(Recipe.noop()), new SourcesFileResults(Recipe.noop()),
                             new SourcesFileErrors(Recipe.noop()), LargeSourceSet::edit));