[LL API] Pull smart casts from the CFG in 'ContextCollector'

^KT-61599 Fixed

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt

@@ -225,7 +225,8 @@ private class LLFirBodyTargetResolver(
                     withPsiEntry("contextPsiElement", contextPsiElement)
                 }
 
-            LLFirCodeFragmentContext(elementContext.towerDataContext.withExtraScopes(), elementContext.smartCasts)
+            val variables = elementContext.smartCasts.mapKeys { it.key.identifier.symbol }
+            LLFirCodeFragmentContext(elementContext.towerDataContext.withExtraScopes(), variables)
         } else {
             val towerDataContext = FirTowerDataContext().withExtraScopes()
             LLFirCodeFragmentContext(towerDataContext, emptyMap())

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt

@@ -7,8 +7,6 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.util
 
 import com.intellij.openapi.progress.ProgressManager
 import com.intellij.psi.PsiElement
-import kotlinx.collections.immutable.PersistentMap
-import kotlinx.collections.immutable.persistentMapOf
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignation
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.ContextCollector.ContextKind
@@ -23,11 +21,17 @@ import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.SessionHolder
 import org.jetbrains.kotlin.fir.resolve.dfa.DataFlowAnalyzerContext
+import org.jetbrains.kotlin.fir.resolve.dfa.PropertyStability
+import org.jetbrains.kotlin.fir.resolve.dfa.RealVariable
+import org.jetbrains.kotlin.fir.resolve.dfa.cfg.CFGNode
+import org.jetbrains.kotlin.fir.resolve.dfa.cfg.ClassExitNode
+import org.jetbrains.kotlin.fir.resolve.dfa.controlFlowGraph
+import org.jetbrains.kotlin.fir.resolve.dfa.smartCastedType
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculatorForFullBodyResolve
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.BodyResolveContext
-import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.typeContext
 import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitorVoid
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
@@ -43,7 +47,7 @@ internal object ContextCollector {
 
     class Context(
         val towerDataContext: FirTowerDataContext,
-        val smartCasts: Map<FirBasedSymbol<*>, Set<ConeKotlinType>>,
+        val smartCasts: Map<RealVariable, Set<ConeKotlinType>>,
     )
 
     enum class FilterResponse {
@@ -171,8 +175,6 @@ private class ContextCollectorVisitor(
         dataFlowAnalyzerContext = DataFlowAnalyzerContext(session)
     )
 
-    private var smartCasts: PersistentMap<FirBasedSymbol<*>, Set<ConeKotlinType>> = persistentMapOf()
-
     private val result = HashMap<ContextKey, Context>()
 
     override fun visitElement(element: FirElement) {
@@ -199,7 +201,7 @@ private class ContextCollectorVisitor(
 
         val response = filter(psi)
         if (response != FilterResponse.SKIP) {
-            result[key] = Context(context.towerDataContext, smartCasts)
+            result[key] = computeContext(fir)
         }
 
         if (response == FilterResponse.STOP) {
@@ -207,6 +209,75 @@ private class ContextCollectorVisitor(
         }
     }
 
+    private fun computeContext(fir: FirElement): Context {
+        val implicitReceiverStack = context.towerDataContext.implicitReceiverStack
+
+        val smartCasts = mutableMapOf<RealVariable, Set<ConeKotlinType>>()
+
+        // Receiver types cannot be updated in an immutable snapshot.
+        // So here we modify the types inside the 'context', then make a snapshot, and restore the types back.
+        val oldReceiverTypes = mutableListOf<Pair<Int, ConeKotlinType>>()
+
+        val cfgNode = getControlFlowNode(fir)
+
+        if (cfgNode != null) {
+            val flow = cfgNode.flow
+
+            for (realVariable in flow.knownVariables) {
+                val typeStatement = flow.getTypeStatement(realVariable) ?: continue
+                if (realVariable.stability != PropertyStability.STABLE_VALUE && realVariable.stability != PropertyStability.LOCAL_VAR) {
+                    continue
+                }
+
+                smartCasts[typeStatement.variable] = typeStatement.exactType
+
+                // The compiler pushes smart-cast types for implicit receivers to ease later lookups.
+                // Here we emulate such behavior. Unlike the compiler, though, modified types are only reflected in the created snapshot.
+                // See other usages of 'replaceReceiverType()' for more information.
+                if (realVariable.isThisReference) {
+                    val identifier = typeStatement.variable.identifier
+                    val receiverIndex = implicitReceiverStack.getReceiverIndex(identifier.symbol)
+                    if (receiverIndex != null) {
+                        oldReceiverTypes.add(receiverIndex to implicitReceiverStack.getType(receiverIndex))
+
+                        val originalType = implicitReceiverStack.getOriginalType(receiverIndex)
+                        val smartCastedType = typeStatement.smartCastedType(session.typeContext, originalType)
+                        implicitReceiverStack.replaceReceiverType(receiverIndex, smartCastedType)
+                    }
+                }
+            }
+        }
+
+        val towerDataContextSnapshot = context.towerDataContext.createSnapshot()
+
+        for ((index, oldType) in oldReceiverTypes) {
+            implicitReceiverStack.replaceReceiverType(index, oldType)
+        }
+
+        return Context(towerDataContextSnapshot, smartCasts)
+    }
+
+    private fun getControlFlowNode(fir: FirElement): CFGNode<*>? {
+        for (container in context.containers.asReversed()) {
+            val cfgOwner = container as? FirControlFlowGraphOwner ?: continue
+            val cfgReference = cfgOwner.controlFlowGraphReference ?: continue
+            val cfg = cfgReference.controlFlowGraph ?: continue
+
+            val node = cfg.nodes.lastOrNull { isAcceptedControlFlowNode(it) && it.fir === fir }
+            if (node != null) {
+                return node
+            } else if (!cfg.isSubGraph) {
+                return null
+            }
+        }
+
+        return null
+    }
+
+    private fun isAcceptedControlFlowNode(node: CFGNode<*>): Boolean {
+        return node !is ClassExitNode
+    }
+
     override fun visitFile(file: FirFile) {
         context.withFile(file, holder) {
             withInterceptor {
@@ -477,24 +548,6 @@ private class ContextCollectorVisitor(
         }
     }
 
-    override fun visitSmartCastExpression(smartCastExpression: FirSmartCastExpression) {
-        if (smartCastExpression.isStable) {
-            val symbol = smartCastExpression.originalExpression.toResolvedCallableSymbol()
-            if (symbol != null) {
-                val previousSmartCasts = smartCasts
-                try {
-                    smartCasts = smartCasts.put(symbol, smartCastExpression.typesFromSmartCast.toSet())
-                    super.visitSmartCastExpression(smartCastExpression)
-                    return
-                } finally {
-                    smartCasts = previousSmartCasts
-                }
-            }
-        }
-
-        super.visitSmartCastExpression(smartCastExpression)
-    }
-
     @ContextCollectorDsl
     private fun Processor.processSignatureAnnotations(declaration: FirDeclaration) {
         for (annotation in declaration.annotations) {

analysis/low-level-api-fir/testdata/contextCollector/localClass.kt

@@ -0,0 +1,9 @@
+fun test() {
+    class Local {
+        fun foo(): String {
+            <expr>return "foo"</expr>
+        }
+    }
+
+    Local().foo()
+}

analysis/low-level-api-fir/testdata/contextCollector/localClass.txt

@@ -0,0 +1,45 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirLocalScope
+    Element 8
+        Scope: FirLocalScope
+            Classifiers:
+                FirRegularClassSymbol local final class Local : R|kotlin/Any|
+    Element 9
+        Implicit receiver:
+            FirRegularClassSymbol local final class Local : R|kotlin/Any|
+                Type: <local>/Local
+    Element 10
+        Scope: FirLocalScope
+    Element 11
+        Scope: FirLocalScope
+
+FILE: [ResolvedTo(IMPORTS)] localClass.kt
+    public final [ResolvedTo(BODY_RESOLVE)] fun test(): R|kotlin/Unit| {
+        local final [ResolvedTo(BODY_RESOLVE)] class Local : R|kotlin/Any| {
+            public [ResolvedTo(BODY_RESOLVE)] constructor(): R|<local>/Local| {
+                super<R|kotlin/Any|>()
+            }
+
+            public final [ResolvedTo(BODY_RESOLVE)] fun foo(): R|kotlin/String| {
+                ^foo String(foo)
+            }
+
+        }
+
+        R|<local>/Local.Local|().R|<local>/foo|()
+    }

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/afterIf.kt

@@ -0,0 +1,10 @@
+interface Foo {}
+
+fun test(obj: Any) {
+    if (obj is Foo) {
+        consume(obj)
+    }
+    <expr>obj.hashCode()</expr>
+}
+
+fun consume(obj: Foo) {}

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/afterIf.txt

@@ -0,0 +1,36 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirLocalScope
+            Properties:
+                FirValueParameterSymbol obj: R|kotlin/Any|
+    Element 8
+        Scope: FirLocalScope
+
+FILE: [ResolvedTo(IMPORTS)] afterIf.kt
+    public abstract [ResolvedTo(STATUS)] interface Foo : R|kotlin/Any| {
+    }
+    public final [ResolvedTo(BODY_RESOLVE)] fun test([ResolvedTo(BODY_RESOLVE)] obj: R|kotlin/Any|): R|kotlin/Unit| {
+        when () {
+            (R|<local>/obj| is R|Foo|) ->  {
+                R|/consume|(R|<local>/obj|)
+            }
+        }
+
+        R|<local>/obj|.R|kotlin/Any.hashCode|()
+    }
+    public final [ResolvedTo(CONTRACTS)] fun consume([ResolvedTo(CONTRACTS)] obj: R|Foo|): R|kotlin/Unit| {
+    }

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/afterLoop.kt

@@ -0,0 +1,17 @@
+interface Node {
+    val shouldProcess: Boolean
+    val parent: Node?
+}
+
+fun test(initial: Node?) {
+    var current = initial
+
+    while (current!!.shouldProcess) {
+        consume(current)
+        current = current.parent
+    }
+
+    <expr>consume(current)</expr>
+}
+
+fun consume(node: Node) {}

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/afterLoop.txt

@@ -0,0 +1,48 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirLocalScope
+            Properties:
+                FirValueParameterSymbol initial: R|Node?|
+    Element 8
+        Scope: FirLocalScope
+            Properties:
+                FirPropertySymbol lvar current: R|Node?|
+Smart Casts:
+    FirPropertySymbol lvar current: R|Node?|
+    Types:
+        kotlin/Any
+
+FILE: [ResolvedTo(IMPORTS)] afterLoop.kt
+    public abstract [ResolvedTo(STATUS)] interface Node : R|kotlin/Any| {
+        public abstract [ResolvedTo(CONTRACTS)] val shouldProcess: R|kotlin/Boolean|
+            public [ResolvedTo(CONTRACTS)] get(): R|kotlin/Boolean|
+
+        public abstract [ResolvedTo(CONTRACTS)] val parent: R|Node?|
+            public [ResolvedTo(CONTRACTS)] get(): R|Node?|
+
+    }
+    public final [ResolvedTo(BODY_RESOLVE)] fun test([ResolvedTo(BODY_RESOLVE)] initial: R|Node?|): R|kotlin/Unit| {
+        [ResolvedTo(BODY_RESOLVE)] lvar current: R|Node?| = R|<local>/initial|
+        while(R|<local>/current|!!.R|/Node.shouldProcess|) {
+            R|/consume|(R|<local>/current|)
+            R|<local>/current| = R|<local>/current|.R|/Node.parent|
+        }
+
+        R|/consume|(R|<local>/current|)
+    }
+    public final [ResolvedTo(CONTRACTS)] fun consume([ResolvedTo(CONTRACTS)] node: R|Node|): R|kotlin/Unit| {
+    }

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/andRight.kt

@@ -0,0 +1,11 @@
+interface Foo {
+    val isValid: Boolean
+}
+
+fun test(obj: Any) {
+    if (obj is Foo && <expr>obj.isValid</expr>) {
+        consume(obj)
+    }
+}
+
+fun consume(obj: Foo) {}

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/andRight.txt

@@ -0,0 +1,43 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirLocalScope
+            Properties:
+                FirValueParameterSymbol obj: R|kotlin/Any|
+    Element 8
+        Scope: FirLocalScope
+Smart Casts:
+    FirValueParameterSymbol obj: R|kotlin/Any|
+    Types:
+        Foo
+        kotlin/Any
+
+FILE: [ResolvedTo(IMPORTS)] andRight.kt
+    public abstract [ResolvedTo(STATUS)] interface Foo : R|kotlin/Any| {
+        public abstract [ResolvedTo(CONTRACTS)] val isValid: R|kotlin/Boolean|
+            public [ResolvedTo(CONTRACTS)] get(): R|kotlin/Boolean|
+
+    }
+    public final [ResolvedTo(BODY_RESOLVE)] fun test([ResolvedTo(BODY_RESOLVE)] obj: R|kotlin/Any|): R|kotlin/Unit| {
+        when () {
+            (R|<local>/obj| is R|Foo|) && R|<local>/obj|.R|/Foo.isValid| ->  {
+                R|/consume|(R|<local>/obj|)
+            }
+        }
+
+    }
+    public final [ResolvedTo(CONTRACTS)] fun consume([ResolvedTo(CONTRACTS)] obj: R|Foo|): R|kotlin/Unit| {
+    }

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/argument.kt

@@ -2,8 +2,8 @@ interface Foo {}
 
 fun test(obj: Any) {
     if (obj is Foo) {
-        consume(<expr>obj</expr>)
+        <expr>consume(obj)</expr>
     }
 }
 
-fun consume(obj: Any) {}
+fun consume(obj: Foo) {}

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/argument.txt

@@ -27,7 +27,7 @@ Smart Casts:
         Foo
         kotlin/Any
 
-FILE: [ResolvedTo(IMPORTS)] smartCastArgument.kt
+FILE: [ResolvedTo(IMPORTS)] argument.kt
     public abstract [ResolvedTo(STATUS)] interface Foo : R|kotlin/Any| {
     }
     public final [ResolvedTo(BODY_RESOLVE)] fun test([ResolvedTo(BODY_RESOLVE)] obj: R|kotlin/Any|): R|kotlin/Unit| {
@@ -38,5 +38,5 @@ FILE: [ResolvedTo(IMPORTS)] smartCastArgument.kt
         }
 
     }
-    public final [ResolvedTo(CONTRACTS)] fun consume([ResolvedTo(CONTRACTS)] obj: R|kotlin/Any|): R|kotlin/Unit| {
+    public final [ResolvedTo(CONTRACTS)] fun consume([ResolvedTo(CONTRACTS)] obj: R|Foo|): R|kotlin/Unit| {
     }

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/argumentAsReceiver.kt

@@ -4,7 +4,7 @@ interface Bar {}
 fun test(obj: Any) {
     if (obj is Foo) {
         if (obj is Bar) {
-            obj.consume(<expr>obj</expr>)
+            <expr>obj.consume(obj)</expr>
         }
     }
 }

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/argumentAsReceiver.txt

@@ -30,7 +30,7 @@ Smart Casts:
         kotlin/Any
         Bar
 
-FILE: [ResolvedTo(IMPORTS)] smartCastReceiverArgument.kt
+FILE: [ResolvedTo(IMPORTS)] argumentAsReceiver.kt
     public abstract [ResolvedTo(STATUS)] interface Foo : R|kotlin/Any| {
     }
     public abstract [ResolvedTo(STATUS)] interface Bar : R|kotlin/Any| {

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/beforeIf.kt

@@ -0,0 +1,10 @@
+interface Foo {}
+
+fun test(obj: Any) {
+    <expr>obj.hashCode()</expr>
+    if (obj is Foo) {
+        consume(obj)
+    }
+}
+
+fun consume(obj: Foo) {}

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/beforeIf.txt

@@ -0,0 +1,36 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirLocalScope
+            Properties:
+                FirValueParameterSymbol obj: R|kotlin/Any|
+    Element 8
+        Scope: FirLocalScope
+
+FILE: [ResolvedTo(IMPORTS)] beforeIf.kt
+    public abstract [ResolvedTo(STATUS)] interface Foo : R|kotlin/Any| {
+    }
+    public final [ResolvedTo(BODY_RESOLVE)] fun test([ResolvedTo(BODY_RESOLVE)] obj: R|kotlin/Any|): R|kotlin/Unit| {
+        R|<local>/obj|.R|kotlin/Any.hashCode|()
+        when () {
+            (R|<local>/obj| is R|Foo|) ->  {
+                R|/consume|(R|<local>/obj|)
+            }
+        }
+
+    }
+    public final [ResolvedTo(CONTRACTS)] fun consume([ResolvedTo(CONTRACTS)] obj: R|Foo|): R|kotlin/Unit| {
+    }

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/beforeLoop.kt

@@ -0,0 +1,14 @@
+interface Node {
+    val parent: Node?
+}
+
+fun test(initial: Node) { 
+    <expr>var current = initial</expr>
+    
+    while (current != null) {
+        consume(current)
+        current = current.parent
+    }
+}
+
+fun consume(node: Node) {}

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/beforeLoop.txt

@@ -0,0 +1,38 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirLocalScope
+            Properties:
+                FirValueParameterSymbol initial: R|Node|
+    Element 8
+        Scope: FirLocalScope
+
+FILE: [ResolvedTo(IMPORTS)] beforeLoop.kt
+    public abstract [ResolvedTo(STATUS)] interface Node : R|kotlin/Any| {
+        public abstract [ResolvedTo(CONTRACTS)] val parent: R|Node?|
+            public [ResolvedTo(CONTRACTS)] get(): R|Node?|
+
+    }
+    public final [ResolvedTo(BODY_RESOLVE)] fun test([ResolvedTo(BODY_RESOLVE)] initial: R|Node|): R|kotlin/Unit| {
+        [ResolvedTo(BODY_RESOLVE)] lvar current: R|Node| = R|<local>/initial|
+        while(!=(R|<local>/current|, Null(null))) {
+            R|/consume|(R|<local>/current|)
+            R|<local>/current| = R|<local>/current|.R|/Node.parent|
+        }
+
+    }
+    public final [ResolvedTo(CONTRACTS)] fun consume([ResolvedTo(CONTRACTS)] node: R|Node|): R|kotlin/Unit| {
+    }

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/dispatchReceiver.kt

@@ -0,0 +1,11 @@
+open class Foo {
+    fun foo() {
+        if (this is Bar) {
+            <expr>consume(this)</expr>
+        }
+    }
+}
+
+class Bar : Foo()
+
+fun consume(obj: Bar) {}

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/dispatchReceiver.txt

@@ -0,0 +1,55 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Implicit receiver:
+            FirRegularClassSymbol public open class Foo : R|kotlin/Any|
+                Type: Bar
+    Element 8
+        Scope: FirLocalScope
+    Element 9
+        Scope: FirLocalScope
+    Element 10
+        Scope: FirLocalScope
+Smart Casts:
+    FirRegularClassSymbol public open class Foo : R|kotlin/Any|
+    Types:
+        Bar
+        kotlin/Any
+
+FILE: [ResolvedTo(IMPORTS)] dispatchReceiver.kt
+    public open [ResolvedTo(STATUS)] class Foo : R|kotlin/Any| {
+        public [ResolvedTo(STATUS)] constructor(): R|Foo| {
+            LAZY_super<R|kotlin/Any|>
+        }
+
+        public final [ResolvedTo(BODY_RESOLVE)] fun foo(): R|kotlin/Unit| {
+            when () {
+                (this@R|/Foo| is R|Bar|) ->  {
+                    R|/consume|(this@R|/Foo|)
+                }
+            }
+
+        }
+
+    }
+    public final [ResolvedTo(STATUS)] class Bar : R|Foo| {
+        public [ResolvedTo(STATUS)] constructor(): R|Bar| {
+            LAZY_super<R|Foo|>
+        }
+
+    }
+    public final [ResolvedTo(CONTRACTS)] fun consume([ResolvedTo(CONTRACTS)] obj: R|Bar|): R|kotlin/Unit| {
+    }

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/extensionReceiver.kt

@@ -0,0 +1,9 @@
+interface Foo {}
+
+fun Any.test() {
+    if (this is Foo) {
+        <expr>consume(this)</expr>
+    }
+}
+
+fun consume(obj: Foo) {}

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/extensionReceiver.txt

@@ -0,0 +1,44 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirLocalScope
+    Element 8
+        Implicit receiver:
+            FirNamedFunctionSymbol public final fun R|kotlin/Any|.test(): R|kotlin/Unit|
+                Type: Foo
+    Element 9
+        Scope: FirLocalScope
+    Element 10
+        Scope: FirLocalScope
+Smart Casts:
+    FirNamedFunctionSymbol public final fun R|kotlin/Any|.test(): R|kotlin/Unit|
+    Types:
+        Foo
+        kotlin/Any
+
+FILE: [ResolvedTo(IMPORTS)] extensionReceiver.kt
+    public abstract [ResolvedTo(STATUS)] interface Foo : R|kotlin/Any| {
+    }
+    public final [ResolvedTo(BODY_RESOLVE)] fun R|kotlin/Any|.test(): R|kotlin/Unit| {
+        when () {
+            (this@R|/test| is R|Foo|) ->  {
+                R|/consume|(this@R|/test|)
+            }
+        }
+
+    }
+    public final [ResolvedTo(CONTRACTS)] fun consume([ResolvedTo(CONTRACTS)] obj: R|Foo|): R|kotlin/Unit| {
+    }

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/insideLoop.kt

@@ -0,0 +1,14 @@
+interface Node {
+    val parent: Node?
+}
+
+fun test(initial: Node) { 
+    var current = initial
+    
+    while (current != null) {
+        <expr>consume(current)</expr>
+        current = current.parent
+    }
+}
+
+fun consume(node: Node) {}

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/insideLoop.txt

@@ -0,0 +1,46 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirLocalScope
+            Properties:
+                FirValueParameterSymbol initial: R|Node|
+    Element 8
+        Scope: FirLocalScope
+            Properties:
+                FirPropertySymbol lvar current: R|Node|
+    Element 9
+        Scope: FirLocalScope
+Smart Casts:
+    FirPropertySymbol lvar current: R|Node|
+    Types:
+        kotlin/Any
+
+FILE: [ResolvedTo(IMPORTS)] insideLoop.kt
+    public abstract [ResolvedTo(STATUS)] interface Node : R|kotlin/Any| {
+        public abstract [ResolvedTo(CONTRACTS)] val parent: R|Node?|
+            public [ResolvedTo(CONTRACTS)] get(): R|Node?|
+
+    }
+    public final [ResolvedTo(BODY_RESOLVE)] fun test([ResolvedTo(BODY_RESOLVE)] initial: R|Node|): R|kotlin/Unit| {
+        [ResolvedTo(BODY_RESOLVE)] lvar current: R|Node| = R|<local>/initial|
+        while(!=(R|<local>/current|, Null(null))) {
+            R|/consume|(R|<local>/current|)
+            R|<local>/current| = R|<local>/current|.R|/Node.parent|
+        }
+
+    }
+    public final [ResolvedTo(CONTRACTS)] fun consume([ResolvedTo(CONTRACTS)] node: R|Node|): R|kotlin/Unit| {
+    }

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/orRight.kt

@@ -0,0 +1,13 @@
+interface Foo {
+    val isValid: Boolean
+}
+
+fun test(obj: Any) {
+    if (obj is Foo || <expr>isValid(obj)</expr>) {
+        consume(obj)
+    }
+}
+
+fun isValid(obj: Any): Boolean = true
+
+fun consume(obj: Foo) {}

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/orRight.txt

@@ -0,0 +1,41 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirLocalScope
+            Properties:
+                FirValueParameterSymbol obj: R|kotlin/Any|
+    Element 8
+        Scope: FirLocalScope
+
+FILE: [ResolvedTo(IMPORTS)] orRight.kt
+    public? final? [ResolvedTo(RAW_FIR)] interface Foo : R|kotlin/Any| {
+        public? final? [ResolvedTo(RAW_FIR)] val isValid: Boolean
+            public? [ResolvedTo(RAW_FIR)] get(): Boolean
+
+    }
+    public final [ResolvedTo(BODY_RESOLVE)] fun test([ResolvedTo(BODY_RESOLVE)] obj: R|kotlin/Any|): R|kotlin/Unit| {
+        when () {
+            (R|<local>/obj| is R|Foo|) || R|/isValid|(R|<local>/obj|) ->  {
+                R|/consume<Inapplicable(INAPPLICABLE): /consume>#|(R|<local>/obj|)
+            }
+        }
+
+    }
+    public final [ResolvedTo(CONTRACTS)] fun isValid([ResolvedTo(CONTRACTS)] obj: R|kotlin/Any|): R|kotlin/Boolean| {
+        ^isValid Boolean(true)
+    }
+    public final [ResolvedTo(CONTRACTS)] fun consume([ResolvedTo(CONTRACTS)] obj: R|Foo|): R|kotlin/Unit| {
+    }

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/plainCheck.kt

@@ -0,0 +1,9 @@
+interface Foo {}
+
+fun test(obj: Any, another: Foo) {
+    if (obj is Foo) {
+        <expr>consume(another)</expr>
+    }
+}
+
+fun consume(obj: Foo) {}

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/plainCheck.txt

@@ -0,0 +1,43 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirLocalScope
+            Properties:
+                FirValueParameterSymbol another: R|Foo|
+                FirValueParameterSymbol obj: R|kotlin/Any|
+    Element 8
+        Scope: FirLocalScope
+    Element 9
+        Scope: FirLocalScope
+Smart Casts:
+    FirValueParameterSymbol obj: R|kotlin/Any|
+    Types:
+        Foo
+        kotlin/Any
+
+FILE: [ResolvedTo(IMPORTS)] plainCheck.kt
+    public abstract [ResolvedTo(STATUS)] interface Foo : R|kotlin/Any| {
+    }
+    public final [ResolvedTo(BODY_RESOLVE)] fun test([ResolvedTo(BODY_RESOLVE)] obj: R|kotlin/Any|, [ResolvedTo(BODY_RESOLVE)] another: R|Foo|): R|kotlin/Unit| {
+        when () {
+            (R|<local>/obj| is R|Foo|) ->  {
+                R|/consume|(R|<local>/another|)
+            }
+        }
+
+    }
+    public final [ResolvedTo(CONTRACTS)] fun consume([ResolvedTo(CONTRACTS)] obj: R|Foo|): R|kotlin/Unit| {
+    }

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/require.kt

@@ -0,0 +1,10 @@
+// WITH_STDLIB
+
+interface Foo
+
+fun test(obj: Any) {
+    require(obj is Foo)
+    <expr>consume(obj)</expr>
+}
+
+fun consume(obj: Any) {}

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/require.txt

@@ -0,0 +1,36 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirLocalScope
+            Properties:
+                FirValueParameterSymbol obj: R|kotlin/Any|
+    Element 8
+        Scope: FirLocalScope
+Smart Casts:
+    FirValueParameterSymbol obj: R|kotlin/Any|
+    Types:
+        Foo
+        kotlin/Any
+
+FILE: [ResolvedTo(IMPORTS)] require.kt
+    public abstract [ResolvedTo(STATUS)] interface Foo : R|kotlin/Any| {
+    }
+    public final [ResolvedTo(BODY_RESOLVE)] fun test([ResolvedTo(BODY_RESOLVE)] obj: R|kotlin/Any|): R|kotlin/Unit| {
+        R|kotlin/require|((R|<local>/obj| is R|Foo|))
+        R|/consume|(R|<local>/obj|)
+    }
+    public final [ResolvedTo(CONTRACTS)] fun consume([ResolvedTo(CONTRACTS)] obj: R|kotlin/Any|): R|kotlin/Unit| {
+    }

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/when.kt

@@ -0,0 +1,14 @@
+interface Foo {
+    fun foo()
+}
+
+interface Bar {
+    fun bar()
+}
+
+fun test(obj: Any) {
+    when (obj) {
+        is Foo -> obj.foo()
+        is Bar -> <expr>obj.bar()</expr>
+    }
+}

analysis/low-level-api-fir/testdata/contextCollector/smartCasts/when.txt

@@ -0,0 +1,49 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirLocalScope
+            Properties:
+                FirValueParameterSymbol obj: R|kotlin/Any|
+    Element 8
+        Scope: FirLocalScope
+    Element 9
+        Scope: FirLocalScope
+Smart Casts:
+    FirValueParameterSymbol obj: R|kotlin/Any|
+    Types:
+        Bar
+        kotlin/Any
+
+FILE: [ResolvedTo(IMPORTS)] when.kt
+    public abstract [ResolvedTo(STATUS)] interface Foo : R|kotlin/Any| {
+        public abstract [ResolvedTo(CONTRACTS)] fun foo(): R|kotlin/Unit|
+
+    }
+    public abstract [ResolvedTo(STATUS)] interface Bar : R|kotlin/Any| {
+        public abstract [ResolvedTo(CONTRACTS)] fun bar(): R|kotlin/Unit|
+
+    }
+    public final [ResolvedTo(BODY_RESOLVE)] fun test([ResolvedTo(BODY_RESOLVE)] obj: R|kotlin/Any|): R|kotlin/Unit| {
+        when (R|<local>/obj|) {
+            ($subj$ is R|Foo|) ->  {
+                R|<local>/obj|.R|/Foo.foo|()
+            }
+            ($subj$ is R|Bar|) ->  {
+                R|<local>/obj|.R|/Bar.bar|()
+            }
+        }
+
+    }

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractContextCollectorTest.kt

@@ -16,6 +16,8 @@ import org.jetbrains.kotlin.fir.declarations.FirResolvedImport
 import org.jetbrains.kotlin.fir.declarations.FirTowerDataContext
 import org.jetbrains.kotlin.fir.renderer.FirRenderer
 import org.jetbrains.kotlin.fir.resolve.SessionHolderImpl
+import org.jetbrains.kotlin.fir.resolve.dfa.Identifier
+import org.jetbrains.kotlin.fir.resolve.dfa.RealVariable
 import org.jetbrains.kotlin.fir.scopes.*
 import org.jetbrains.kotlin.fir.scopes.impl.*
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
@@ -180,14 +182,14 @@ internal object ElementContextRenderer {
         }
     }
 
-    private fun StringBuilder.renderSmartCasts(smartCasts: Map<FirBasedSymbol<*>, Set<ConeKotlinType>>) {
+    private fun StringBuilder.renderSmartCasts(smartCasts: Map<RealVariable, Set<ConeKotlinType>>) {
         if (smartCasts.isEmpty()) {
             return
         }
 
         appendBlock("Smart Casts:") {
-            for ((symbol, types) in smartCasts) {
+            for ((realVariable, types) in smartCasts) {
-                appendSymbol(symbol).appendLine()
+                appendSymbol(realVariable.identifier.symbol).appendLine()
 
                 appendBlock("Types:") {
                     for (type in types) {

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ContextCollectorTestGenerated.java

@@ -72,6 +72,12 @@ public class ContextCollectorTestGenerated extends AbstractContextCollectorTest
         runTest("analysis/low-level-api-fir/testdata/contextCollector/lambdaArguments.kt");
     }
 
+    @Test
+    @TestMetadata("localClass.kt")
+    public void testLocalClass() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/contextCollector/localClass.kt");
+    }
+
     @Test
     @TestMetadata("nestedClasses.kt")
     public void testNestedClasses() throws Exception {
@@ -90,15 +96,97 @@ public class ContextCollectorTestGenerated extends AbstractContextCollectorTest
         runTest("analysis/low-level-api-fir/testdata/contextCollector/simple.kt");
     }
 
-    @Test
+    @Nested
-    @TestMetadata("smartCastArgument.kt")
+    @TestMetadata("analysis/low-level-api-fir/testdata/contextCollector/smartCasts")
-    public void testSmartCastArgument() throws Exception {
+    @TestDataPath("$PROJECT_ROOT")
-        runTest("analysis/low-level-api-fir/testdata/contextCollector/smartCastArgument.kt");
+    public class SmartCasts {
-    }
+        @Test
+        @TestMetadata("afterIf.kt")
+        public void testAfterIf() throws Exception {
+            runTest("analysis/low-level-api-fir/testdata/contextCollector/smartCasts/afterIf.kt");
+        }
 
-    @Test
+        @Test
-    @TestMetadata("smartCastReceiverArgument.kt")
+        @TestMetadata("afterLoop.kt")
-    public void testSmartCastReceiverArgument() throws Exception {
+        public void testAfterLoop() throws Exception {
-        runTest("analysis/low-level-api-fir/testdata/contextCollector/smartCastReceiverArgument.kt");
+            runTest("analysis/low-level-api-fir/testdata/contextCollector/smartCasts/afterLoop.kt");
+        }
+
+        @Test
+        public void testAllFilesPresentInSmartCasts() throws Exception {
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/low-level-api-fir/testdata/contextCollector/smartCasts"), Pattern.compile("^(.+)\\.kt$"), null, true);
+        }
+
+        @Test
+        @TestMetadata("andRight.kt")
+        public void testAndRight() throws Exception {
+            runTest("analysis/low-level-api-fir/testdata/contextCollector/smartCasts/andRight.kt");
+        }
+
+        @Test
+        @TestMetadata("argument.kt")
+        public void testArgument() throws Exception {
+            runTest("analysis/low-level-api-fir/testdata/contextCollector/smartCasts/argument.kt");
+        }
+
+        @Test
+        @TestMetadata("argumentAsReceiver.kt")
+        public void testArgumentAsReceiver() throws Exception {
+            runTest("analysis/low-level-api-fir/testdata/contextCollector/smartCasts/argumentAsReceiver.kt");
+        }
+
+        @Test
+        @TestMetadata("beforeIf.kt")
+        public void testBeforeIf() throws Exception {
+            runTest("analysis/low-level-api-fir/testdata/contextCollector/smartCasts/beforeIf.kt");
+        }
+
+        @Test
+        @TestMetadata("beforeLoop.kt")
+        public void testBeforeLoop() throws Exception {
+            runTest("analysis/low-level-api-fir/testdata/contextCollector/smartCasts/beforeLoop.kt");
+        }
+
+        @Test
+        @TestMetadata("dispatchReceiver.kt")
+        public void testDispatchReceiver() throws Exception {
+            runTest("analysis/low-level-api-fir/testdata/contextCollector/smartCasts/dispatchReceiver.kt");
+        }
+
+        @Test
+        @TestMetadata("extensionReceiver.kt")
+        public void testExtensionReceiver() throws Exception {
+            runTest("analysis/low-level-api-fir/testdata/contextCollector/smartCasts/extensionReceiver.kt");
+        }
+
+        @Test
+        @TestMetadata("insideLoop.kt")
+        public void testInsideLoop() throws Exception {
+            runTest("analysis/low-level-api-fir/testdata/contextCollector/smartCasts/insideLoop.kt");
+        }
+
+        @Test
+        @TestMetadata("orRight.kt")
+        public void testOrRight() throws Exception {
+            runTest("analysis/low-level-api-fir/testdata/contextCollector/smartCasts/orRight.kt");
+        }
+
+        @Test
+        @TestMetadata("plainCheck.kt")
+        public void testPlainCheck() throws Exception {
+            runTest("analysis/low-level-api-fir/testdata/contextCollector/smartCasts/plainCheck.kt");
+        }
+
+        @Test
+        @TestMetadata("require.kt")
+        public void testRequire() throws Exception {
+            runTest("analysis/low-level-api-fir/testdata/contextCollector/smartCasts/require.kt");
+        }
+
+        @Test
+        @TestMetadata("when.kt")
+        public void testWhen() throws Exception {
+            runTest("analysis/low-level-api-fir/testdata/contextCollector/smartCasts/when.kt");
+        }
     }
 }

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/PersistentImplicitReceiverStack.kt

@@ -98,6 +98,10 @@ class PersistentImplicitReceiverStack private constructor(
         return originalTypes[index]
     }
 
+    fun getType(index: Int): ConeKotlinType {
+        return stack[index].type
+    }
+
     // This method is only used from DFA and it's in some sense breaks persistence contracts of the data structure
     // But it's ok since DFA handles everything properly yet, but still may be it should be rewritten somehow
     @OptIn(ImplicitReceiverValue.ImplicitReceiverInternals::class)