FIR: Rework FirRegularTowerDataContexts

- Get rid of SPECIAL mode (just left REGULAR and class-related instead)
- Clear naming
- Restore contexts after lambda/callable reference are processed

The test has been failing before this change because after
callable reference is resolved, its tower data context has been left
erroneously in the SPECIAL-related entry

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosisCompilerTestFE10TestdataTestGenerated.java

@@ -7164,6 +7164,12 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                 runTest("compiler/testData/diagnostics/tests/dataFlow/IsExpression.kt");
             }
 
+            @Test
+            @TestMetadata("smartCastWithLambdaAndCallableReference.kt")
+            public void testSmartCastWithLambdaAndCallableReference() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/dataFlow/smartCastWithLambdaAndCallableReference.kt");
+            }
+
             @Test
             @TestMetadata("WhenSubject.kt")
             public void testWhenSubject() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsTestGenerated.java

@@ -7164,6 +7164,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                 runTest("compiler/testData/diagnostics/tests/dataFlow/IsExpression.kt");
             }
 
+            @Test
+            @TestMetadata("smartCastWithLambdaAndCallableReference.kt")
+            public void testSmartCastWithLambdaAndCallableReference() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/dataFlow/smartCastWithLambdaAndCallableReference.kt");
+            }
+
             @Test
             @TestMetadata("WhenSubject.kt")
             public void testWhenSubject() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsWithLightTreeTestGenerated.java

@@ -7164,6 +7164,12 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac
                 runTest("compiler/testData/diagnostics/tests/dataFlow/IsExpression.kt");
             }
 
