Remove the need for cursor messaging

openrewrite · Dec 13, 2024 · 9d2ed73 · 9d2ed73
1 parent 40a44b1
commit 9d2ed73
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 38 deletions.
diff --git a/src/main/java/org/openrewrite/java/migrate/lombok/LombokUtils.java b/src/main/java/org/openrewrite/java/migrate/lombok/LombokUtils.java
@@ -32,6 +32,9 @@
 class LombokUtils {
 
     static boolean isGetter(J.MethodDeclaration method) {
+        if (method.getMethodType() == null) {
+            return false;
+        }
         // Check signature: no parameters
         if (!(method.getParameters().get(0) instanceof J.Empty) || method.getReturnTypeExpression() == null) {
             return false;
@@ -43,24 +46,20 @@ static boolean isGetter(J.MethodDeclaration method) {
             return false;
         }
         // Check field is declared on method type
-        JavaType.Method methodType = method.getMethodType();
-        if (methodType == null) {
-            return false;
-        }
-        JavaType.FullyQualified declaringType = methodType.getDeclaringType();
-
-        // Check return: type and matching field name
+        JavaType.FullyQualified declaringType = method.getMethodType().getDeclaringType();
         Expression returnExpression = ((J.Return) method.getBody().getStatements().get(0)).getExpression();
         if (returnExpression instanceof J.Identifier) {
             J.Identifier identifier = (J.Identifier) returnExpression;
             if (identifier.getFieldType() != null && declaringType == identifier.getFieldType().getOwner()) {
+                // Check return: type and matching field name
                 return hasMatchingTypeAndName(method, identifier.getType(), identifier.getSimpleName());
             }
         } else if (returnExpression instanceof J.FieldAccess) {
             J.FieldAccess fieldAccess = (J.FieldAccess) returnExpression;
             Expression target = fieldAccess.getTarget();
             if (target instanceof J.Identifier && ((J.Identifier) target).getFieldType() != null &&
                     declaringType == ((J.Identifier) target).getFieldType().getOwner()) {
+                // Check return: type and matching field name
                 return hasMatchingTypeAndName(method, fieldAccess.getType(), fieldAccess.getSimpleName());
             }
         }

diff --git a/src/main/java/org/openrewrite/java/migrate/lombok/UseLombokGetter.java b/src/main/java/org/openrewrite/java/migrate/lombok/UseLombokGetter.java
@@ -68,44 +68,39 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
     @Value
     @EqualsAndHashCode(callSuper = false)
     private static class MethodRemover extends JavaIsoVisitor<ExecutionContext> {
-        private static final String FIELDS_TO_DECORATE_KEY = "FIELDS_TO_DECORATE";
-
         @Override
         public J.ClassDeclaration visitClassDeclaration(J.ClassDeclaration classDecl, ExecutionContext ctx) {
-
-            //initialize set of fields to annotate
-            getCursor().putMessage(FIELDS_TO_DECORATE_KEY, new HashSet<Finding>());
-
-            //delete methods, note down corresponding fields
-            J.ClassDeclaration classDeclAfterVisit = super.visitClassDeclaration(classDecl, ctx);
+            Set<Finding> fieldsToDecorate = new HashSet<>();
+
+            J.Block oldBody = classDecl.getBody();
+            J.Block newBody = (J.Block) new JavaIsoVisitor<ExecutionContext>() {
+                //delete methods, note down corresponding fields
+                @Override
+                public J.@Nullable MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, ExecutionContext ctx) {
+                    if (LombokUtils.isGetter(method)) {
+                        Expression returnExpression = ((J.Return) method.getBody().getStatements().get(0)).getExpression();
+                        if (returnExpression instanceof J.Identifier) {
+                            fieldsToDecorate.add(new Finding(
+                                    ((J.Identifier) returnExpression).getSimpleName(),
+                                    LombokUtils.getAccessLevel(method.getModifiers())));
+                            return null;
+                        } else if (returnExpression instanceof J.FieldAccess) {
+                            fieldsToDecorate.add(new Finding(
+                                    ((J.FieldAccess) returnExpression).getSimpleName(),
+                                    LombokUtils.getAccessLevel(method.getModifiers())));
+                            return null;
+                        }
+                    }
+                    return method;
+                }
+            }.visitNonNull(oldBody, ctx);
 
             //only thing that can have changed is removal of getter methods
-            if (classDeclAfterVisit != classDecl) {
+            if (oldBody != newBody) {
                 //this set collects the fields for which existing methods have already been removed
-                Set<Finding> fieldsToDecorate = getCursor().pollNearestMessage(FIELDS_TO_DECORATE_KEY);
                 doAfterVisit(new FieldAnnotator(fieldsToDecorate));
             }
-            return classDeclAfterVisit;
-        }
-
-        @Override
-        public J.@Nullable MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, ExecutionContext ctx) {
-            if (method.getMethodType() != null && LombokUtils.isGetter(method)) {
-                Set<Finding> set = getCursor().getNearestMessage(FIELDS_TO_DECORATE_KEY);
-                Expression returnExpression = ((J.Return) method.getBody().getStatements().get(0)).getExpression();
-                if (returnExpression instanceof J.Identifier) {
-                    set.add(new Finding(
-                            ((J.Identifier) returnExpression).getSimpleName(),
-                            LombokUtils.getAccessLevel(method.getModifiers())));
-                    return null;
-                } else if (returnExpression instanceof J.FieldAccess) {
-                    set.add(new Finding(
-                            ((J.FieldAccess) returnExpression).getSimpleName(),
-                            LombokUtils.getAccessLevel(method.getModifiers())));
-                    return null;
-                }
-            }
-            return method;
+            return super.visitClassDeclaration(classDecl.withBody(newBody), ctx);
         }
     }