From a4c5e1bc873eb59462b4fb8e3c6bf75f3172320c Mon Sep 17 00:00:00 2001
From: "Denis.Zharkov" <denis.zharkov@jetbrains.com>
Date: Wed, 7 Dec 2022 15:43:48 +0100
Subject: [PATCH] K2: Unwrap fake overrides before using them as DFA
 identifiers

It's necessary because even for stable a.b.c.d we can't guarantee that
this reference will always point to the same symbol because
different capture type instantiations generate different scopes
with different resulting symbol instances.
---
 .../jetbrains/kotlin/fir/resolve/dfa/util.kt  | 20 ++++++++++++++++---
 .../tests/smartCasts/inference/kt25432.fir.kt |  2 +-
 ...rtCastOnValueBoundToSubjectVariable.fir.kt |  2 +-
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/util.kt b/compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/util.kt
index 3b03053cfa0..a2945630b5f 100644
--- a/compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/util.kt
+++ b/compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/util.kt
@@ -12,9 +12,13 @@ import org.jetbrains.kotlin.fir.references.FirNamedReferenceWithCandidateBase
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.references.FirThisReference
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirSyntheticPropertySymbol
-import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.ConeTypeContext
+import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.fir.unwrapFakeOverrides
 
 fun TypeStatement?.smartCastedType(context: ConeTypeContext, originalType: ConeKotlinType): ConeKotlinType =
     if (this != null && exactType.isNotEmpty()) {
@@ -39,9 +43,9 @@ val FirExpression.coneType: ConeKotlinType
 @DfaInternals
 val FirElement.symbol: FirBasedSymbol<*>?
     get() = when (this) {
-        is FirResolvable -> symbol
+        is FirResolvable -> symbol.unwrapFakeOverridesIfNecessary()
         is FirVariableAssignment -> unwrapLValue()?.symbol
-        is FirDeclaration -> symbol
+        is FirDeclaration -> symbol.unwrapFakeOverridesIfNecessary()
         is FirWhenSubjectExpression -> whenRef.value.subject?.symbol
         is FirSafeCallExpression -> selector.symbol
         is FirSmartCastExpression -> originalExpression.symbol
@@ -52,6 +56,16 @@ val FirElement.symbol: FirBasedSymbol<*>?
                 (it !is FirFunctionSymbol<*> && it !is FirSyntheticPropertySymbol)
     }
 
+private fun FirBasedSymbol<*>?.unwrapFakeOverridesIfNecessary(): FirBasedSymbol<*>? {
+    if (this !is FirCallableSymbol) return this
+    // This is necessary only for sake of optimizations necessary because this is a really hot place.
+    // Not having `dispatchReceiverType` means that this is a local/top-level property that can't be a fake override.
+    // And at the same time, checking a field is much faster than a couple of attributes (0.3 secs at MT Full Kotlin)
+    if (this.dispatchReceiverType == null) return this
+
+    return this.unwrapFakeOverrides()
+}
+
 @DfaInternals
 internal val FirResolvable.symbol: FirBasedSymbol<*>?
     get() = when (val reference = calleeReference) {
diff --git a/compiler/testData/diagnostics/tests/smartCasts/inference/kt25432.fir.kt b/compiler/testData/diagnostics/tests/smartCasts/inference/kt25432.fir.kt
index 15e4c2c716c..fad47bd5e27 100644
--- a/compiler/testData/diagnostics/tests/smartCasts/inference/kt25432.fir.kt
+++ b/compiler/testData/diagnostics/tests/smartCasts/inference/kt25432.fir.kt
@@ -4,7 +4,7 @@ class Data<T>(val s: T)
 
 fun test(d: Data<out Any>) {
     if (d.s is String) {
-        d.s.<!UNRESOLVED_REFERENCE!>length<!>
+        d.s.length
     }
 }
 
diff --git a/compiler/testData/diagnostics/tests/when/withSubjectVariable/smartCastOnValueBoundToSubjectVariable.fir.kt b/compiler/testData/diagnostics/tests/when/withSubjectVariable/smartCastOnValueBoundToSubjectVariable.fir.kt
index d98e52b8cf0..472cdcd7b70 100644
--- a/compiler/testData/diagnostics/tests/when/withSubjectVariable/smartCastOnValueBoundToSubjectVariable.fir.kt
+++ b/compiler/testData/diagnostics/tests/when/withSubjectVariable/smartCastOnValueBoundToSubjectVariable.fir.kt
@@ -43,6 +43,6 @@ fun test5(x: Inv<out Any?>) {
 
 fun test6(x: Inv<out String?>) {
     when (val y = x.data) {
-        is String -> x.data<!UNSAFE_CALL!>.<!>length // should be ok
+        is String -> x.data.length // should be ok
     }
 }