FIR: add mapping to partially resolved arg list

Initially I tried adding `mapping` field to `FirArgumentList` but it
seems to be very difficult to make it flexible enough to do what I want.

So instead, I am creating a `FirPartialResolvedArgumentList`, which
seems to be very simple.

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/generators/CallAndReferenceGenerator.kt

@@ -465,7 +465,7 @@ class CallAndReferenceGenerator(
                             ((calleeReference as? FirResolvedNamedReference)?.resolvedSymbol as? FirFunctionSymbol<*>)?.fir
                         }
                         val valueParameters = function?.valueParameters
-                        val argumentMapping = call.argumentMapping
+                        val argumentMapping = call.resolvedArgumentMapping
                         val substitutor = (call as? FirFunctionCall)?.buildSubstitutorByCalledFunction(function) ?: ConeSubstitutor.Empty
                         if (argumentMapping != null && (annotationMode || argumentMapping.isNotEmpty())) {
                             if (valueParameters != null) {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt

@@ -16,7 +16,6 @@ import org.jetbrains.kotlin.fir.references.builder.buildResolvedCallableReferenc
 import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.calls.*
-import org.jetbrains.kotlin.fir.resolve.calls.CallableReferenceAdaptation
 import org.jetbrains.kotlin.fir.resolve.dfa.FirDataFlowAnalyzer
 import org.jetbrains.kotlin.fir.resolve.inference.*
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
@@ -161,7 +160,11 @@ class FirCallCompletionResultsWriterTransformer(
                 resultType = typeRef.substituteTypeRef(subCandidate)
                 val expectedArgumentsTypeMapping = runIf(!calleeReference.isError) { subCandidate.createArgumentsMapping() }
                 result.argumentList.transformArguments(this, expectedArgumentsTypeMapping)
-                if (!calleeReference.isError) {
+                if (calleeReference.isError) {
+                    subCandidate.argumentMapping?.let {
+                        result.replaceArgumentList(buildPartiallyResolvedArgumentList(result.argumentList, it))
+                    }
+                } else {
                     subCandidate.handleVarargs()
                     subCandidate.argumentMapping?.let {
                         result.replaceArgumentList(buildResolvedArgumentList(it))
@@ -212,7 +215,11 @@ class FirCallCompletionResultsWriterTransformer(
                 }
             }
         }
-        if (!calleeReference.isError) {
+        if (calleeReference.isError) {
+            subCandidate.argumentMapping?.let {
+                annotationCall.replaceArgumentList(buildPartiallyResolvedArgumentList(annotationCall.argumentList, it))
+            }
+        } else {
             subCandidate.handleVarargs()
             subCandidate.argumentMapping?.let {
                 annotationCall.replaceArgumentList(buildResolvedArgumentList(it))
@@ -367,7 +374,11 @@ class FirCallCompletionResultsWriterTransformer(
 
         val argumentsMapping = runIf(!calleeReference.isError) { calleeReference.candidate.createArgumentsMapping() }
         delegatedConstructorCall.argumentList.transformArguments(this, argumentsMapping)
-        if (!calleeReference.isError) {
+        if (calleeReference.isError) {
+            subCandidate.argumentMapping?.let {
+                delegatedConstructorCall.replaceArgumentList(buildPartiallyResolvedArgumentList(delegatedConstructorCall.argumentList, it))
+            }
+        } else {
             subCandidate.handleVarargs()
             subCandidate.argumentMapping?.let {
                 delegatedConstructorCall.replaceArgumentList(buildResolvedArgumentList(it))

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirArgumentUtil.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.declarations.FirValueParameter
 import org.jetbrains.kotlin.fir.expressions.builder.buildArgumentList
 import org.jetbrains.kotlin.fir.expressions.impl.FirArraySetArgumentList
+import org.jetbrains.kotlin.fir.expressions.impl.FirPartiallyResolvedArgumentList
 import org.jetbrains.kotlin.fir.expressions.impl.FirResolvedArgumentList
 
 fun buildUnaryArgumentList(argument: FirExpression): FirArgumentList = buildArgumentList {
@@ -26,6 +27,16 @@ fun buildArraySetArgumentList(rValue: FirExpression, indexes: List<FirExpression
 fun buildResolvedArgumentList(mapping: LinkedHashMap<FirExpression, FirValueParameter>): FirArgumentList =
     FirResolvedArgumentList(mapping)
 
+fun buildPartiallyResolvedArgumentList(
+    original: FirArgumentList,
+    mapping: LinkedHashMap<FirExpression, FirValueParameter>
+): FirArgumentList {
+    return FirPartiallyResolvedArgumentList(
+        original.source,
+        original.arguments.map { key -> key to mapping[key] }.toMap(LinkedHashMap())
+    )
+}
+
 object FirEmptyArgumentList : FirAbstractArgumentList() {
     override val arguments: List<FirExpression>
         get() = emptyList()

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirExpressionUtil.kt

@@ -11,8 +11,10 @@ import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
 import org.jetbrains.kotlin.fir.expressions.builder.buildConstExpression
 import org.jetbrains.kotlin.fir.expressions.builder.buildErrorExpression
 import org.jetbrains.kotlin.fir.expressions.builder.buildErrorLoop
-import org.jetbrains.kotlin.fir.expressions.impl.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirBlockImpl
+import org.jetbrains.kotlin.fir.expressions.impl.FirPartiallyResolvedArgumentList
+import org.jetbrains.kotlin.fir.expressions.impl.FirResolvedArgumentList
+import org.jetbrains.kotlin.fir.expressions.impl.FirSingleExpressionBlock
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.types.ConeClassLikeType
@@ -41,8 +43,18 @@ inline val FirCall.arguments: List<FirExpression> get() = argumentList.arguments
 
 inline val FirCall.argument: FirExpression get() = argumentList.arguments.first()
 
+inline val FirCall.resolvedArgumentMapping: LinkedHashMap<FirExpression, FirValueParameter>?
+    get() = when (val argumentList = argumentList) {
+        is FirResolvedArgumentList -> argumentList.mapping
+        else -> null
+    }
+
 inline val FirCall.argumentMapping: LinkedHashMap<FirExpression, FirValueParameter>?
-    get() = (argumentList as? FirResolvedArgumentList)?.mapping
+    get() = when (val argumentList = argumentList) {
+        is FirResolvedArgumentList -> argumentList.mapping
+        is FirPartiallyResolvedArgumentList -> argumentList.mapping
+        else -> null
+    }
 
 fun FirExpression.toResolvedCallableReference(): FirResolvedNamedReference? {
     if (this is FirWrappedArgumentExpression) return expression.toResolvedCallableReference()

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/impl/FirPartiallyResolvedArgumentList.kt

@@ -0,0 +1,46 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.expressions.impl
+
+import org.jetbrains.kotlin.fir.FirSourceElement
+import org.jetbrains.kotlin.fir.declarations.FirValueParameter
+import org.jetbrains.kotlin.fir.expressions.FirArgumentList
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.visitors.FirVisitor
+import org.jetbrains.kotlin.fir.visitors.transformSingle
+
+class FirPartiallyResolvedArgumentList internal constructor(
+    override var source: FirSourceElement?,
+    private var _mapping: LinkedHashMap<FirExpression, FirValueParameter?>
+) : FirArgumentList() {
+
+    @Suppress("UNCHECKED_CAST")
+    val mapping: LinkedHashMap<FirExpression, FirValueParameter> =
+        _mapping.filterValues { it != null } as LinkedHashMap<FirExpression, FirValueParameter>
+
+
+    override val arguments: List<FirExpression>
+        get() = _mapping.keys.toList()
+
+    override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
+        _mapping.forEach { (k, _) -> k.accept(visitor, data) }
+    }
+
+    override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirPartiallyResolvedArgumentList {
+        transformArguments(transformer, data)
+        return this
+    }
+
+    override fun <D> transformArguments(transformer: FirTransformer<D>, data: D): FirPartiallyResolvedArgumentList {
+        _mapping = _mapping.mapKeys { (k, _) -> k.transformSingle(transformer, data) } as LinkedHashMap<FirExpression, FirValueParameter?>
+        return this
+    }
+
+    override fun replaceSource(newSource: FirSourceElement?) {
+        source = newSource
+    }
+}

compiler/testData/diagnostics/tests/annotations/annotationParameterMustBeConstant/enumConst_after.fir.kt

@@ -11,5 +11,5 @@ enum class MyEnum {
     A
 }
 
-<!INAPPLICABLE_CANDIDATE!>@AnnE(Test())<!>
+<!INAPPLICABLE_CANDIDATE!>@AnnE(<!ANNOTATION_ARGUMENT_MUST_BE_CONST!>Test()<!>)<!>
 class Test2

compiler/testData/diagnostics/tests/annotations/annotationParameterMustBeConstant/enumConst_before.fir.kt

@@ -11,5 +11,5 @@ enum class MyEnum {
     A
 }
 
-<!INAPPLICABLE_CANDIDATE!>@AnnE(Test())<!>
+<!INAPPLICABLE_CANDIDATE!>@AnnE(<!ANNOTATION_ARGUMENT_MUST_BE_CONST!>Test()<!>)<!>
 class Test2

compiler/testData/diagnostics/tests/annotations/dontReportWarningAboutChangingExecutionOrderForVararg.fir.kt

@@ -7,7 +7,7 @@ fun foo1() {}
 @Anno(x = ["a", "b"], y = "a")
 fun foo2() {}
 
-<!INAPPLICABLE_CANDIDATE!>@Anno(x = arrayOf(arrayOf("a"), arrayOf("b")), y = "a")<!>
+<!INAPPLICABLE_CANDIDATE!>@Anno(x = <!NON_CONST_VAL_USED_IN_CONSTANT_EXPRESSION!>arrayOf(<!ANNOTATION_ARGUMENT_MUST_BE_CONST!>arrayOf("a")<!>, <!ANNOTATION_ARGUMENT_MUST_BE_CONST!>arrayOf("b")<!>)<!>, y = "a")<!>
 fun foo3() {}
 
 @Anno(x = arrayOf("a", "b"), y = "a")