From b8b9502db4a2611c6faf6e7a51b6021ad02cd09e Mon Sep 17 00:00:00 2001
From: Tianyu Geng <tgeng@google.com>
Date: Tue, 7 Sep 2021 09:10:15 -0700
Subject: [PATCH] FIR checker: fix
 ASSIGNING_SINGLE_ELEMENT_TO_VARARG_IN_NAMED_FORM_ANNOTATION

Currently this checker is missing `arrayOf` calls inside annotation.
---
 .../expression/FirNamedVarargChecker.kt       | 24 ++++++++++++-------
 ...tAssigningSingleElementsInNamedForm.fir.kt |  4 ++++
 ...hibitAssigningSingleElementsInNamedForm.kt |  4 ++++
 ...ibitAssigningSingleElementsInNamedForm.txt |  2 ++
 4 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirNamedVarargChecker.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirNamedVarargChecker.kt
index e9c93aee499..a6998927ef6 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirNamedVarargChecker.kt
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirNamedVarargChecker.kt
@@ -22,7 +22,10 @@ import org.jetbrains.kotlin.fir.types.isArrayType
 
 object FirNamedVarargChecker : FirCallChecker() {
     override fun check(expression: FirCall, context: CheckerContext, reporter: DiagnosticReporter) {
-        if (expression !is FirFunctionCall && expression !is FirAnnotation && expression !is FirDelegatedConstructorCall) return
+        if (expression !is FirFunctionCall &&
+            expression !is FirAnnotation &&
+            expression !is FirDelegatedConstructorCall &&
+            expression !is FirArrayOfCall) return
         val isAnnotation = expression is FirAnnotation
         val errorFactory =
             if (isAnnotation) FirErrors.ASSIGNING_SINGLE_ELEMENT_TO_VARARG_IN_NAMED_FORM_ANNOTATION
@@ -44,13 +47,18 @@ object FirNamedVarargChecker : FirCallChecker() {
             reporter.reportOn(argument.expression.source, errorFactory, context)
         }
 
-        val argumentMap = expression.argumentMapping ?: return
-        for ((argument, parameter) in argumentMap) {
-            if (!parameter.isVararg) continue
-            if (argument is FirVarargArgumentsExpression) {
-                argument.arguments.forEach(::checkArgument)
-            } else {
-                checkArgument(argument)
+        if (expression is FirArrayOfCall) {
+            // FirArrayOfCall has the `vararg` argument expression pre-flattened and doesn't have an argument mapping.
+            expression.arguments.forEach(::checkArgument)
+        } else {
+            val argumentMap = expression.argumentMapping ?: return
+            for ((argument, parameter) in argumentMap) {
+                if (!parameter.isVararg) continue
+                if (argument is FirVarargArgumentsExpression) {
+                    argument.arguments.forEach(::checkArgument)
+                } else {
+                    checkArgument(argument)
+                }
             }
         }
     }
diff --git a/compiler/testData/diagnostics/tests/varargs/prohibitAssigningSingleElementsInNamedForm.fir.kt b/compiler/testData/diagnostics/tests/varargs/prohibitAssigningSingleElementsInNamedForm.fir.kt
index 4aea436cf72..8b03955bb68 100644
--- a/compiler/testData/diagnostics/tests/varargs/prohibitAssigningSingleElementsInNamedForm.fir.kt
+++ b/compiler/testData/diagnostics/tests/varargs/prohibitAssigningSingleElementsInNamedForm.fir.kt
@@ -12,6 +12,10 @@ fun f1() {}
 @Anno2(i = intArrayOf(1))
 fun f2() {}
 
+@Anno1(s = arrayOf(elements = <!ASSIGNING_SINGLE_ELEMENT_TO_VARARG_IN_NAMED_FORM_FUNCTION_ERROR!>"foo"<!>))
+@Anno2(i = intArrayOf(elements = *intArrayOf(1)))
+fun f3() {}
+
 fun foo(vararg ints: Int) {}
 
 fun test() {
diff --git a/compiler/testData/diagnostics/tests/varargs/prohibitAssigningSingleElementsInNamedForm.kt b/compiler/testData/diagnostics/tests/varargs/prohibitAssigningSingleElementsInNamedForm.kt
index 4eb8af851f5..8d44de46f76 100644
--- a/compiler/testData/diagnostics/tests/varargs/prohibitAssigningSingleElementsInNamedForm.kt
+++ b/compiler/testData/diagnostics/tests/varargs/prohibitAssigningSingleElementsInNamedForm.kt
@@ -12,6 +12,10 @@ fun f1() {}
 @Anno2(i = intArrayOf(1))
 fun f2() {}
 
+@Anno1(s = arrayOf(elements = <!ASSIGNING_SINGLE_ELEMENT_TO_VARARG_IN_NAMED_FORM_FUNCTION_ERROR, TYPE_MISMATCH!>"foo"<!>))
+@Anno2(i = intArrayOf(elements = *<!REDUNDANT_SPREAD_OPERATOR_IN_NAMED_FORM_IN_FUNCTION!>intArrayOf(1)<!>))
+fun f3() {}
+
 fun foo(vararg ints: Int) {}
 
 fun test() {
diff --git a/compiler/testData/diagnostics/tests/varargs/prohibitAssigningSingleElementsInNamedForm.txt b/compiler/testData/diagnostics/tests/varargs/prohibitAssigningSingleElementsInNamedForm.txt
index 25f1681a889..eabd7615ddc 100644
--- a/compiler/testData/diagnostics/tests/varargs/prohibitAssigningSingleElementsInNamedForm.txt
+++ b/compiler/testData/diagnostics/tests/varargs/prohibitAssigningSingleElementsInNamedForm.txt
@@ -2,6 +2,7 @@ package
 
 @Anno1(s = {"foo"}) @Anno2(i = {1}) public fun f1(): kotlin.Unit
 @Anno1(s = {"foo"}) @Anno2(i = {1}) public fun f2(): kotlin.Unit
+@Anno1(s = {"foo"}) @Anno2(i = {{1}}) public fun f3(): kotlin.Unit
 public fun foo(/*0*/ vararg ints: kotlin.Int /*kotlin.IntArray*/): kotlin.Unit
 public fun test(): kotlin.Unit
 
@@ -20,3 +21,4 @@ public final annotation class Anno2 : kotlin.Annotation {
     public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
     public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
 }
+