Set the correct MethodType on names of method invocation and declarat…

…ions.
openrewrite · Dec 7, 2023 · 6a77148 · 6a77148
1 parent c24a822
commit 6a77148
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 17 deletions.
diff --git a/src/main/java/org/openrewrite/kotlin/internal/KotlinTreeParserVisitor.java b/src/main/java/org/openrewrite/kotlin/internal/KotlinTreeParserVisitor.java
@@ -194,7 +194,7 @@ public J visitArrayAccessExpression(KtArrayAccessExpression expression, Executio
         }
 
         JContainer<Expression> args = JContainer.build(Space.EMPTY, expressions, markers);
-        return new J.MethodInvocation(
+        return mapType(new J.MethodInvocation(
                 randomId(),
                 prefix(expression),
                 markers,
@@ -203,7 +203,7 @@ public J visitArrayAccessExpression(KtArrayAccessExpression expression, Executio
                 name,
                 args,
                 type
-        );
+        ));
     }
 
     @Override
@@ -2370,7 +2370,7 @@ public J visitDotQualifiedExpression(KtDotQualifiedExpression expression, Execut
                         padLeft(suffix(expression.getReceiverExpression()), (J.Identifier) pt.getClazz()),
                         pt.getType()
                 ));
-                return pt.withClazz(newName).withPrefix(prefix);
+                return mapType(pt.withClazz(newName).withPrefix(prefix));
             } else if (j instanceof J.MethodInvocation) {
                 J.MethodInvocation m = (J.MethodInvocation) j;
                 return m.getPadding().withSelect(padRight(receiver, suffix(expression.getReceiverExpression())))
@@ -3307,7 +3307,7 @@ private J.MethodInvocation mapFunctionCall(KtBinaryExpression expression, Execut
         JContainer<Expression> args = JContainer.build(Space.EMPTY, expressions, paramMarkers);
         JavaType.Method methodType = methodInvocationType(expression);
 
-        return new J.MethodInvocation(
+        return mapType(new J.MethodInvocation(
                 randomId(),
                 prefix(expression),
                 markers,
@@ -3316,7 +3316,7 @@ private J.MethodInvocation mapFunctionCall(KtBinaryExpression expression, Execut
                 name,
                 args,
                 methodType
-        );
+        ));
     }
 
     private J.ControlParentheses<Expression> buildIfCondition(KtIfExpression expression) {
@@ -3482,19 +3482,11 @@ private J.FieldAccess mapType(J.FieldAccess tree) {
     }
 
     private J.MethodDeclaration mapType(J.MethodDeclaration tree) {
-        J.MethodDeclaration m = tree;
-        if (!(m.getName().getType() instanceof JavaType.Method)) {
-            m = m.withName(m.getName().withType(m.getMethodType()));
-        }
-        return m;
+        return tree.withName(tree.getName().withType(tree.getMethodType()));
     }
 
     private J.MethodInvocation mapType(J.MethodInvocation tree) {
-        J.MethodInvocation m = tree;
-        if (!(m.getName().getType() instanceof JavaType.Method)) {
-            m = m.withName(m.getName().withType(m.getMethodType()));
-        }
-        return m;
+        return tree.withName(tree.getName().withType(tree.getMethodType()));
     }
 
     private J.NewClass mapType(J.NewClass tree) {

diff --git a/src/test/java/org/openrewrite/kotlin/KotlinTypeMappingTest.java b/src/test/java/org/openrewrite/kotlin/KotlinTypeMappingTest.java
@@ -450,11 +450,11 @@ void implicitInvoke() {
                   """, spec -> spec.afterRecipe(cu -> {
                     AtomicBoolean found = new AtomicBoolean(false);
                     new KotlinIsoVisitor<AtomicBoolean>() {
-                        final MethodMatcher matcher = new MethodMatcher("kotlin.Function1 invoke(..)");
+                        final MethodMatcher matcher = new MethodMatcher("kotlin.Function1 block(..)");
                         @Override
                         public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, AtomicBoolean atomicBoolean) {
                             if (matcher.matches(method)) {
-                                assertThat(method.getMethodType().toString()).isEqualTo("kotlin.Function1<kotlin.collections.Collection<kotlin.Any>, kotlin.Unit>{name=invoke,return=kotlin.Unit,parameters=[kotlin.collections.Collection<kotlin.Any>]}");
+                                assertThat(method.getMethodType().toString()).isEqualTo("kotlin.Function1<kotlin.collections.Collection<kotlin.Any>, kotlin.Unit>{name=block,return=kotlin.Unit,parameters=[kotlin.collections.Collection<kotlin.Any>]}");
                                 found.set(true);
                             }
                             return super.visitMethodInvocation(method, atomicBoolean);
@@ -1275,5 +1275,28 @@ public J.Annotation visitAnnotation(J.Annotation annotation, Integer integer) {
               )
             );
         }
+
+        @Issue("https://github.com/openrewrite/rewrite-kotlin/issues/506")
+        @Test
+        void methodTypeOnMethodInvocation() {
+            //noinspection RemoveRedundantBackticks
+            rewriteRun(
+              kotlin(
+                "val arr = listOf(1, 2, 3)",
+                spec -> spec.afterRecipe(cu -> {
+                    AtomicBoolean found = new AtomicBoolean(false);
+                    new KotlinIsoVisitor<Integer>() {
+                        @Override
+                        public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, Integer integer) {
+                            assertThat(method.getName().getType()).isEqualTo(method.getMethodType());
+                            found.set(true);
+                            return super.visitMethodInvocation(method, integer);
+                        }
+                    }.visit(cu, 0);
+                    assertThat(found.get()).isTrue();
+                })
+              )
+            );
+        }
     }
 }