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/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
+    }
+}