[FIR] Fix smartcasts on member vals in init blocks

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.ImplicitReceiverStackImpl
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
+import org.jetbrains.kotlin.fir.resolve.calls.ClassDispatchReceiverValue
 import org.jetbrains.kotlin.fir.resolve.calls.ConeInferenceContext
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.*
 import org.jetbrains.kotlin.fir.resolve.dfa.contracts.buildContractFir
@@ -629,14 +630,23 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         graphBuilder.exitConstExpresion(constExpression).mergeIncomingFlow()
     }
 
-    fun exitVariableDeclaration(variable: FirProperty) {
+    fun exitLocalVariableDeclaration(variable: FirProperty) {
         val node = graphBuilder.exitVariableDeclaration(variable).mergeIncomingFlow()
         val initializer = variable.initializer ?: return
-        exitVariableInitialization(node, initializer, variable, isVariableDeclaration = true)
+        exitVariableInitialization(node, initializer, variable, assignment = null)
     }
 
-    private fun exitVariableInitialization(node: CFGNode<*>, initializer: FirExpression, variable: FirProperty, isVariableDeclaration: Boolean) {
+    fun exitVariableAssignment(assignment: FirVariableAssignment) {
-        var propertyVariable = variableStorage.getOrCreateRealVariable(variable.symbol, variable)
+        val node = graphBuilder.exitVariableAssignment(assignment).mergeIncomingFlow()
+        val property = (assignment.lValue as? FirResolvedNamedReference)?.resolvedSymbol?.fir as? FirProperty ?: return
+        // TODO: add unstable smartcast
+        if (!property.isLocal && property.isVar) return
+        exitVariableInitialization(node, assignment.rValue, property, assignment)
+    }
+
+    private fun exitVariableInitialization(node: CFGNode<*>, initializer: FirExpression, variable: FirProperty, assignment: FirVariableAssignment?) {
+        var propertyVariable = variableStorage.getOrCreateRealVariable(variable.symbol, assignment ?: variable)
+        val isVariableDeclaration = assignment == null
         if (!isVariableDeclaration) {
             // Don't remove statements for variable under alias
             if (propertyVariable.identifier.symbol == variable.symbol) {
@@ -644,7 +654,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
             } else {
                 variableStorage.unboundPossiblyAliasedVariable(variable.symbol)
                 // We should create new dataFlowVariable for local variable without alias
-                propertyVariable = variableStorage.getOrCreateRealVariable(variable.symbol, variable)
+                propertyVariable = variableStorage.getOrCreateRealVariable(variable.symbol, assignment ?: variable)
             }
         }
 
@@ -674,12 +684,6 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         }
     }
 
-    fun exitVariableAssignment(assignment: FirVariableAssignment) {
-        val node = graphBuilder.exitVariableAssignment(assignment).mergeIncomingFlow()
-        val property = (assignment.lValue as? FirResolvedNamedReference)?.resolvedSymbol?.fir as? FirProperty ?: return //error("left side of assignment should have symbol")
-        if (!property.isLocal) return
-        exitVariableInitialization(node, assignment.rValue, property, isVariableDeclaration = false)
-    }
 
     fun exitThrowExceptionNode(throwExpression: FirThrowExpression) {
         graphBuilder.exitThrowExceptionNode(throwExpression).mergeIncomingFlow()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/VariableStorage.kt

@@ -6,10 +6,7 @@
 package org.jetbrains.kotlin.fir.resolve.dfa
 
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.expressions.FirExpressionWithSmartcast
-import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
-import org.jetbrains.kotlin.fir.expressions.FirWhenSubjectExpression
 import org.jetbrains.kotlin.fir.expressions.impl.FirNoReceiverExpression
 import org.jetbrains.kotlin.fir.references.FirThisReference
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
@@ -38,7 +35,7 @@ class VariableStorage {
         fir: FirElement,
     ): Identifier {
         return localVariableAliases[symbol] ?: run {
-            val expression = fir as? FirQualifiedAccessExpression
+            val expression = fir as? FirQualifiedAccess
             Identifier(
                 symbol,
                 expression?.dispatchReceiver?.takeIf { it != FirNoReceiverExpression }?.let(this::getOrCreateVariable),
@@ -53,9 +50,11 @@ class VariableStorage {
     private fun createRealVariableInternal(identifier: Identifier, originalFir: FirElement): RealVariable {
         val receiver: FirExpression?
         val isThisReference: Boolean
-        val expression = when (originalFir) {
+        val expression: FirQualifiedAccess? = when (originalFir) {
             is FirQualifiedAccessExpression -> originalFir
             is FirWhenSubjectExpression -> originalFir.whenSubject.whenExpression.subject as? FirQualifiedAccessExpression
+            is FirVariableAssignment -> originalFir
+            is FirArraySetCall -> originalFir
             else -> null
         }
 

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

@@ -129,7 +129,7 @@ class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer)
         variable.transformAccessors()
         localScopes.lastOrNull()?.storeVariable(variable)
         variable.replaceResolvePhase(transformerPhase)
-        dataFlowAnalyzer.exitVariableDeclaration(variable)
+        dataFlowAnalyzer.exitLocalVariableDeclaration(variable)
         return variable.compose()
     }
 

compiler/fir/resolve/testData/resolve/smartcasts/smartCastInInit.dot

@@ -0,0 +1,55 @@
+digraph smartCastInInit_kt {
+    graph [splines=ortho nodesep=3]
+    node [shape=box penwidth=2]
+    edge [penwidth=2]
+
+    subgraph cluster_0 {
+        color=red
+        0 [label="Enter function foo" style="filled" fillcolor=red];
+        1 [label="Exit function foo" style="filled" fillcolor=red];
+    }
+
+    0 -> {1};
+
+    subgraph cluster_1 {
+        color=red
+        2 [label="Enter function s" style="filled" fillcolor=red];
+        3 [label="Function call: R|kotlin/TODO|()"];
+        4 [label="Stub" style="filled" fillcolor=gray];
+        5 [label="Jump: ^s R|kotlin/TODO|()" style="filled" fillcolor=gray];
+        6 [label="Stub" style="filled" fillcolor=gray];
+        7 [label="Exit function s" style="filled" fillcolor=red];
+    }
+
+    2 -> {3};
+    3 -> {7};
+    3 -> {4} [style=dotted];
+    4 -> {5} [style=dotted];
+    5 -> {7 6} [style=dotted];
+    6 -> {7} [style=dotted];
+
+    subgraph cluster_2 {
+        color=red
+        8 [label="Enter function <init>" style="filled" fillcolor=red];
+        9 [label="Exit function <init>" style="filled" fillcolor=red];
+    }
+
+    8 -> {9};
+
+    subgraph cluster_3 {
+        color=red
+        10 [label="Enter function getter" style="filled" fillcolor=red];
+        11 [label="Exit function getter" style="filled" fillcolor=red];
+    }
+
+    10 -> {11};
+
+    subgraph cluster_4 {
+        color=red
+        12 [label="Enter property" style="filled" fillcolor=red];
+        13 [label="Exit property" style="filled" fillcolor=red];
+    }
+
+    12 -> {13};
+
+}

compiler/fir/resolve/testData/resolve/smartcasts/smartCastInInit.kt

@@ -9,6 +9,6 @@ class Main {
     private val x: I
     init {
         x = s()
-        x.<!UNRESOLVED_REFERENCE!>foo<!>()
+        x.foo()
     }
 }

compiler/fir/resolve/testData/resolve/smartcasts/smartCastInInit.txt

@@ -18,7 +18,7 @@ FILE: smartCastInInit.kt
 
         init {
             this@R|/Main|.R|/Main.x| = R|/s|()
-            this@R|/Main|.R|/Main.x|.<Unresolved name: foo>#()
+            this@R|/Main|.R|/Main.x|.R|/S.foo|()
         }
 
     }

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsTestGenerated.java

@@ -1176,11 +1176,6 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
             runTest("compiler/fir/resolve/testData/resolve/problems/samConversionInConstructorCall.kt");
         }
 
-        @TestMetadata("smartCastInInit.kt")
-        public void testSmartCastInInit() throws Exception {
-            runTest("compiler/fir/resolve/testData/resolve/problems/smartCastInInit.kt");
-        }
-
         @TestMetadata("syntheticsVsNormalProperties.kt")
         public void testSyntheticsVsNormalProperties() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/problems/syntheticsVsNormalProperties.kt");

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithCfgTestGenerated.java

@@ -218,6 +218,11 @@ public class FirDiagnosticsWithCfgTestGenerated extends AbstractFirDiagnosticsWi
             runTest("compiler/fir/resolve/testData/resolve/smartcasts/simpleIf.kt");
         }
 
+        @TestMetadata("smartCastInInit.kt")
+        public void testSmartCastInInit() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/smartcasts/smartCastInInit.kt");
+        }
+
         @TestMetadata("smartcastAfterReassignment.kt")
         public void testSmartcastAfterReassignment() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/smartcasts/smartcastAfterReassignment.kt");

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithLightTreeTestGenerated.java

@@ -1176,11 +1176,6 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
             runTest("compiler/fir/resolve/testData/resolve/problems/samConversionInConstructorCall.kt");
         }
 
-        @TestMetadata("smartCastInInit.kt")
-        public void testSmartCastInInit() throws Exception {
-            runTest("compiler/fir/resolve/testData/resolve/problems/smartCastInInit.kt");
-        }
-
         @TestMetadata("syntheticsVsNormalProperties.kt")
         public void testSyntheticsVsNormalProperties() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/problems/syntheticsVsNormalProperties.kt");