Fix trailing comment of K.FunctionType

openrewrite · Nov 10, 2023 · 346ad78 · 346ad78
1 parent f4d68b6
commit 346ad78
Show file tree

Hide file tree

Showing 7 changed files with 31 additions and 19 deletions.
diff --git a/src/main/java/org/openrewrite/kotlin/KotlinVisitor.java b/src/main/java/org/openrewrite/kotlin/KotlinVisitor.java
@@ -159,7 +159,7 @@ public J visitFunctionType(K.FunctionType functionType, P p) {
         if (f.getPadding().getParameters() != null) {
             f = f.getPadding().withParameters(this.visitContainer(f.getPadding().getParameters(), KContainer.Location.FUNCTION_TYPE_PARAMETERS, p));
         }
-        f = f.withReturnType(visitAndCast(f.getReturnType(), p));
+        f = f.withReturnType(visitRightPadded(f.getReturnType(), p));
         return f;
     }
 

diff --git a/src/main/java/org/openrewrite/kotlin/format/SpacesVisitor.java b/src/main/java/org/openrewrite/kotlin/format/SpacesVisitor.java
@@ -951,7 +951,9 @@ public K.FunctionType visitFunctionType(K.FunctionType functionType, P p) {
         if (kf.getArrow() != null) {
             kf = kf.withArrow(updateSpace(kf.getArrow(), style.getOther().getAroundArrowInFunctionTypes()));
         }
-        kf = kf.withReturnType(spaceBefore(kf.getReturnType(), style.getOther().getAroundArrowInFunctionTypes()));
+
+        JRightPadded<TypedTree> rpTypedTree = kf.getReturnType();
+        kf = kf.withReturnType(rpTypedTree.withElement(spaceBefore(rpTypedTree.getElement(), style.getOther().getAroundArrowInFunctionTypes())));
         return kf;
     }
 

diff --git a/src/main/java/org/openrewrite/kotlin/internal/KotlinPrinter.java b/src/main/java/org/openrewrite/kotlin/internal/KotlinPrinter.java
@@ -234,7 +234,8 @@ public J visitFunctionType(K.FunctionType functionType, PrintOutputCapture<P> p)
         delegate.visitContainer("(", functionType.getPadding().getParameters(), JContainer.Location.TYPE_PARAMETERS, ",", ")", p);
         visitSpace(functionType.getArrow() != null ? functionType.getArrow() : Space.SINGLE_SPACE, KSpace.Location.FUNCTION_TYPE_ARROW_PREFIX, p);
         p.append("->");
-        visit(functionType.getReturnType(), p);
+
+        visitRightPadded(functionType.getReturnType(), p);
         if (nullable) {
             p.append(")");
         }

diff --git a/src/main/java/org/openrewrite/kotlin/internal/KotlinTreeParserVisitor.java b/src/main/java/org/openrewrite/kotlin/internal/KotlinTreeParserVisitor.java
@@ -618,13 +618,13 @@ public J visitFunctionType(KtFunctionType type, ExecutionContext data) {
         return new K.FunctionType(
                 randomId(),
                 prefix(type, consumedSpaces),
-                Markers.EMPTY, //.addIfAbsent(new IsNullable(randomId(), Space.EMPTY)), // TODO
+                Markers.EMPTY,
                 emptyList(), // TODO
                 emptyList(), // TODO
                 type.getReceiver() != null ? padRight((NameTree) requireNonNull(type.getReceiverTypeReference()).accept(this, data), suffix(type.getReceiver())) : null,
                 parameters,
                 suffix(type.getParameterList()),
-                requireNonNull(type.getReturnTypeReference()).accept(this, data).withPrefix(prefix(type.getReturnTypeReference()))
+                padRight((TypedTree) requireNonNull(type.getReturnTypeReference()).accept(this, data), suffix(type))
         );
     }
 
