Correctly set Cursor in Repeat (#4015)

The `Cursor` wasn't being set in `Repeat` causing issues when
used for non `SourceFile` trees.

Signed-off-by: Jonathan Leitschuh <[email protected]>

rewrite-core/src/main/java/org/openrewrite/Cursor.java

@@ -26,6 +26,9 @@
 
 import static java.util.stream.StreamSupport.stream;
 
+/**
+ * A cursor is linked path of LST elements that can be used to traverse down the tree towards the root.
+ */
 @EqualsAndHashCode(exclude = "messages")
 public class Cursor {
     public static final String ROOT_VALUE = "root";
@@ -51,6 +54,13 @@ public Cursor getRoot() {
         return c;
     }
 
+    /**
+     * @return true if this cursor is the root of the tree, false otherwise
+     */
+    final boolean isRoot() {
+        return ROOT_VALUE.equals(value);
+    }
+
     public Iterator<Cursor> getPathAsCursors() {
         return new CursorPathIterator(this);
     }

rewrite-core/src/main/java/org/openrewrite/Repeat.java

@@ -43,6 +43,16 @@ public static TreeVisitor<?, ExecutionContext> repeatUntilStable(TreeVisitor<?,
                     return tree;
                 }
 
+                if (tree != null && !(tree instanceof SourceFile) && getCursor().isRoot()) {
+                    throw new IllegalArgumentException(
+                            String.format(
+                                    "Repeat visitor called on a non-source file tree without a cursor pointing to the root of the tree. " +
+                                    "Passed tree type: `%s`. " +
+                                    "This is likely a bug in the calling code. Use a `visit` method that accepts a cursor instead.",
+                                    tree.getClass().getName()
+                            ));
+                }
+
                 Tree previous = tree;
                 Tree current = null;
                 for (int i = 0; i < maxCycles; i++) {
@@ -55,6 +65,25 @@ public static TreeVisitor<?, ExecutionContext> repeatUntilStable(TreeVisitor<?,
 
                 return current;
             }
+
+            @Override
+            public @Nullable Tree visit(@Nullable Tree tree, ExecutionContext ctx, Cursor parent) {
+                if (tree instanceof SourceFile && !v.isAcceptable((SourceFile) tree, ctx)) {
+                    return tree;
+                }
+
+                Tree previous = tree;
+                Tree current = null;
+                for (int i = 0; i < maxCycles; i++) {
+                    current = v.visit(previous, ctx, parent);
+                    if (current == previous) {
+                        break;
+                    }
+                    previous = current;
+                }
+
+                return current;
+            }
         };
     }
 }

rewrite-java-test/src/test/java/org/openrewrite/RepeatTest.java

@@ -0,0 +1,109 @@
+/*
+ * Copyright 2024 the original author or authors.
+ * <p>
+ * Licensed under the Apache License, Version 2.0 (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://www.apache.org/licenses/LICENSE-2.0
+ * <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.
+ */
+package org.openrewrite;
+
+import org.junit.jupiter.api.Test;
+import org.openrewrite.internal.RecipeRunException;
+import org.openrewrite.internal.lang.Nullable;
+import org.openrewrite.java.JavaVisitor;
+import org.openrewrite.java.tree.J;
+import org.openrewrite.java.tree.JavaSourceFile;
+import org.openrewrite.test.RewriteTest;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.api.Assertions.*;
+import static org.openrewrite.java.Assertions.java;
+
+public class RepeatTest implements RewriteTest {
+
+    /**
+     * This visitor verifies that the cursor is well-formed.
+     */
+    static class VerifyCursorWellFormed<P> extends JavaVisitor<P> {
+        @Override
+        public J preVisit(J tree, P p) {
+            assertNotNull(getCursor().firstEnclosing(JavaSourceFile.class), "JavaSourceFile should be accessible");
+            assertNotEquals(getCursor().getParentOrThrow().getValue(), tree, "Tree should not be the same as its parent");
+            return tree;
+        }
+    }
+
+    static class VerifyCursorWellFormedInRepeat extends JavaVisitor<ExecutionContext> {
+        @Override
+        public J preVisit(J tree, ExecutionContext executionContext) {
+            return (J) Repeat.repeatUntilStable(new VerifyCursorWellFormed<>()).visitNonNull(tree, executionContext, getCursor().getParentTreeCursor());
+        }
+    }
+
+    @Test
+    void repeatInPreVisit() {
+        rewriteRun(
+          spec -> spec.recipe(RewriteTest.toRecipe(VerifyCursorWellFormedInRepeat::new)),
+          java("class A {}")
+        );
+    }
+
+    public static class VerifyCursorWellFormedRecipe extends Recipe {
+        @Override
+        public String getDisplayName() {
+            return "Verify cursor well-formed";
+        }
+
+        @Override
+        public String getDescription() {
+            return "This recipe verifies that the cursor is well-formed.";
+        }
+
+        @Override
+        public TreeVisitor<?, ExecutionContext> getVisitor() {
+            return Repeat.repeatUntilStable(new VerifyCursorWellFormed<>());
+        }
+    }
+
+    @Test
+    void repeatInRecipe() {
+        rewriteRun(
+          spec -> spec.recipe(new VerifyCursorWellFormedRecipe()),
+          java("class A {}")
+        );
+    }
+
+    static class VisitorThatFailsToSetCursor extends JavaVisitor<ExecutionContext> {
+        @Override
+        public @Nullable J preVisit(J tree, ExecutionContext executionContext) {
+            return (J) Repeat.repeatUntilStable(new JavaVisitor<>()).visitNonNull(tree, executionContext);
+        }
+    }
+
+    @Test
+    void repeatValidatesCursorIsPassed() {
+        AssertionError assertionError = assertThrows(AssertionError.class, () -> {
+            rewriteRun(
+              spec -> spec.recipe(RewriteTest.toRecipe(VisitorThatFailsToSetCursor::new)),
+              java("class A {}")
+            );
+        });
+        assertThat(assertionError).cause().isInstanceOf(RecipeRunException.class);
+        RecipeRunException e = (RecipeRunException) assertionError.getCause();
+        assertThat(e.getMessage())
+          .contains(
+            "Repeat visitor called on a non-source file tree without a cursor pointing to the root of the tree. " +
+            "Passed tree type: `org.openrewrite.java.tree.J$ClassDeclaration`. " +
+            "This is likely a bug in the calling code. " +
+            "Use a `visit` method that accepts a cursor instead."
+          );
+    }
+}