Add withModifier to base java visitors

openrewrite · Oct 9, 2024 · 8d38278 · 8d38278
1 parent 7522567
commit 8d38278
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 8 deletions.
diff --git a/rewrite-java/src/main/java/org/openrewrite/java/JavaIsoVisitor.java b/rewrite-java/src/main/java/org/openrewrite/java/JavaIsoVisitor.java
@@ -229,6 +229,11 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, P p)
         return (J.MethodInvocation) super.visitMethodInvocation(method, p);
     }
 
+    @Override
+    public J.Modifier visitModifier(J.Modifier modifier, P p) {
+        return (J.Modifier) super.visitModifier(modifier, p);
+    }
+
     @Override
     public J.MultiCatch visitMultiCatch(J.MultiCatch multiCatch, P p) {
         return (J.MultiCatch) super.visitMultiCatch(multiCatch, p);

diff --git a/rewrite-java/src/main/java/org/openrewrite/java/JavaPrinter.java b/rewrite-java/src/main/java/org/openrewrite/java/JavaPrinter.java
@@ -92,7 +92,8 @@ protected void visitRightPadded(@Nullable JRightPadded<? extends J> rightPadded,
         }
     }
 
-    protected void visitModifier(Modifier mod, PrintOutputCapture<P> p) {
+    @Override
+    public J visitModifier(Modifier mod, PrintOutputCapture<P> p) {
         visit(mod.getAnnotations(), p);
         String keyword = "";
         switch (mod.getType()) {
@@ -154,6 +155,7 @@ protected void visitModifier(Modifier mod, PrintOutputCapture<P> p) {
         beforeSyntax(mod, Space.Location.MODIFIER_PREFIX, p);
         p.append(keyword);
         afterSyntax(mod, p);
+        return mod;
     }
 
     @Override

diff --git a/rewrite-java/src/main/java/org/openrewrite/java/JavaVisitor.java b/rewrite-java/src/main/java/org/openrewrite/java/JavaVisitor.java
@@ -449,8 +449,6 @@ public J visitClassDeclaration(J.ClassDeclaration classDecl, P p) {
             c = (J.ClassDeclaration) temp;
         }
         c = c.withLeadingAnnotations(ListUtils.map(c.getLeadingAnnotations(), a -> visitAndCast(a, p)));
-        c = c.withModifiers(ListUtils.map(c.getModifiers(),
-                mod -> mod.withPrefix(visitSpace(mod.getPrefix(), Space.Location.MODIFIER_PREFIX, p))));
         c = c.withModifiers(ListUtils.map(c.getModifiers(), m -> visitAndCast(m, p)));
         //Kind can have annotations associated with it, need to visit those.
         c = c.getPadding().withKind(
@@ -846,8 +844,6 @@ public J visitMethodDeclaration(J.MethodDeclaration method, P p) {
         }
         m = m.withLeadingAnnotations(ListUtils.map(m.getLeadingAnnotations(), a -> visitAndCast(a, p)));
         m = m.withModifiers(ListUtils.map(m.getModifiers(), e -> visitAndCast(e, p)));
-        m = m.withModifiers(ListUtils.map(m.getModifiers(),
-                mod -> mod.withPrefix(visitSpace(mod.getPrefix(), Space.Location.MODIFIER_PREFIX, p))));
         J.TypeParameters typeParameters = m.getAnnotations().getTypeParameters();
         if (typeParameters != null) {
             m = m.getAnnotations().withTypeParameters(typeParameters.withAnnotations(
@@ -920,6 +916,14 @@ public J visitMethodInvocation(J.MethodInvocation method, P p) {
         return m;
     }
 
+    public J visitModifier(J.Modifier modifer, P p) {
+        J.Modifier m = modifer;
+        m = m.withPrefix(visitSpace(m.getPrefix(), Space.Location.MODIFIER_PREFIX, p));
+        m = m.withMarkers(visitMarkers(m.getMarkers(), p));
+        m = m.withAnnotations(ListUtils.map(m.getAnnotations(), a -> visitAndCast(a, p)));
+        return m;
+    }
+
     public J visitMultiCatch(J.MultiCatch multiCatch, P p) {
         J.MultiCatch m = multiCatch;
         m = m.withPrefix(visitSpace(m.getPrefix(), Space.Location.MULTI_CATCH_PREFIX, p));
@@ -941,8 +945,6 @@ public J visitVariableDeclarations(J.VariableDeclarations multiVariable, P p) {
         }
         m = m.withLeadingAnnotations(ListUtils.map(m.getLeadingAnnotations(), a -> visitAndCast(a, p)));
         m = m.withModifiers(Objects.requireNonNull(ListUtils.map(m.getModifiers(), e -> visitAndCast(e, p))));
-        m = m.withModifiers(ListUtils.map(m.getModifiers(),
-                mod -> mod.withPrefix(visitSpace(mod.getPrefix(), Space.Location.MODIFIER_PREFIX, p))));
         m = m.withTypeExpression(visitAndCast(m.getTypeExpression(), p));
         m = m.withTypeExpression(m.getTypeExpression() == null ?
                 null :

diff --git a/rewrite-java/src/main/java/org/openrewrite/java/UpdateSourcePositions.java b/rewrite-java/src/main/java/org/openrewrite/java/UpdateSourcePositions.java
@@ -66,14 +66,15 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
             }
 
             @Override
-            protected void visitModifier(J.Modifier modifier, PrintOutputCapture<ExecutionContext> p) {
+            public J visitModifier(J.Modifier modifier, PrintOutputCapture<ExecutionContext> p) {
                 PositionPrintOutputCapture prefix = new PositionPrintOutputCapture(ppoc.pos, ppoc.line, ppoc.column);
                 spacePrinter.visitSpace(modifier.getPrefix(), Space.Location.ANY, prefix);
 
                 Range.Position startPosition = new Range.Position(prefix.pos, prefix.line, prefix.column);
                 super.visitModifier(modifier, p);
                 Range.Position endPosition = new Range.Position(ppoc.pos, ppoc.line, ppoc.column);
                 positionMap.put(modifier, new Range(randomId(), startPosition, endPosition));
+                return modifier;
             }
 
         };