From 5892646a275e3f051e569124f48bb448835bf2cc Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Wed, 17 Feb 2021 16:25:49 +0300
Subject: [PATCH] Change FirQualifiedAccess source to KtQualifiedExpression (if
 any)

---
 .../checkers/expression/FirExhaustiveWhenChecker.kt | 13 +++++++++----
 .../fir/lightTree/converter/ExpressionsConverter.kt |  1 +
 .../jetbrains/kotlin/fir/builder/RawFirBuilder.kt   |  1 +
 .../testData/diagnostics/tests/DeferredTypes.fir.kt |  2 +-
 .../tests/FunctionCalleeExpressions.fir.kt          |  2 +-
 .../initializerScopeOfExtensionProperty.fir.kt      |  2 +-
 6 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirExhaustiveWhenChecker.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirExhaustiveWhenChecker.kt
index ace54dbcc92..913081955df 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirExhaustiveWhenChecker.kt
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirExhaustiveWhenChecker.kt
@@ -18,15 +18,20 @@ import org.jetbrains.kotlin.fir.expressions.isExhaustive
 object FirExhaustiveWhenChecker : FirWhenExpressionChecker() {
     override fun check(expression: FirWhenExpression, context: CheckerContext, reporter: DiagnosticReporter) {
         if (expression.usedAsExpression && !expression.isExhaustive) {
-            if (expression.source?.isIfExpression == true) {
-                reporter.reportOn(expression.source, FirErrors.INVALID_IF_AS_EXPRESSION, context)
-            } else {
+            val source = expression.source ?: return
+            if (source.isIfExpression) {
+                reporter.reportOn(source, FirErrors.INVALID_IF_AS_EXPRESSION, context)
+                return
+            } else if (source.isWhenExpression) {
                 val missingCases = (expression.exhaustivenessStatus as ExhaustivenessStatus.NotExhaustive).reasons
-                reporter.reportOn(expression.source, FirErrors.NO_ELSE_IN_WHEN, missingCases, context)
+                reporter.reportOn(source, FirErrors.NO_ELSE_IN_WHEN, missingCases, context)
             }
         }
     }
 
     private val FirSourceElement.isIfExpression: Boolean
         get() = elementType == KtNodeTypes.IF
+
+    private val FirSourceElement.isWhenExpression: Boolean
+        get() = elementType == KtNodeTypes.WHEN
 }
diff --git a/compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt b/compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt
index 394b650ec39..0473bc56df1 100644
--- a/compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt
+++ b/compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt
@@ -499,6 +499,7 @@ class ExpressionsConverter(
 
             it.replaceExplicitReceiver(firReceiver)
         }
+        firSelector.replaceSource(dotQualifiedExpression.toFirSourceElement())
         return firSelector
     }
 
diff --git a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt b/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt
index 7b47a505f95..25e02f27f9c 100644
--- a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt
+++ b/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt
@@ -1985,6 +1985,7 @@ class RawFirBuilder(
 
                 firSelector.replaceExplicitReceiver(receiver)
             }
+            firSelector.replaceSource(expression.toFirSourceElement())
             return firSelector
         }
 
diff --git a/compiler/testData/diagnostics/tests/DeferredTypes.fir.kt b/compiler/testData/diagnostics/tests/DeferredTypes.fir.kt
index 3647a682a51..7febd7aed4c 100644
--- a/compiler/testData/diagnostics/tests/DeferredTypes.fir.kt
+++ b/compiler/testData/diagnostics/tests/DeferredTypes.fir.kt
@@ -1,5 +1,5 @@
 // NI_EXPECTED_FILE
 
 interface T {
-  val a = <!UNRESOLVED_REFERENCE!>Foo<!>.<!PROPERTY_INITIALIZER_IN_INTERFACE!><!UNRESOLVED_REFERENCE!>bar<!>()<!>
+  val a = <!PROPERTY_INITIALIZER_IN_INTERFACE!><!UNRESOLVED_REFERENCE!>Foo<!>.<!UNRESOLVED_REFERENCE!>bar<!>()<!>
 }
\ No newline at end of file
diff --git a/compiler/testData/diagnostics/tests/FunctionCalleeExpressions.fir.kt b/compiler/testData/diagnostics/tests/FunctionCalleeExpressions.fir.kt
index 14ccbb2caa7..6a62617d900 100644
--- a/compiler/testData/diagnostics/tests/FunctionCalleeExpressions.fir.kt
+++ b/compiler/testData/diagnostics/tests/FunctionCalleeExpressions.fir.kt
@@ -72,7 +72,7 @@ fun main1() {
 
     1."sdf"
     1.{}
-    1.<!INVALID_IF_AS_EXPRESSION!>if<!> (true) {}
+    1.if (true) {}
 }
 
 fun test() {
diff --git a/compiler/testData/diagnostics/tests/scopes/initializerScopeOfExtensionProperty.fir.kt b/compiler/testData/diagnostics/tests/scopes/initializerScopeOfExtensionProperty.fir.kt
index 7ce78c898ac..8ec4257a2cd 100644
--- a/compiler/testData/diagnostics/tests/scopes/initializerScopeOfExtensionProperty.fir.kt
+++ b/compiler/testData/diagnostics/tests/scopes/initializerScopeOfExtensionProperty.fir.kt
@@ -26,6 +26,6 @@ class C {
 
 val C.foo : C.D = <!EXTENSION_PROPERTY_WITH_BACKING_FIELD!><!UNRESOLVED_REFERENCE!>D<!>()<!>
 
-val C.bar : C.D = C().<!EXTENSION_PROPERTY_WITH_BACKING_FIELD!>D()<!>
+val C.bar : C.D = <!EXTENSION_PROPERTY_WITH_BACKING_FIELD!>C().D()<!>
 
 val C.foo1 : C.D get() = D()