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.
2021-03-08 20:36:37 -08:00
parent 5d3afbad54
commit 9476175cc2
8 changed files with 90 additions and 10 deletions
@@ -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) {
@@ -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))
@@ -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()
@@ -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()
@@ -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
+    }
+}
@@ -11,5 +11,5 @@ enum class MyEnum {
    A
 }

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

-<!INAPPLICABLE_CANDIDATE!>@AnnE(Test())<!>
+<!INAPPLICABLE_CANDIDATE!>@AnnE(<!ANNOTATION_ARGUMENT_MUST_BE_CONST!>Test()<!>)<!>
 class Test2
@@ -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")