[FIR] Create snapshot of tower data context for postponed callable referneces

This fixes an issue if type of some implicit receiver was changed because
  of changed smartcast in outer scope

^KT-51228 Fixed

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

@@ -3741,6 +3741,12 @@ public class DiagnosisCompilerFirTestdataTestGenerated extends AbstractDiagnosis
                 runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/smartcastInByClause.kt");
             }
 
+            @Test
+            @TestMetadata("smartcastOnImplicitReceiverOfCallableReference.kt")
+            public void testSmartcastOnImplicitReceiverOfCallableReference() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/smartcastOnImplicitReceiverOfCallableReference.kt");
+            }
+
             @Test
             @TestMetadata("smartcastToNothing.kt")
             public void testSmartcastToNothing() throws Exception {

compiler/fir/analysis-tests/legacy-fir-tests/tests-gen/org/jetbrains/kotlin/fir/LazyBodyIsNotTouchedTilContractsPhaseTestGenerated.java

@@ -3302,6 +3302,11 @@ public class LazyBodyIsNotTouchedTilContractsPhaseTestGenerated extends Abstract
             runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/smartcastInByClause.kt");
         }
 
+        @TestMetadata("smartcastOnImplicitReceiverOfCallableReference.kt")
+        public void testSmartcastOnImplicitReceiverOfCallableReference() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/smartcastOnImplicitReceiverOfCallableReference.kt");
+        }
+
         @TestMetadata("smartcastToNothing.kt")
         public void testSmartcastToNothing() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/smartcastToNothing.kt");

compiler/fir/analysis-tests/testData/resolve/smartcasts/smartcastOnImplicitReceiverOfCallableReference.fir.txt

@@ -0,0 +1,39 @@
+FILE: smartcastOnImplicitReceiverOfCallableReference.kt
+    public abstract interface A : R|kotlin/Any| {
+        public abstract fun foo(): R|kotlin/Unit|
+
+    }
+    public final fun R|kotlin/Any|.test_1(): R|kotlin/Unit| {
+        when (this@R|/test_1|) {
+            ($subj$ is R|A|) ->  {
+                this@R|/test_1|.R|/A.foo|()
+                ::R|/A.foo|
+            }
+            else ->  {
+                throw R|java/lang/Exception.Exception|()
+            }
+        }
+
+    }
+    public final fun R|kotlin/Any|.test_2(): R|kotlin/Unit| {
+        when (this@R|/test_2|) {
+            ($subj$ is R|A|) ->  {
+                this@R|/test_2|.R|/A.foo|()
+                fun <anonymous>(): R|kotlin/Unit| <inline=Unknown>  {
+                    this@R|/test_2|.R|/A.foo|()
+                }
+
+            }
+            else ->  {
+            }
+        }
+
+    }
+    public final fun R|kotlin/Any|.test_3(): R|kotlin/Unit| {
+        when () {
+            (this@R|/test_3| is R|A|) ->  {
+                lval f: R|kotlin/reflect/KFunction0<kotlin/Unit>| = ::R|/A.foo|
+            }
+        }
+
+    }

compiler/fir/analysis-tests/testData/resolve/smartcasts/smartcastOnImplicitReceiverOfCallableReference.kt

@@ -0,0 +1,31 @@
+// ISSUE: KT-51228
+
+interface A {
+    fun foo()
+}
+
+fun Any.test_1() {
+    when(this) {
+        is A -> {
+            foo() // ok
+            ::foo // UNRESOLVED_REFERENCE, should be ok
+        }
+        else -> throw Exception()
+    }
+}
+
+fun Any.test_2() {
+    when(this) {
+        is A -> {
+            foo(); // ok
+            { foo() }
+        }
+        else -> {}
+    }
+}
+
+fun Any.test_3() {
+    if (this is A) {
+        val f = ::foo // ok
+    }
+}

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

@@ -3741,6 +3741,12 @@ public class FirDiagnosticTestGenerated extends AbstractFirDiagnosticTest {
                 runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/smartcastInByClause.kt");
             }
 
+            @Test
+            @TestMetadata("smartcastOnImplicitReceiverOfCallableReference.kt")
+            public void testSmartcastOnImplicitReceiverOfCallableReference() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/smartcastOnImplicitReceiverOfCallableReference.kt");
+            }
+
             @Test
             @TestMetadata("smartcastToNothing.kt")
             public void testSmartcastToNothing() throws Exception {

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

@@ -3741,6 +3741,12 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
                 runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/smartcastInByClause.kt");
             }
 
+            @Test
+            @TestMetadata("smartcastOnImplicitReceiverOfCallableReference.kt")
+            public void testSmartcastOnImplicitReceiverOfCallableReference() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/smartcastOnImplicitReceiverOfCallableReference.kt");
+            }
+
             @Test
             @TestMetadata("smartcastToNothing.kt")
             public void testSmartcastToNothing() throws Exception {

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

@@ -101,8 +101,8 @@ sealed class ImplicitReceiverValue<S : FirBasedSymbol<*>>(
      * Should be called only in ImplicitReceiverStack
      */
     fun replaceType(type: ConeKotlinType) {
-        if (!mutable) throw IllegalStateException("Cannot mutate an immutable ImplicitReceiverValue")
         if (type == this.type) return
+        if (!mutable) throw IllegalStateException("Cannot mutate an immutable ImplicitReceiverValue")
         this.type = type
         receiverExpression = if (type == originalReceiverExpression.typeRef.coneType) {
             originalReceiverExpression

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

@@ -806,7 +806,7 @@ class BodyResolveContext(
 
     @OptIn(PrivateForInline::class)
     fun storeCallableReferenceContext(callableReferenceAccess: FirCallableReferenceAccess) {
-        specialTowerDataContexts.storeCallableReferenceContext(callableReferenceAccess, towerDataContext)
+        specialTowerDataContexts.storeCallableReferenceContext(callableReferenceAccess, towerDataContext.createSnapshot())
     }
 
     @OptIn(PrivateForInline::class)