FIR IDE: collect snapshot of FirTowerDataContext for statements

ImplicitReceiverValue is mutable and FIR body resolve could alter it
while analysing code with smartcast. Hence, previously the IDE may see
inconsistent receiver values for a local scope. For example

```
open class A
interface Foo {
    fun foo()
}
fun A.bar() {
    if (this is Foo) {
        // scope here has implicit receiver type to be `A` rather than `it<A, Foo>`
    }
}
```

This change creates snapshots for local statements so later changes
during body resolve won't affect the collected context.

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.resolve
 
 import kotlinx.collections.immutable.PersistentList
 import kotlinx.collections.immutable.persistentListOf
+import kotlinx.collections.immutable.toPersistentList
 import org.jetbrains.kotlin.fir.FirCallResolver
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
@@ -137,6 +138,16 @@ class FirTowerDataContext private constructor(
             nonLocalTowerDataElements.add(element)
         )
     }
+
+    fun createSnapshot(): FirTowerDataContext {
+        return FirTowerDataContext(
+            towerDataElements.map { FirTowerDataElement(it.scope, it.implicitReceiver?.createSnapshot(), it.isLocal) }.toPersistentList(),
+            implicitReceiverStack.createSnapshot(),
+            localScopes.toPersistentList(),
+            nonLocalTowerDataElements.map { FirTowerDataElement(it.scope, it.implicitReceiver?.createSnapshot(), it.isLocal) }
+                .toPersistentList()
+        )
+    }
 }
 
 class FirTowerDataElement(val scope: FirScope?, val implicitReceiver: ImplicitReceiverValue<*>?, val isLocal: Boolean)

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

@@ -5,10 +5,7 @@
 
 package org.jetbrains.kotlin.fir.resolve
 
-import kotlinx.collections.immutable.PersistentList
-import kotlinx.collections.immutable.PersistentMap
-import kotlinx.collections.immutable.persistentListOf
-import kotlinx.collections.immutable.persistentMapOf
+import kotlinx.collections.immutable.*
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitDispatchReceiverValue
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
@@ -84,4 +81,13 @@ class PersistentImplicitReceiverStack private constructor(
         assert(index >= 0 && index < stack.size)
         stack[index].replaceType(type)
     }
+
+    fun createSnapshot(): PersistentImplicitReceiverStack {
+        return PersistentImplicitReceiverStack(
+            stack.map { it.createSnapshot() }.toPersistentList(),
+            indexesPerLabel,
+            indexesPerSymbol,
+            originalTypes
+        )
+    }
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/FirReceivers.kt

@@ -76,7 +76,8 @@ sealed class ImplicitReceiverValue<S : FirBasedSymbol<*>>(
     val boundSymbol: S,
     type: ConeKotlinType,
     protected val useSiteSession: FirSession,
-    protected val scopeSession: ScopeSession
+    protected val scopeSession: ScopeSession,
+    private val mutable: Boolean,
 ) : ReceiverValue {
     final override var type: ConeKotlinType = type
         private set
@@ -96,6 +97,7 @@ sealed class ImplicitReceiverValue<S : FirBasedSymbol<*>>(
      * Should be called only in ImplicitReceiverStack
      */
     internal fun replaceType(type: ConeKotlinType) {
+        if (!mutable) throw IllegalStateException("Cannot mutate an immutable ImplicitReceiverValue")
         if (type == this.type) return
         this.type = type
         receiverExpression = if (type == originalReceiverExpression.typeRef.coneType) {
@@ -110,6 +112,8 @@ sealed class ImplicitReceiverValue<S : FirBasedSymbol<*>>(
         }
         implicitScope = type.scope(useSiteSession, scopeSession, FakeOverrideTypeCalculator.DoNothing)
     }
+
+    abstract fun createSnapshot(): ImplicitReceiverValue<S>
 }
 
 private fun receiverExpression(symbol: FirBasedSymbol<*>, type: ConeKotlinType): FirThisReceiverExpression =
@@ -129,27 +133,42 @@ class ImplicitDispatchReceiverValue internal constructor(
     boundSymbol: FirClassSymbol<*>,
     type: ConeKotlinType,
     useSiteSession: FirSession,
-    scopeSession: ScopeSession
-) : ImplicitReceiverValue<FirClassSymbol<*>>(boundSymbol, type, useSiteSession, scopeSession) {
+    scopeSession: ScopeSession,
+    mutable: Boolean = true,
+) : ImplicitReceiverValue<FirClassSymbol<*>>(boundSymbol, type, useSiteSession, scopeSession, mutable) {
     internal constructor(
         boundSymbol: FirClassSymbol<*>, useSiteSession: FirSession, scopeSession: ScopeSession
     ) : this(
         boundSymbol, boundSymbol.constructType(typeArguments = emptyArray(), isNullable = false),
         useSiteSession, scopeSession
     )
+
+    override fun createSnapshot(): ImplicitReceiverValue<FirClassSymbol<*>> {
+        return ImplicitDispatchReceiverValue(boundSymbol, type, useSiteSession, scopeSession, false)
+    }
 }
 
 class ImplicitExtensionReceiverValue(
     boundSymbol: FirCallableSymbol<*>,
     type: ConeKotlinType,
     useSiteSession: FirSession,
-    scopeSession: ScopeSession
-) : ImplicitReceiverValue<FirCallableSymbol<*>>(boundSymbol, type, useSiteSession, scopeSession)
+    scopeSession: ScopeSession,
+    mutable: Boolean = true,
+) : ImplicitReceiverValue<FirCallableSymbol<*>>(boundSymbol, type, useSiteSession, scopeSession, mutable) {
+    override fun createSnapshot(): ImplicitReceiverValue<FirCallableSymbol<*>> {
+        return ImplicitExtensionReceiverValue(boundSymbol, type, useSiteSession, scopeSession, false)
+    }
+}
 
 
 class InaccessibleImplicitReceiverValue(
     boundSymbol: FirClassSymbol<*>,
     type: ConeKotlinType,
     useSiteSession: FirSession,
-    scopeSession: ScopeSession
-) : ImplicitReceiverValue<FirClassSymbol<*>>(boundSymbol, type, useSiteSession, scopeSession)
+    scopeSession: ScopeSession,
+    mutable: Boolean = true,
+) : ImplicitReceiverValue<FirClassSymbol<*>>(boundSymbol, type, useSiteSession, scopeSession, mutable) {
+    override fun createSnapshot(): ImplicitReceiverValue<FirClassSymbol<*>> {
+        return InaccessibleImplicitReceiverValue(boundSymbol, type, useSiteSession, scopeSession, false)
+    }
+}

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirTowerDataContextCollector.kt

@@ -37,7 +37,8 @@ internal class FirTowerDataContextAllElementsCollector : FirTowerDataContextColl
 
     override fun addStatementContext(statement: FirStatement, context: FirTowerDataContext) {
         val closestStatementInBlock = statement.psi?.closestBlockLevelOrInitializerExpression() ?: return
-        elementsToContext[closestStatementInBlock] = context
+        // FIR body transform may alter the context if there implicit receivers with smartcast
+        elementsToContext[closestStatementInBlock] = context.createSnapshot()
     }
 
     override fun addDeclarationContext(declaration: FirDeclaration, context: FirTowerDataContext) {