From b0b198119da8b5cd3bfbab79434100bd7b2e285f Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Wed, 22 Apr 2020 16:19:42 +0300
Subject: [PATCH] FIR resolve completion: transform receivers for variable
 accesses

---
 .../invoke/threeReceiversCorrect.txt          |  2 +-
 ...rCallCompletionResultsWriterTransformer.kt | 44 +++++++++++++------
 2 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/compiler/fir/analysis-tests/testData/resolve/expresssions/invoke/threeReceiversCorrect.txt b/compiler/fir/analysis-tests/testData/resolve/expresssions/invoke/threeReceiversCorrect.txt
index 93bf34fe658..9de0198e18e 100644
--- a/compiler/fir/analysis-tests/testData/resolve/expresssions/invoke/threeReceiversCorrect.txt
+++ b/compiler/fir/analysis-tests/testData/resolve/expresssions/invoke/threeReceiversCorrect.txt
@@ -32,7 +32,7 @@ FILE: threeReceiversCorrect.kt
             }
             )
             ^ R|/with|<R|B|, R|kotlin/Unit|>(R|<local>/b|, <L> = with@fun R|B|.<anonymous>(): R|kotlin/Unit| {
-                this@R|special/anonymous|.R|/A.foo|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/Unit|>|(R|<local>/c|)
+                (this@R|special/anonymous|, this@R|special/anonymous|).R|/A.foo|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/Unit|>|(R|<local>/c|)
                 ^ R|/with|<R|C|, R|kotlin/Unit|>(R|<local>/c|, <L> = with@fun R|C|.<anonymous>(): R|kotlin/Unit| {
                     (this@R|special/anonymous|, this@R|special/anonymous|).R|/A.foo|.R|FakeOverride<kotlin/Function1.invoke: R|kotlin/Unit|>|(this@R|special/anonymous|)
                 }
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt
index 1b0029cb50b..6fb522e2cfa 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt
@@ -56,27 +56,43 @@ class FirCallCompletionResultsWriterTransformer(
         val calleeReference =
             qualifiedAccessExpression.calleeReference as? FirNamedReferenceWithCandidate ?: return qualifiedAccessExpression.compose()
 
+        val subCandidate = calleeReference.candidate
         val candidateFir = calleeReference.candidateSymbol.phasedFir
         val typeArguments = computeTypeArguments(qualifiedAccessExpression, calleeReference.candidate)
-        val typeRef = (candidateFir as? FirTypedDeclaration)?.let {
-            typeCalculator.tryCalculateReturnType(it)
+        val typeRef = (candidateFir as? FirTypedDeclaration)?.let { typedDeclaration ->
+            typeCalculator.tryCalculateReturnType(typedDeclaration).let {
+                if (qualifiedAccessExpression.safe) {
+                    val nullableType = it.coneTypeUnsafe<ConeKotlinType>().withNullability(
+                        ConeNullability.NULLABLE, session.inferenceContext
+                    )
+                    it.withReplacedConeType(nullableType)
+                } else {
+                    it
+                }
+            }
         } ?: buildErrorTypeRef {
             source = calleeReference.source
             diagnostic = ConeSimpleDiagnostic("Callee reference to candidate without return type: ${candidateFir.render()}")
         }
 
-        qualifiedAccessExpression.replaceTypeRefWithSubstituted(calleeReference, typeRef).replaceTypeArguments(typeArguments)
+        val result = qualifiedAccessExpression
+            .transformCalleeReference(
+                StoreCalleeReference,
+                buildResolvedNamedReference {
+                    source = calleeReference.source
+                    name = calleeReference.name
+                    resolvedSymbol = calleeReference.candidateSymbol
+                },
+            )
+            .transformDispatchReceiver(StoreReceiver, subCandidate.dispatchReceiverExpression())
+            .transformExtensionReceiver(StoreReceiver, subCandidate.extensionReceiverExpression())
+            .transformExplicitReceiver(integerApproximator, null)
 
-        return qualifiedAccessExpression.transformCalleeReference(
-            StoreCalleeReference,
-            buildResolvedNamedReference {
-                source = calleeReference.source
-                name = calleeReference.name
-                resolvedSymbol = calleeReference.candidateSymbol
-            },
-        ).apply {
-            replaceTypeArguments(typeArguments)
-        }.compose()
+        val resultType = typeRef.substituteTypeRef(subCandidate)
+        result.replaceTypeRef(resultType)
+        result.replaceTypeArguments(typeArguments)
+
+        return result.compose()
     }
 
     private fun <D : FirExpression> D.replaceTypeRefWithSubstituted(
@@ -146,7 +162,7 @@ class FirCallCompletionResultsWriterTransformer(
                 name = calleeReference.name
                 resolvedSymbol = calleeReference.candidateSymbol
             },
-        ).apply {
+        ).transformExplicitReceiver(integerApproximator, null).apply {
             replaceTypeArguments(typeArguments)
         }.compose()
     }