[Analysis API] Keep the collected 'FirTowerDataContext' mutable

'ContextCollector' is used for computing context of 'FirCodeFragment's.
Code fragments themselves might contain additional smart cast operations
that modify the context receiver stack.

^KT-63056 Fixed

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisLibrarySourceModuleCompilerFacilityTestGenerated.java

@@ -196,6 +196,12 @@ public class FirIdeNormalAnalysisLibrarySourceModuleCompilerFacilityTestGenerate
                 runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/extensionReceiverPropertyLabeled.kt");
             }
 
+            @Test
+            @TestMetadata("extensionReceiverSmartCasted.kt")
+            public void testExtensionReceiverSmartCasted() throws Exception {
+                runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/extensionReceiverSmartCasted.kt");
+            }
+
             @Test
             @TestMetadata("initializer.kt")
             public void testInitializer() throws Exception {

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisSourceModuleCompilerFacilityTestGenerated.java

@@ -196,6 +196,12 @@ public class FirIdeNormalAnalysisSourceModuleCompilerFacilityTestGenerated exten
                 runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/extensionReceiverPropertyLabeled.kt");
             }
 
+            @Test
+            @TestMetadata("extensionReceiverSmartCasted.kt")
+            public void testExtensionReceiverSmartCasted() throws Exception {
+                runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/extensionReceiverSmartCasted.kt");
+            }
+
             @Test
             @TestMetadata("initializer.kt")
             public void testInitializer() throws Exception {

analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/extensionReceiverSmartCasted.capturing.txt

@@ -0,0 +1,3 @@
+ExtensionReceiver[name: apply; isMutated: false; displayText: this@apply]
+    apply@fun R|Foo|.<anonymous>(): R|kotlin/Unit| <inline=Inline, kind=EXACTLY_ONCE>
+    R|Foo|

analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/extensionReceiverSmartCasted.fragment.kt

@@ -0,0 +1,2 @@
+this as FooImpl
+n

analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/extensionReceiverSmartCasted.ir.txt

@@ -0,0 +1,16 @@
+MODULE_FRAGMENT
+  FILE fqName:<root> fileName:fragment.kt
+    CLASS CLASS name:CodeFragment modality:FINAL visibility:public superTypes:[kotlin.Any]
+      $this: VALUE_PARAMETER INSTANCE_RECEIVER name:<this> type:<root>.CodeFragment
+      CONSTRUCTOR visibility:public <> () returnType:<root>.CodeFragment [primary]
+        BLOCK_BODY
+          DELEGATING_CONSTRUCTOR_CALL 'public constructor <init> () [primary] declared in kotlin.Any'
+      FUN name:run visibility:public modality:FINAL <> (p0:<root>.Foo) returnType:kotlin.Int
+        VALUE_PARAMETER name:p0 index:0 type:<root>.Foo
+        EXPRESSION_BODY
+          BLOCK type=kotlin.Int origin=null
+            TYPE_OP type=<root>.FooImpl origin=CAST typeOperand=<root>.FooImpl
+              GET_VAR 'p0: <root>.Foo declared in <root>.CodeFragment.run' type=<root>.Foo origin=null
+            CALL 'public final fun <get-n> (): kotlin.Int declared in <root>.FooImpl' type=kotlin.Int origin=GET_PROPERTY
+              $this: TYPE_OP type=<root>.FooImpl origin=IMPLICIT_CAST typeOperand=<root>.FooImpl
+                GET_VAR 'p0: <root>.Foo declared in <root>.CodeFragment.run' type=<root>.Foo origin=null

analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/extensionReceiverSmartCasted.kt

@@ -0,0 +1,13 @@
+interface Foo
+
+class FooImpl : Foo {
+    val n: Int = 5
+}
+
+fun makeFoo(): Foo = FooImpl()
+
+fun main() {
+    makeFoo().apply {
+        <caret>Unit
+    }
+}

analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/extensionReceiverSmartCasted.txt

@@ -0,0 +1,5 @@
+public final class CodeFragment {
+    // source: 'fragment.kt'
+    public method <init>(): void
+    public final static method run(p0: Foo): int
+}

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

@@ -51,7 +51,7 @@ internal class FirTowerDataContextAllElementsCollector : FirResolveContextCollec
     override fun addStatementContext(statement: FirStatement, context: BodyResolveContext) {
         val closestParentExpression = statement.psi?.closestParentExpressionWithSameContextOrSelf() ?: return
         // FIR body transform may alter the context if there are implicit receivers with smartcast
-        elementsToContext[closestParentExpression] = context.towerDataContext.createSnapshot()
+        elementsToContext[closestParentExpression] = context.towerDataContext.createSnapshot(keepMutable = false)
     }
 
     override fun addDeclarationContext(declaration: FirDeclaration, context: BodyResolveContext) {
@@ -62,7 +62,7 @@ internal class FirTowerDataContextAllElementsCollector : FirResolveContextCollec
         // we do not collect contexts for such declarations
         if (psi is KtDelegatedSuperTypeEntry) return
 
-        elementsToContext[psi] = context.towerDataContext.createSnapshot()
+        elementsToContext[psi] = context.towerDataContext.createSnapshot(keepMutable = false)
     }
 
     override fun addClassHeaderContext(declaration: FirRegularClass, context: FirTowerDataContext) {

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

@@ -264,7 +264,7 @@ private class ContextCollectorVisitor(
             }
         }
 
-        val towerDataContextSnapshot = context.towerDataContext.createSnapshot()
+        val towerDataContextSnapshot = context.towerDataContext.createSnapshot(keepMutable = true)
 
         for ((index, oldType) in oldReceiverTypes) {
             implicitReceiverStack.replaceReceiverType(index, oldType)

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

@@ -72,6 +72,12 @@ public class CodeFragmentCapturingTestGenerated extends AbstractCodeFragmentCapt
         runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/extensionReceiverPropertyLabeled.kt");
     }
 
+    @Test
+    @TestMetadata("extensionReceiverSmartCasted.kt")
+    public void testExtensionReceiverSmartCasted() throws Exception {
+        runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/capturing/extensionReceiverSmartCasted.kt");
+    }
+
     @Test
     @TestMetadata("initializer.kt")
     public void testInitializer() throws Exception {

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

@@ -176,7 +176,7 @@ sealed class ImplicitReceiverValue<S : FirBasedSymbol<*>>(
         )
     }
 
-    abstract fun createSnapshot(): ImplicitReceiverValue<S>
+    abstract fun createSnapshot(keepMutable: Boolean): ImplicitReceiverValue<S>
 }
 
 private fun receiverExpression(
@@ -219,8 +219,8 @@ class ImplicitDispatchReceiverValue(
         useSiteSession, scopeSession
     )
 
-    override fun createSnapshot(): ImplicitReceiverValue<FirClassSymbol<*>> {
+    override fun createSnapshot(keepMutable: Boolean): ImplicitReceiverValue<FirClassSymbol<*>> {
-        return ImplicitDispatchReceiverValue(boundSymbol, type, useSiteSession, scopeSession, false)
+        return ImplicitDispatchReceiverValue(boundSymbol, type, useSiteSession, scopeSession, keepMutable)
     }
 
     override val isContextReceiver: Boolean
@@ -234,8 +234,8 @@ class ImplicitExtensionReceiverValue(
     scopeSession: ScopeSession,
     mutable: Boolean = true,
 ) : ImplicitReceiverValue<FirCallableSymbol<*>>(boundSymbol, type, useSiteSession, scopeSession, mutable) {
-    override fun createSnapshot(): ImplicitReceiverValue<FirCallableSymbol<*>> {
+    override fun createSnapshot(keepMutable: Boolean): ImplicitReceiverValue<FirCallableSymbol<*>> {
-        return ImplicitExtensionReceiverValue(boundSymbol, type, useSiteSession, scopeSession, false)
+        return ImplicitExtensionReceiverValue(boundSymbol, type, useSiteSession, scopeSession, keepMutable)
     }
 
     override val isContextReceiver: Boolean
@@ -250,8 +250,8 @@ class InaccessibleImplicitReceiverValue(
     scopeSession: ScopeSession,
     mutable: Boolean = true,
 ) : ImplicitReceiverValue<FirClassSymbol<*>>(boundSymbol, type, useSiteSession, scopeSession, mutable, inaccessibleReceiver = true) {
-    override fun createSnapshot(): ImplicitReceiverValue<FirClassSymbol<*>> {
+    override fun createSnapshot(keepMutable: Boolean): ImplicitReceiverValue<FirClassSymbol<*>> {
-        return InaccessibleImplicitReceiverValue(boundSymbol, type, useSiteSession, scopeSession, false)
+        return InaccessibleImplicitReceiverValue(boundSymbol, type, useSiteSession, scopeSession, keepMutable)
     }
 
     override val isContextReceiver: Boolean
@@ -269,7 +269,7 @@ sealed class ContextReceiverValue<S : FirBasedSymbol<*>>(
 ) : ImplicitReceiverValue<S>(
     boundSymbol, type, useSiteSession, scopeSession, mutable, contextReceiverNumber,
 ) {
-    abstract override fun createSnapshot(): ContextReceiverValue<S>
+    abstract override fun createSnapshot(keepMutable: Boolean): ContextReceiverValue<S>
 
     override val isContextReceiver: Boolean
         get() = true
@@ -286,8 +286,8 @@ class ContextReceiverValueForCallable(
 ) : ContextReceiverValue<FirCallableSymbol<*>>(
     boundSymbol, type, labelName, useSiteSession, scopeSession, mutable, contextReceiverNumber
 ) {
-    override fun createSnapshot(): ContextReceiverValue<FirCallableSymbol<*>> =
+    override fun createSnapshot(keepMutable: Boolean): ContextReceiverValue<FirCallableSymbol<*>> =
-        ContextReceiverValueForCallable(boundSymbol, type, labelName, useSiteSession, scopeSession, mutable = false, contextReceiverNumber)
+        ContextReceiverValueForCallable(boundSymbol, type, labelName, useSiteSession, scopeSession, keepMutable, contextReceiverNumber)
 }
 
 class ContextReceiverValueForClass(
@@ -301,8 +301,8 @@ class ContextReceiverValueForClass(
 ) : ContextReceiverValue<FirClassSymbol<*>>(
     boundSymbol, type, labelName, useSiteSession, scopeSession, mutable, contextReceiverNumber
 ) {
-    override fun createSnapshot(): ContextReceiverValue<FirClassSymbol<*>> =
+    override fun createSnapshot(keepMutable: Boolean): ContextReceiverValue<FirClassSymbol<*>> =
-        ContextReceiverValueForClass(boundSymbol, type, labelName, useSiteSession, scopeSession, mutable = false, contextReceiverNumber)
+        ContextReceiverValueForClass(boundSymbol, type, labelName, useSiteSession, scopeSession, keepMutable, contextReceiverNumber)
 }
 
 class ImplicitReceiverValueForScript(
@@ -316,7 +316,7 @@ class ImplicitReceiverValueForScript(
 ) : ContextReceiverValue<FirScriptSymbol>(
     boundSymbol, type, labelName, useSiteSession, scopeSession, mutable, contextReceiverNumber
 ) {
-    override fun createSnapshot(): ContextReceiverValue<FirScriptSymbol> =
+    override fun createSnapshot(keepMutable: Boolean): ContextReceiverValue<FirScriptSymbol> =
-        ImplicitReceiverValueForScript(boundSymbol, type, labelName, useSiteSession, scopeSession, mutable = false, contextReceiverNumber)
+        ImplicitReceiverValueForScript(boundSymbol, type, labelName, useSiteSession, scopeSession, keepMutable, contextReceiverNumber)
 }
 

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

@@ -949,7 +949,10 @@ class BodyResolveContext(
 
     @OptIn(PrivateForInline::class)
     fun storeCallableReferenceContext(callableReferenceAccess: FirCallableReferenceAccess) {
-        specialTowerDataContexts.storeCallableReferenceContext(callableReferenceAccess, towerDataContext.createSnapshot())
+        specialTowerDataContexts.storeCallableReferenceContext(
+            callableReferenceAccess,
+            towerDataContext.createSnapshot(keepMutable = false)
+        )
     }
 
     @OptIn(PrivateForInline::class)

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/declarations/ImplicitReceiverUtils.kt

@@ -236,12 +236,12 @@ class FirTowerDataContext private constructor(
         )
     }
 
-    fun createSnapshot(): FirTowerDataContext {
+    fun createSnapshot(keepMutable: Boolean): FirTowerDataContext {
         return FirTowerDataContext(
-            towerDataElements.map(FirTowerDataElement::createSnapshot).toPersistentList(),
+            towerDataElements.map { it.createSnapshot(keepMutable) }.toPersistentList(),
-            implicitReceiverStack.createSnapshot(),
+            implicitReceiverStack.createSnapshot(keepMutable),
             localScopes.toPersistentList(),
-            nonLocalTowerDataElements.map(FirTowerDataElement::createSnapshot).toPersistentList()
+            nonLocalTowerDataElements.map { it.createSnapshot(keepMutable) }.toPersistentList()
         )
     }
 }
@@ -254,11 +254,11 @@ class FirTowerDataElement(
     val isLocal: Boolean,
     val staticScopeOwnerSymbol: FirRegularClassSymbol? = null
 ) {
-    fun createSnapshot(): FirTowerDataElement =
+    fun createSnapshot(keepMutable: Boolean): FirTowerDataElement =
         FirTowerDataElement(
             scope,
-            implicitReceiver?.createSnapshot(),
+            implicitReceiver?.createSnapshot(keepMutable),
-            contextReceiverGroup?.map { it.createSnapshot() },
+            contextReceiverGroup?.map { it.createSnapshot(keepMutable) },
             isLocal,
             staticScopeOwnerSymbol
         )

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

@@ -110,9 +110,9 @@ class PersistentImplicitReceiverStack private constructor(
         stack[index].updateTypeFromSmartcast(type)
     }
 
-    fun createSnapshot(): PersistentImplicitReceiverStack {
+    fun createSnapshot(keepMutable: Boolean): PersistentImplicitReceiverStack {
         return PersistentImplicitReceiverStack(
-            stack.map { it.createSnapshot() }.toPersistentList(),
+            stack.map { it.createSnapshot(keepMutable) }.toPersistentList(),
             receiversPerLabel,
             indexesPerSymbol,
             originalTypes