[FIR] Unify FirResolvedArgumentList with FirPartiallyResolvedArgumentList

Effectively FirPartiallyResolvedArgumentList was resolved argument list
  for error calls, so not all arguments had proper value parameter. So
  FirPartiallyResolvedArgumentList now inherits FirResolvedArgumentList
  and renamed to FirResolvedArgumentListForErrorCall

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

@@ -205,7 +205,7 @@ class FirCallCompletionResultsWriterTransformer(
                 result.argumentList.transformArguments(this, expectedArgumentsTypeMapping)
                 if (calleeReference.isError) {
                     subCandidate.argumentMapping?.let {
-                        result.replaceArgumentList(buildPartiallyResolvedArgumentList(result.argumentList, it))
+                        result.replaceArgumentList(buildArgumentListForErrorCall(result.argumentList, it))
                     }
                 } else {
                     subCandidate.handleVarargs()
@@ -282,7 +282,7 @@ class FirCallCompletionResultsWriterTransformer(
         }
         if (calleeReference.isError) {
             subCandidate.argumentMapping?.let {
-                annotationCall.replaceArgumentList(buildPartiallyResolvedArgumentList(annotationCall.argumentList, it))
+                annotationCall.replaceArgumentList(buildArgumentListForErrorCall(annotationCall.argumentList, it))
             }
         } else {
             subCandidate.handleVarargs()
@@ -465,7 +465,7 @@ class FirCallCompletionResultsWriterTransformer(
         delegatedConstructorCall.argumentList.transformArguments(this, argumentsMapping)
         if (calleeReference.isError) {
             subCandidate.argumentMapping?.let {
-                delegatedConstructorCall.replaceArgumentList(buildPartiallyResolvedArgumentList(delegatedConstructorCall.argumentList, it))
+                delegatedConstructorCall.replaceArgumentList(buildArgumentListForErrorCall(delegatedConstructorCall.argumentList, it))
             }
         } else {
             subCandidate.handleVarargs()

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

@@ -9,8 +9,9 @@ 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.FirResolvedArgumentListForErrorCall
 import org.jetbrains.kotlin.fir.expressions.impl.FirResolvedArgumentList
+import org.jetbrains.kotlin.fir.expressions.impl.FirResolvedArgumentListImpl
 
 fun buildUnaryArgumentList(argument: FirExpression): FirArgumentList = buildArgumentList {
     arguments += argument
@@ -28,13 +29,13 @@ fun buildResolvedArgumentList(
     mapping: LinkedHashMap<FirExpression, FirValueParameter>,
     source: FirSourceElement? = null
 ): FirResolvedArgumentList =
-    FirResolvedArgumentList(mapping, source)
+    FirResolvedArgumentListImpl(source, mapping)
 
-fun buildPartiallyResolvedArgumentList(
+fun buildArgumentListForErrorCall(
     original: FirArgumentList,
-    mapping: LinkedHashMap<FirExpression, FirValueParameter>
+    mapping: Map<FirExpression, FirValueParameter?>
 ): FirArgumentList {
-    return FirPartiallyResolvedArgumentList(
+    return FirResolvedArgumentListForErrorCall(
         original.source,
         original.arguments.map { key -> key to mapping[key] }.toMap(LinkedHashMap())
     )

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

@@ -12,7 +12,6 @@ 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.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.FirReference
@@ -53,7 +52,6 @@ inline val FirCall.resolvedArgumentMapping: Map<FirExpression, FirValueParameter
 inline val FirCall.argumentMapping: LinkedHashMap<FirExpression, FirValueParameter>?
     get() = when (val argumentList = argumentList) {
         is FirResolvedArgumentList -> argumentList.mapping
-        is FirPartiallyResolvedArgumentList -> argumentList.mapping
         else -> null
     }
 

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

@@ -1,42 +0,0 @@
-/*
- * 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
-    }
-}

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.expressions.impl
 
+import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.declarations.FirValueParameter
 import org.jetbrains.kotlin.fir.expressions.FirAbstractArgumentList
@@ -14,25 +15,61 @@ import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 
-class FirResolvedArgumentList internal constructor(
+abstract class FirResolvedArgumentList : FirAbstractArgumentList() {
-    mapping: LinkedHashMap<FirExpression, FirValueParameter>,
+    abstract override val source: FirSourceElement?
-    override val source: FirSourceElement?
+    abstract val mapping: LinkedHashMap<FirExpression, FirValueParameter>
-) : FirAbstractArgumentList() {
-
-    var mapping: LinkedHashMap<FirExpression, FirValueParameter> = mapping
-        private set
 
     override val arguments: List<FirExpression>
         get() = mapping.keys.toList()
 
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
-        for (argument in mapping.keys) {
+        for (argument in arguments) {
             argument.accept(visitor, data)
         }
     }
 
+    abstract override fun <D> transformArguments(transformer: FirTransformer<D>, data: D): FirArgumentList
+
+    override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirElement {
+        transformArguments(transformer, data)
+        return this
+    }
+}
+
+
+internal class FirResolvedArgumentListImpl(
+    override val source: FirSourceElement?,
+    mapping: LinkedHashMap<FirExpression, FirValueParameter>
+) : FirResolvedArgumentList() {
+
+    override var mapping: LinkedHashMap<FirExpression, FirValueParameter> = mapping
+        private set
+
     override fun <D> transformArguments(transformer: FirTransformer<D>, data: D): FirArgumentList {
         mapping = mapping.mapKeys { (k, _) -> k.transformSingle(transformer, data) } as LinkedHashMap<FirExpression, FirValueParameter>
         return this
     }
 }
+
+internal class FirResolvedArgumentListForErrorCall(
+    override var source: FirSourceElement?,
+    private var _mapping: LinkedHashMap<FirExpression, FirValueParameter?>
+) : FirResolvedArgumentList() {
+
+    override var mapping: LinkedHashMap<FirExpression, FirValueParameter> = computeMapping()
+        private set
+
+    private fun computeMapping(): LinkedHashMap<FirExpression, FirValueParameter> {
+        @Suppress("UNCHECKED_CAST")
+        return _mapping.filterValues { it != null } as LinkedHashMap<FirExpression, FirValueParameter>
+    }
+
+    override val arguments: List<FirExpression>
+        get() = _mapping.keys.toList()
+
+    override fun <D> transformArguments(transformer: FirTransformer<D>, data: D): FirResolvedArgumentListForErrorCall {
+        _mapping = _mapping.mapKeys { (k, _) -> k.transformSingle(transformer, data) } as LinkedHashMap<FirExpression, FirValueParameter?>
+        mapping = computeMapping()
+        return this
+    }
+}