From 6969a5350d40afd8b7bbb25fad10731061636e4e Mon Sep 17 00:00:00 2001
From: Nicolay Mitropolsky <nicolay.mitropolsky@jetbrains.com>
Date: Tue, 13 Feb 2018 18:14:13 +0300
Subject: [PATCH] Uast: `guessCorrespondingParameter` support for varargs
 (IDEA-CR-29257)

---
 .../TestDataFileReferenceTest.kt              | 21 +++++++++++++++++++
 .../src/org/jetbrains/uast/UastUtils.kt       |  7 ++++++-
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/plugins/devkit/testSources/testAssistant/TestDataFileReferenceTest.kt b/plugins/devkit/testSources/testAssistant/TestDataFileReferenceTest.kt
index a3e5c6a0d683..b1b62d133a52 100644
--- a/plugins/devkit/testSources/testAssistant/TestDataFileReferenceTest.kt
+++ b/plugins/devkit/testSources/testAssistant/TestDataFileReferenceTest.kt
@@ -29,6 +29,27 @@ class TestDataFileReferenceTest : TestDataReferenceTestCase() {
     assertResolvedTo(javaFile, "TestClass.java")
   }
 
+  fun testDataFileResolveVarags() {
+    val testClass1 = myContentRootSubdir.writeChild("TestClass1.java", "some java code here")
+    val testClass2 = myContentRootSubdir.writeChild("TestClass2.java", "some java code here")
+
+    myFixture.configureByText("ATest.java", """
+      import com.intellij.testFramework.fixtures.LightCodeInsightFixtureTestCase;
+
+      @com.intellij.testFramework.TestDataPath("${"\$"}CONTENT_ROOT/contentRootSubdir/")
+      public class ATest extends LightCodeInsightFixtureTestCase {
+        protected void doTest() {
+          configureByFiles("TestClass1.java", "TestClass2.java");
+        }
+
+        void configureByFiles(@com.intellij.testFramework.TestDataFile String... files){}
+      }
+    """.trimIndent())
+
+    assertResolvedTo(testClass1, "TestClass1.java")
+    assertResolvedTo(testClass2, "TestClass2.java")
+  }
+
 }
 
 
diff --git a/uast/uast-common/src/org/jetbrains/uast/UastUtils.kt b/uast/uast-common/src/org/jetbrains/uast/UastUtils.kt
index 106426d3ed48..cb7ddeb5e182 100644
--- a/uast/uast-common/src/org/jetbrains/uast/UastUtils.kt
+++ b/uast/uast-common/src/org/jetbrains/uast/UastUtils.kt
@@ -184,7 +184,12 @@ fun guessCorrespondingParameter(callExpression: UCallExpression, arg: UExpressio
   val parameters = psiMethod.parameterList.parameters
 
   if (callExpression is UCallExpressionEx)
-    return parameters.withIndex().find { (i, _) -> callExpression.getArgumentForParameter(i) == arg }?.value
+    return parameters.withIndex().find { (i, p) ->
+      val argumentForParameter = callExpression.getArgumentForParameter(i) ?: return@find false
+      if (argumentForParameter == arg) return@find true
+      if (p.isVarArgs && argumentForParameter is UExpressionList) return@find argumentForParameter.expressions.contains(arg)
+      return@find false
+    }?.value
 
   // not everyone implements UCallExpressionEx, lets try to guess
   val indexInArguments = callExpression.valueArguments.indexOf(arg)