@@ -777,9 +777,9 @@ public J visitNamedDeclaration(KtNamedDeclaration declaration, ExecutionContext
 
     @Override
     public J visitNullableType(KtNullableType nullableType, ExecutionContext data) {
-        J j = requireNonNull(nullableType.getInnerType()).accept(this, data);
-        return j.withPrefix(deepPrefix(nullableType))
-                .withMarkers(j.getMarkers().addIfAbsent(new IsNullable(randomId(),
+        TypeTree typeTree = (TypeTree) requireNonNull(nullableType.getInnerType()).accept(this, data);
+        return typeTree.withPrefix(deepPrefix(nullableType))
+                .withMarkers(typeTree.getMarkers().addIfAbsent(new IsNullable(randomId(),
                         prefix((PsiElement) nullableType.getQuestionMarkNode())
                 )));
     }

diff --git a/src/main/java/org/openrewrite/kotlin/tree/K.java b/src/main/java/org/openrewrite/kotlin/tree/K.java
@@ -1023,7 +1023,7 @@ class FunctionType implements K, TypeTree, Expression {
 
         public FunctionType(UUID id, Space prefix, Markers markers, List<Annotation> leadingAnnotations,
                             List<Modifier> modifiers, @Nullable JRightPadded<NameTree> receiver,
-                            @Nullable JContainer<TypeTree> parameters, @Nullable Space arrow, TypedTree returnType) {
+                            @Nullable JContainer<TypeTree> parameters, @Nullable Space arrow, JRightPadded<TypedTree> returnType) {
             this.id = id;
             this.prefix = prefix;
             this.markers = markers;
@@ -1036,9 +1036,8 @@ public FunctionType(UUID id, Space prefix, Markers markers, List<Annotation> lea
         }
 
         public Space getPrefix() {
-            // For backwards compatibility with older LST before there was a prefix field
             //noinspection ConstantConditions
-            return prefix == null ? returnType.getPrefix() : prefix;
+            return prefix == null ? returnType.getElement().getPrefix() : prefix;
         }
 
         @With
@@ -1089,21 +1088,19 @@ public FunctionType withParameters(List<TypeTree> parameters) {
         @Getter
         Space arrow;
 
-        // backwards compatibility
-        @JsonAlias("typedTree")
         @With
         @Getter
-        TypedTree returnType;
+        JRightPadded<TypedTree> returnType;
 
         @Override
         public @Nullable JavaType getType() {
-            return returnType.getType();
+            return returnType.getElement().getType();
         }
 
         public <T extends J> T withType(@Nullable JavaType type) {
-            TypeTree newType = returnType.withType(type);
+            TypeTree newType = returnType.getElement().withType(type);
             //noinspection unchecked
-            return (T) (newType == type ? this : withReturnType(newType));
+            return (T) (newType == type ? this : withReturnType(returnType.withElement(newType)));
         }
 
         @Override

diff --git a/src/main/kotlin/org/openrewrite/kotlin/internal/KotlinParserVisitor.kt b/src/main/kotlin/org/openrewrite/kotlin/internal/KotlinParserVisitor.kt
@@ -1883,7 +1883,7 @@ class KotlinParserVisitor(
                 receiver,
                 JContainer.build(before, refParams as List<JRightPadded<TypeTree>>, Markers.EMPTY),
                 arrow,
-                returnType
+                padRight(returnType, Space.EMPTY)
         )
     }
 

diff --git a/src/test/java/org/openrewrite/kotlin/tree/FunctionTypeTest.java b/src/test/java/org/openrewrite/kotlin/tree/FunctionTypeTest.java
@@ -88,7 +88,19 @@ void parenthesizedNullableType() {
         rewriteRun(
           kotlin(
             """
-              val v: (  (   Any )  -> Any) /*c1*/  ? = null
+              val v: (  (   Int )  -> Any) /*c1*/  ? = null
+              """
+          )
+        );
+    }
+
+    @Test
+    @Issue("https://github.com/openrewrite/rewrite-kotlin/issues/365")
+    void parenthesizedNullableTypeWithTrailingComment() {
+        rewriteRun(
+          kotlin(
+            """
+              val v: (  (   Int )  -> Any /*C*/)  ? = null
               """
           )
         );