Fix argument list parsing in null-safe calls

Fixes: #308
openrewrite · Sep 11, 2023 · 65dfd26 · 65dfd26
1 parent 1cc57bd
commit 65dfd26
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/src/main/kotlin/org/openrewrite/kotlin/internal/KotlinParserVisitor.kt b/src/main/kotlin/org/openrewrite/kotlin/internal/KotlinParserVisitor.kt
@@ -1394,7 +1394,7 @@ class KotlinParserVisitor(
 
     private fun mapFunctionalCallArguments(firCall: FirCall, skipFirstArgument: Boolean = false): JContainer<Expression> {
         var callPsi = getPsiElement(firCall)!!
-        callPsi = if (callPsi is KtDotQualifiedExpression) callPsi.lastChild else callPsi
+        callPsi = if (callPsi is KtDotQualifiedExpression || callPsi is KtSafeQualifiedExpression) callPsi.lastChild else callPsi
         val firArguments = if (skipFirstArgument) firCall.argumentList.arguments.subList(1, firCall.argumentList.arguments.size) else firCall.argumentList.arguments
         val flattenedExpressions = firArguments.stream()
             .map { e -> if (e is FirVarargArgumentsExpression) e.arguments else listOf(e) }

diff --git a/src/test/java/org/openrewrite/kotlin/tree/MethodInvocationTest.java b/src/test/java/org/openrewrite/kotlin/tree/MethodInvocationTest.java
@@ -676,4 +676,19 @@ void spaceAfterLambdaParameter() {
           )
         );
     }
+
+    @Test
+    @Issue("https://github.com/openrewrite/rewrite-kotlin/issues/308")
+    void trailingLambdaAfterNullSafe() {
+        rewriteRun(
+          kotlin(
+            """
+              val x = "x"
+                   ?.associateTo(mutableMapOf()) { p ->
+                       p to listOfNotNull(p.uppercase())
+                   }
+              """
+          )
+        );
+    }
 }