+            @Test
+            @TestMetadata("smartCastWithLambdaAndCallableReference.kt")
+            public void testSmartCastWithLambdaAndCallableReference() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/dataFlow/smartCastWithLambdaAndCallableReference.kt");
+            }
+
             @Test
             @TestMetadata("WhenSubject.kt")
             public void testWhenSubject() throws Exception {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/FirRegularTowerDataContexts.kt

@@ -10,7 +10,7 @@ import org.jetbrains.kotlin.fir.scopes.impl.FirLocalScope
 import java.util.*
 
 enum class FirTowerDataMode {
-    MEMBER_DECLARATION,
+    REGULAR,
     NESTED_CLASS,
     COMPANION_OBJECT,
     CONSTRUCTOR_HEADER,
@@ -22,65 +22,64 @@ class FirRegularTowerDataContexts private constructor(
     private val modeMap: EnumMap<FirTowerDataMode, FirTowerDataContext>,
     val primaryConstructorPureParametersScope: FirLocalScope?,
     val primaryConstructorAllParametersScope: FirLocalScope?,
-    val mode: FirTowerDataMode,
+    val activeMode: FirTowerDataMode,
 ) {
     constructor(
-        forMemberDeclarations: FirTowerDataContext,
+        regular: FirTowerDataContext,
         forNestedClasses: FirTowerDataContext? = null,
         forCompanionObject: FirTowerDataContext? = null,
         forConstructorHeaders: FirTowerDataContext? = null,
         forEnumEntries: FirTowerDataContext? = null,
-        forSpecial: FirTowerDataContext? = null,
         primaryConstructorPureParametersScope: FirLocalScope? = null,
         primaryConstructorAllParametersScope: FirLocalScope? = null,
     ) : this(
-        enumMap(forMemberDeclarations, forNestedClasses, forCompanionObject, forConstructorHeaders, forEnumEntries, forSpecial),
+        enumMap(regular, forNestedClasses, forCompanionObject, forConstructorHeaders, forEnumEntries),
         primaryConstructorPureParametersScope,
         primaryConstructorAllParametersScope,
-        FirTowerDataMode.MEMBER_DECLARATION
+        FirTowerDataMode.REGULAR
     )
 
     val currentContext: FirTowerDataContext?
-        get() = modeMap[mode]
+        get() = modeMap[activeMode]
 
-    fun copy(newContext: FirTowerDataContext): FirRegularTowerDataContexts {
+    fun replaceCurrentlyActiveContext(newContext: FirTowerDataContext): FirRegularTowerDataContexts {
         val modeMap = EnumMap<FirTowerDataMode, FirTowerDataContext>(FirTowerDataMode::class.java)
         modeMap.putAll(this.modeMap)
-        modeMap[mode] = newContext
+        modeMap[activeMode] = newContext
-        return FirRegularTowerDataContexts(modeMap, primaryConstructorPureParametersScope, primaryConstructorAllParametersScope, mode)
+        return FirRegularTowerDataContexts(modeMap, primaryConstructorPureParametersScope, primaryConstructorAllParametersScope, activeMode)
     }
 
-    fun copy(newMode: FirTowerDataMode): FirRegularTowerDataContexts {
+    fun replaceTowerDataMode(newMode: FirTowerDataMode): FirRegularTowerDataContexts {
-        if (newMode == mode) return this
+        if (newMode == activeMode) return this
         return FirRegularTowerDataContexts(modeMap, primaryConstructorPureParametersScope, primaryConstructorAllParametersScope, newMode)
     }
 
-    fun copyWithSpecial(newContext: FirTowerDataContext): FirRegularTowerDataContexts {
+    // Effectively equal to replaceTowerDataMode(REGULAR) + replaceCurrentlyActiveContext(newContext)
+    // But left just for sake of optimization
+    fun replaceAndSetActiveRegularContext(newContext: FirTowerDataContext): FirRegularTowerDataContexts {
         val modeMap = EnumMap<FirTowerDataMode, FirTowerDataContext>(FirTowerDataMode::class.java)
         modeMap.putAll(this.modeMap)
-        modeMap[FirTowerDataMode.SPECIAL] = newContext
+        modeMap[FirTowerDataMode.REGULAR] = newContext
         return FirRegularTowerDataContexts(
-            modeMap, primaryConstructorPureParametersScope, primaryConstructorAllParametersScope, FirTowerDataMode.SPECIAL
+            modeMap, primaryConstructorPureParametersScope, primaryConstructorAllParametersScope, FirTowerDataMode.REGULAR
         )
     }
 
     companion object {
         private fun enumMap(
-            forMemberDeclarations: FirTowerDataContext,
+            regular: FirTowerDataContext,
             forNestedClasses: FirTowerDataContext?,
             forCompanionObject: FirTowerDataContext?,
             forConstructorHeaders: FirTowerDataContext?,
             forEnumEntries: FirTowerDataContext?,
-            forSpecial: FirTowerDataContext?,
         ): EnumMap<FirTowerDataMode, FirTowerDataContext> {
             val modeMap = EnumMap<FirTowerDataMode, FirTowerDataContext>(FirTowerDataMode::class.java)
-            modeMap[FirTowerDataMode.MEMBER_DECLARATION] = forMemberDeclarations
+            modeMap[FirTowerDataMode.REGULAR] = regular
             modeMap[FirTowerDataMode.NESTED_CLASS] = forNestedClasses
             modeMap[FirTowerDataMode.COMPANION_OBJECT] = forCompanionObject
             modeMap[FirTowerDataMode.CONSTRUCTOR_HEADER] = forConstructorHeaders
             modeMap[FirTowerDataMode.ENUM_ENTRY] = forEnumEntries
-            modeMap[FirTowerDataMode.SPECIAL] = forSpecial
             return modeMap
         }
     }
-}
+}

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt

@@ -52,7 +52,7 @@ class BodyResolveContext(
     lateinit var file: FirFile
 
     @PrivateForInline
-    var regularTowerDataContexts = FirRegularTowerDataContexts(forMemberDeclarations = FirTowerDataContext())
+    var regularTowerDataContexts = FirRegularTowerDataContexts(regular = FirTowerDataContext())
 
     @PrivateForInline
     val specialTowerDataContexts = FirSpecialTowerDataContexts()
@@ -64,9 +64,9 @@ class BodyResolveContext(
 
     @OptIn(PrivateForInline::class)
     var towerDataMode: FirTowerDataMode
-        get() = regularTowerDataContexts.mode
+        get() = regularTowerDataContexts.activeMode
         set(value) {
-            regularTowerDataContexts = regularTowerDataContexts.copy(newMode = value)
+            regularTowerDataContexts = regularTowerDataContexts.replaceTowerDataMode(newMode = value)
         }
 
     val implicitReceiverStack: ImplicitReceiverStack
@@ -95,7 +95,7 @@ class BodyResolveContext(
         get() = containingClassDeclarations.lastOrNull()
 
     @OptIn(PrivateForInline::class)
-    inline fun <T> withNewTowerDataForClass(newContexts: FirRegularTowerDataContexts, f: () -> T): T {
+    inline fun <T> withTowerDataContexts(newContexts: FirRegularTowerDataContexts, f: () -> T): T {
         val old = regularTowerDataContexts
         regularTowerDataContexts = newContexts
         return try {
@@ -158,14 +158,14 @@ class BodyResolveContext(
 
     @PrivateForInline
     inline fun <T> withTowerDataMode(mode: FirTowerDataMode, f: () -> T): T {
-        return withTowerModeCleanup {
+        return withTowerDataModeCleanup {
             towerDataMode = mode
             f()
         }
     }
 
     @PrivateForInline
-    inline fun <R> withTowerModeCleanup(l: () -> R): R {
+    inline fun <R> withTowerDataModeCleanup(l: () -> R): R {
         val initialMode = towerDataMode
         return try {
             l()
@@ -176,7 +176,7 @@ class BodyResolveContext(
 
     @PrivateForInline
     fun replaceTowerDataContext(newContext: FirTowerDataContext) {
-        regularTowerDataContexts = regularTowerDataContexts.copy(newContext)
+        regularTowerDataContexts = regularTowerDataContexts.replaceCurrentlyActiveContext(newContext)
     }
 
     @PrivateForInline
@@ -315,23 +315,19 @@ class BodyResolveContext(
 
     @OptIn(PrivateForInline::class)
     inline fun <T> withAnonymousFunctionTowerDataContext(symbol: FirAnonymousFunctionSymbol, f: () -> T): T {
-        return withTowerModeCleanup {
+        return withTemporaryRegularContext(specialTowerDataContexts.getAnonymousFunctionContext(symbol), f)
-            val newContext = specialTowerDataContexts.getAnonymousFunctionContext(symbol)
-            if (newContext != null) {
-                regularTowerDataContexts = regularTowerDataContexts.copyWithSpecial(newContext)
-            }
-            f()
-        }
     }
 
     @OptIn(PrivateForInline::class)
     inline fun <T> withCallableReferenceTowerDataContext(access: FirCallableReferenceAccess, f: () -> T): T {
-        return withTowerModeCleanup {
+        return withTemporaryRegularContext(specialTowerDataContexts.getCallableReferenceContext(access), f)
-            val newContext = specialTowerDataContexts.getCallableReferenceContext(access)
+    }
-            if (newContext != null) {
+
-                regularTowerDataContexts = regularTowerDataContexts.copyWithSpecial(newContext)
+    @PrivateForInline
-            }
+    inline fun <T> withTemporaryRegularContext(newContext: FirTowerDataContext?, f: () -> T): T {
-            f()
+        if (newContext == null) return f()
+        return withTowerDataModeCleanup {
+            withTowerDataContexts(regularTowerDataContexts.replaceAndSetActiveRegularContext(newContext), f)
         }
     }
 
@@ -396,7 +392,7 @@ class BodyResolveContext(
                 withContainerClass(regularClass, f)
             }
         }
-        return withTowerModeCleanup {
+        return withTowerDataModeCleanup {
             if (!regularClass.isInner && containerIfAny is FirRegularClass) {
                 towerDataMode = if (regularClass.isCompanion) {
                     FirTowerDataMode.COMPANION_OBJECT
@@ -486,12 +482,11 @@ class BodyResolveContext(
             statics,
             scopeForConstructorHeader,
             scopeForEnumEntries,
-            forSpecial = null,
             primaryConstructorPureParametersScope,
             primaryConstructorAllParametersScope
         )
 
-        return withNewTowerDataForClass(newContexts) {
+        return withTowerDataContexts(newContexts) {
             f()
         }
     }
@@ -521,7 +516,7 @@ class BodyResolveContext(
         val whenSubjectImportingScope = whenSubjectImportingScopes.lastOrNull() ?: return f()
         val newTowerDataContext = towerDataContext.addNonLocalScope(whenSubjectImportingScope)
         val newContexts = FirRegularTowerDataContexts(newTowerDataContext)
-        return withNewTowerDataForClass(newContexts) {
+        return withTowerDataContexts(newContexts) {
             f()
         }
     }

compiler/testData/diagnostics/tests/dataFlow/smartCastWithLambdaAndCallableReference.fir.kt

@@ -0,0 +1,21 @@
+// SKIP_TXT
+
+fun <T> T.myApply(block: T.() -> Unit): T = this
+
+fun bar(): Int = 1
+
+interface A : C
+interface B : C
+interface C {
+    fun baz()
+}
+
+fun Any.foo() = myApply {
+    when (this) {
+        is A -> ::bar
+        is B -> ::bar
+        else -> throw RuntimeException()
+    }
+
+    baz() // Smart cast should work
+}

compiler/testData/diagnostics/tests/dataFlow/smartCastWithLambdaAndCallableReference.kt

@@ -0,0 +1,21 @@
+// SKIP_TXT
+
+fun <T> T.myApply(block: T.() -> Unit): T = this
+
+fun bar(): Int = 1
+
+interface A : C
+interface B : C
+interface C {
+    fun baz()
+}
+
+fun Any.foo() = myApply {
+    when (this) {
+        is A -> ::bar
+        is B -> ::bar
+        else -> throw RuntimeException()
+    }
+
+    <!UNRESOLVED_REFERENCE!>baz<!>() // Smart cast should work
+}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java

@@ -7170,6 +7170,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                 runTest("compiler/testData/diagnostics/tests/dataFlow/IsExpression.kt");
             }
 
+            @Test
+            @TestMetadata("smartCastWithLambdaAndCallableReference.kt")
+            public void testSmartCastWithLambdaAndCallableReference() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/dataFlow/smartCastWithLambdaAndCallableReference.kt");
+            }
+
             @Test
             @TestMetadata("WhenSubject.kt")
             public void testWhenSubject() throws Exception {