FIR DFA: don't update receivers' types if flow is unchanged

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/resolver/resolverUtils.kt

@@ -55,7 +55,7 @@ internal open class StubBodyResolveTransformerComponents(
             override val receiverStack: Iterable<ImplicitReceiverValue<*>>
                 get() = error("Should not be called")
 
-            override fun receiverUpdated(symbol: FirBasedSymbol<*>, types: Set<ConeKotlinType>?) =
+            override fun receiverUpdated(symbol: FirBasedSymbol<*>, info: TypeStatement?) =
                 error("Should not be called")
 
             override fun getTypeUsingSmartcastInfo(expression: FirExpression): Pair<PropertyStability, MutableList<ConeKotlinType>>? =

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

@@ -10,7 +10,6 @@ import org.jetbrains.kotlin.contracts.description.EventOccurrencesRange
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.contracts.FirResolvedContractDescription
-import org.jetbrains.kotlin.fir.contracts.description.ConeBooleanConstantReference
 import org.jetbrains.kotlin.fir.contracts.description.ConeConditionalEffectDeclaration
 import org.jetbrains.kotlin.fir.contracts.description.ConeConstantReference
 import org.jetbrains.kotlin.fir.contracts.description.ConeReturnsEffectDeclaration
@@ -99,10 +98,10 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
                 private val visibilityChecker = components.session.visibilityChecker
                 private val typeContext = components.session.typeContext
 
-                override fun receiverUpdated(symbol: FirBasedSymbol<*>, types: Set<ConeKotlinType>?) {
+                override fun receiverUpdated(symbol: FirBasedSymbol<*>, info: TypeStatement?) {
                     val index = receiverStack.getReceiverIndex(symbol) ?: return
                     val originalType = receiverStack.getOriginalType(index)
-                    receiverStack.replaceReceiverType(index, types.intersectWith(typeContext, originalType))
+                    receiverStack.replaceReceiverType(index, info?.exactType.intersectWith(typeContext, originalType))
                 }
 
                 override val logicSystem: PersistentLogicSystem =
@@ -137,7 +136,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
 
     protected abstract val logicSystem: LogicSystem<FLOW>
     protected abstract val receiverStack: Iterable<ImplicitReceiverValue<*>>
-    protected abstract fun receiverUpdated(symbol: FirBasedSymbol<*>, types: Set<ConeKotlinType>?)
+    protected abstract fun receiverUpdated(symbol: FirBasedSymbol<*>, info: TypeStatement?)
 
     private val graphBuilder get() = context.graphBuilder
     private val variableStorage get() = context.variableStorage
@@ -161,7 +160,8 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     open fun getTypeUsingSmartcastInfo(expression: FirExpression): Pair<PropertyStability, MutableList<ConeKotlinType>>? {
         val flow = graphBuilder.lastNode.flow
         val variable = variableStorage.getRealVariableWithoutUnwrappingAlias(flow, expression) ?: return null
-        return flow.getType(variable)?.takeIf { it.isNotEmpty() }?.let { variable.stability to it.toMutableList() }
+        val types = flow.getTypeStatement(variable)?.exactType?.ifEmpty { null } ?: return null
+        return variable.stability to types.toMutableList()
     }
 
     fun returnExpressionsOfAnonymousFunction(function: FirAnonymousFunction): Collection<FirStatement> {
@@ -223,7 +223,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         if (graphBuilder.isTopLevel()) {
             context.reset()
         } else {
-            resetReceivers(graph.enterNode.flow)
+            resetReceivers()
         }
         return FirControlFlowGraphReferenceImpl(graph, DataFlowInfo(variableStorage, flowOnNodes))
     }
@@ -245,7 +245,6 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
                 enterCapturingStatement(flowOnEntry, anonymousFunction)
             else -> {}
         }
-        resetReceivers(flowOnEntry)
     }
 
     private fun exitAnonymousFunction(anonymousFunction: FirAnonymousFunction): FirControlFlowGraphReference {
@@ -261,8 +260,11 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
             else -> {}
         }
         functionExitNode.mergeIncomingFlow()
-        postponedLambdaExitNode?.mergeIncomingFlow()
+        if (postponedLambdaExitNode != null) {
-        resetReceivers(graph.enterNode.flow)
+            postponedLambdaExitNode.mergeIncomingFlow()
+        } else {
+            resetReceivers()
+        }
         return FirControlFlowGraphReferenceImpl(graph)
     }
 
@@ -617,7 +619,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     private fun CFGNode<*>.mergeWhenBranchEntryFlow() {
         val previousConditionExitNode = previousNodes.singleOrNull()
         if (previousConditionExitNode is WhenBranchConditionExitNode) {
-            val flow = mergeIncomingFlow(updateReceivers = true)
+            val flow = mergeIncomingFlow()
             val previousConditionVariable = context.variablesForWhenConditions.remove(previousConditionExitNode) ?: return
             flow.commitOperationStatement(previousConditionVariable eq false)
         } else { // first branch
@@ -733,7 +735,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     }
 
     fun enterCatchClause(catch: FirCatch) {
-        graphBuilder.enterCatchClause(catch).mergeIncomingFlow(updateReceivers = true)
+        graphBuilder.enterCatchClause(catch).mergeIncomingFlow()
     }
 
     fun exitCatchClause(catch: FirCatch) {
@@ -741,7 +743,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     }
 
     fun enterFinallyBlock() {
-        graphBuilder.enterFinallyBlock().mergeIncomingFlow(updateReceivers = true)
+        graphBuilder.enterFinallyBlock().mergeIncomingFlow()
     }
 
     fun exitFinallyBlock() {
@@ -775,7 +777,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         // is non-null. In theory, this should be unnecessary if the TODOs below are implemented.
         val flowFromPreviousSafeCall = (node.firstPreviousNode as? ExitSafeCallNode)?.lastNodeInNotNullCase?.flow
         if (flowFromPreviousSafeCall != null) {
-            logicSystem.copyAllInformation(flowFromPreviousSafeCall, flow)
+            flow.copyAllInformationFrom(flowFromPreviousSafeCall)
         }
         val receiverVariable = variableStorage.getOrCreateIfReal(node.flow, safeCall.receiver) ?: return
         flow.commitOperationStatement(receiverVariable notEq null)
@@ -851,9 +853,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     }
 
     private fun UnionFunctionCallArgumentsNode.unionFlowFromArguments() {
-        flow = logicSystem.unionFlow(previousNodes.map { it.flow }).also {
+        flow = logicSystem.unionFlow(previousNodes.map { it.flow })
-            resetReceivers(it)
-        }
     }
 
     private fun processConditionalContract(qualifiedAccess: FirQualifiedAccess) {
@@ -913,7 +913,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
                 }
             }
         }
-        graphBuilder.exitContract(qualifiedAccess).mergeIncomingFlow(updateReceivers = true)
+        graphBuilder.exitContract(qualifiedAccess).mergeIncomingFlow()
         contractDescriptionVisitingMode = false
     }
 
@@ -1137,7 +1137,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         val lhsIsNotNullFlow = lhsIsNotNullNode.mergeIncomingFlow()
         val rhsEnterFlow = rhsEnterNode.mergeIncomingFlow()
         val lhsVariable = variableStorage.getOrCreateIfReal(flow, elvisExpression.lhs) ?: return
-        lhsIsNotNullFlow.commitOperationStatement(lhsVariable notEq null, updateReceivers = false)
+        lhsIsNotNullFlow.commitOperationStatement(lhsVariable notEq null)
         rhsEnterFlow.commitOperationStatement(lhsVariable eq null)
     }
 
@@ -1193,35 +1193,58 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
 
     private val CFGNode<*>.origin: CFGNode<*> get() = if (this is StubNode) firstPreviousNode else this
 
-    // `updateReceivers` should be set if the predecessor node is not the last one that was handled:
+    // Smart cast information is taken from `graphBuilder.lastNode`, but the problem with receivers specifically
-    // when entering a new `when` branch condition or a `catch`/`finally` block.
+    // is that they also affect tower resolver's scope stack. To allow accessing members on smart casted receivers,
-    private fun CFGNode<*>.mergeIncomingFlow(updateReceivers: Boolean = false): FLOW {
+    // we explicitly patch up the stack by calling `receiverUpdated` in a way that maintains consistency with
-        var incomingEdgeCount = 0
+    // `getTypeUsingSmartcastInfo`; i.e. at any point between calls to this class' methods the types in the implicit
+    // receiver stack also correspond to the data flow information attached to `graphBuilder.lastNode`.
+    private var currentReceiverState: FLOW? = null
+
+    // Generally when calling some method on `graphBuilder`, one of the nodes it returns is the new `lastNode`.
+    // In that case `mergeIncomingFlow` will automatically ensure consistency once called on that node.
+    private fun CFGNode<*>.mergeIncomingFlow(): FLOW {
         val previousFlows = previousNodes.mapNotNull {
             val incomingEdgeKind = incomingEdges.getValue(it).kind
-            if (incomingEdgeKind.usedInDeadDfa) {
+            it.takeIf { incomingEdgeKind.usedInDfa || (isDead && incomingEdgeKind.usedInDeadDfa) }?.flow
-                incomingEdgeCount++
-            }
-            if (incomingEdgeKind.usedInDfa || (isDead && incomingEdgeKind.usedInDeadDfa)) {
-                it.flow
-            } else {
-                null
-            }
         }
-        return logicSystem.joinFlow(previousFlows).also {
+        val result = logicSystem.joinFlow(previousFlows).also { flow = it }
-            flow = it
+        if (graphBuilder.lastNodeOrNull == this) {
-            if (updateReceivers || incomingEdgeCount > 1) {
+            // Here it is, the new `lastNode`. If the previous state is the only predecessor, then there is actually
-                resetReceivers(it)
+            // nothing to update; `addTypeStatement` has already ensured we have the correct information.
+            if (currentReceiverState == null || previousFlows.singleOrNull() != currentReceiverState) {
+                updateAllReceivers(currentReceiverState, result)
+            }
+            currentReceiverState = result
+        }
+        return result
+    }
+
+    // In rare cases (like after exiting functions) after adding more nodes `graphBuilder` will revert the current
+    // state to a previously created node, so none of the nodes it returned are `lastNode` and `mergeIncomingFlow`
+    // will not ensure consistency. In that case an explicit call to `resetReceivers` is needed to roll back the stack
+    // to that previously created node's state.
+    private fun resetReceivers() {
+        val currentFlow = graphBuilder.lastNodeOrNull?.flow
+        updateAllReceivers(currentReceiverState, currentFlow)
+        currentReceiverState = currentFlow
+    }
+
+    private fun updateAllReceivers(from: FLOW?, to: FLOW?) {
+        receiverStack.forEach {
+            variableStorage.getLocalVariable(it.boundSymbol)?.let { variable ->
+                val newStatement = to?.getTypeStatement(variable)
+                if (newStatement != from?.getTypeStatement(variable)) {
+                    receiverUpdated(it.boundSymbol, newStatement)
+                }
             }
         }
     }
 
-    private fun resetReceivers(flow: FLOW) {
+    private fun FLOW.copyAllInformationFrom(other: FLOW) {
-        receiverStack.forEach {
+        if (this === currentReceiverState) {
-            variableStorage.getLocalVariable(it.boundSymbol)?.let { variable ->
+            updateAllReceivers(this, other)
-                receiverUpdated(it.boundSymbol, flow.getType(variable))
-            }
         }
+        logicSystem.copyAllInformation(other, this)
     }
 
     private fun FLOW.addImplication(statement: Implication) {
@@ -1240,9 +1263,9 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     }
 
     private fun FLOW.addTypeStatement(info: TypeStatement) {
-        logicSystem.addTypeStatement(this, info)
+        val newStatement = logicSystem.addTypeStatement(this, info) ?: return
-        if (info.variable.isThisReference) {
+        if (newStatement.variable.isThisReference && this === currentReceiverState) {
-            receiverUpdated(info.variable.identifier.symbol, getType(info.variable))
+            receiverUpdated(newStatement.variable.identifier.symbol, newStatement)
         }
     }
 
@@ -1250,20 +1273,14 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         statements.values.forEach { addImplication(condition implies it) }
     }
 
-    private fun FLOW.commitOperationStatement(statement: OperationStatement, updateReceivers: Boolean = true) {
+    private fun FLOW.commitOperationStatement(statement: OperationStatement) {
         logicSystem.approveOperationStatement(this, statement, removeApprovedOrImpossible = true).values.forEach {
-            logicSystem.addTypeStatement(this, it)
+            addTypeStatement(it)
-            if (updateReceivers && it.variable.isThisReference) {
-                receiverUpdated(it.variable.identifier.symbol, getType(it.variable))
-            }
         }
         if (statement.operation == Operation.NotEqNull) {
             val variable = statement.variable
             if (variable is RealVariable) {
-                logicSystem.addTypeStatement(this, variable typeEq any)
+                addTypeStatement(variable typeEq any)
-                if (updateReceivers && variable.isThisReference) {
-                    receiverUpdated(variable.identifier.symbol, getType(variable))
-                }
             }
         }
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt

@@ -43,6 +43,9 @@ class ControlFlowGraphBuilder {
     val lastNode: CFGNode<*>
         get() = lastNodes.top()
 
+    val lastNodeOrNull: CFGNode<*>?
+        get() = lastNodes.topOrNull()
+
     var levelCounter: Int = 0
 
     private val modes: Stack<Mode> = stackOf(Mode.TopLevel)

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/Flow.kt

@@ -5,10 +5,8 @@
 
 package org.jetbrains.kotlin.fir.resolve.dfa
 
-import org.jetbrains.kotlin.fir.types.ConeKotlinType
-
 abstract class Flow {
     abstract val approvedTypeStatements: TypeStatements
     abstract fun unwrapVariable(variable: RealVariable): RealVariable
-    abstract fun getType(variable: RealVariable): Set<ConeKotlinType>?
+    abstract fun getTypeStatement(variable: RealVariable): TypeStatement?
 }

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/LogicSystem.kt

@@ -15,7 +15,8 @@ abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceCont
     abstract fun unionFlow(flows: Collection<FLOW>): FLOW
 
     // -------------------------------- Flow mutators --------------------------------
-    abstract fun addTypeStatement(flow: FLOW, statement: TypeStatement)
+    // Returns all known information about the variable, or null if unchanged by this statement:
+    abstract fun addTypeStatement(flow: FLOW, statement: TypeStatement): TypeStatement?
     abstract fun addImplication(flow: FLOW, implication: Implication)
     abstract fun addLocalVariableAlias(flow: FLOW, alias: RealVariable, underlyingVariable: RealVariable)
     abstract fun recordNewAssignment(flow: FLOW, variable: RealVariable, index: Int)
@@ -47,7 +48,7 @@ abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceCont
         right.isEmpty() -> right
         else -> buildMap {
             for ((variable, leftStatement) in left) {
-                put(variable, or(listOf(leftStatement, right[variable] ?: continue)))
+                put(variable, or(listOf(leftStatement, right[variable] ?: continue))!!)
             }
         }
     }
@@ -57,14 +58,14 @@ abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceCont
         right.isEmpty() -> left
         else -> left.toMutableMap().apply {
             for ((variable, rightStatement) in right) {
-                put(variable, { rightStatement }, { and(listOf(it, rightStatement)) })
+                put(variable, { rightStatement }, { and(listOf(it, rightStatement))!! })
             }
         }
     }
 
-    private inline fun Collection<TypeStatement>.singleOrNew(exactType: () -> MutableSet<ConeKotlinType>): TypeStatement =
+    private inline fun Collection<TypeStatement>.singleOrNew(exactType: () -> MutableSet<ConeKotlinType>): TypeStatement? =
         when (size) {
-            0 -> throw AssertionError("need at least one statement")
+            0 -> null
             1 -> first()
             else -> {
                 val variable = first().variable
@@ -84,9 +85,9 @@ abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceCont
         }
     }
 
-    protected fun and(statements: Collection<TypeStatement>): TypeStatement =
+    protected fun and(statements: Collection<TypeStatement>): TypeStatement? =
         statements.singleOrNew { statements.flatMapTo(mutableSetOf()) { it.exactType } }
 
-    protected fun or(statements: Collection<TypeStatement>): TypeStatement =
+    protected fun or(statements: Collection<TypeStatement>): TypeStatement? =
         statements.singleOrNew { unifyTypes(statements.map { it.exactType })?.let { mutableSetOf(it) } ?: mutableSetOf() }
 }

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/PersistentLogicSystem.kt

@@ -64,11 +64,8 @@ class PersistentFlow : Flow {
     override fun unwrapVariable(variable: RealVariable): RealVariable =
         directAliasMap[variable] ?: variable
 
-    fun getTypeStatement(variable: RealVariable): TypeStatement =
+    override fun getTypeStatement(variable: RealVariable): TypeStatement? =
-        approvedTypeStatements[unwrapVariable(variable)]?.copy(variable = variable) ?: MutableTypeStatement(variable)
+        approvedTypeStatements[unwrapVariable(variable)]?.copy(variable = variable)
-
-    override fun getType(variable: RealVariable): Set<ConeKotlinType>? =
-        approvedTypeStatements[unwrapVariable(variable)]?.exactType
 }
 
 abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSystem<PersistentFlow>(context) {
@@ -108,7 +105,7 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         // Computing the statements for these aliases is redundant as the result is equal
         // to the statements for whatever they are aliasing.
         val variables = flows.flatMapTo(mutableSetOf()) { it.approvedTypeStatements.keys + it.directAliasMap.keys } - commonAliases.keys
-        val statements = variables.mapNotNull { variable ->
+        val statements = variables.mapNotNull computeStatement@{ variable ->
             val statement = if (allExecute) {
                 // All input flows execute in some order. If none of the flows reassign, i.e. the only key
                 // in `byAssignment` is `commonFlow`'s assignment index, then all statements are true.
@@ -118,12 +115,14 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
                 if (byAssignment.size > 1) {
                     byAssignment.remove(commonFlow.assignmentIndex[variable] ?: -1)
                 }
-                or(byAssignment.values.map { flowSubset -> and(flowSubset.map { it.getTypeStatement(variable) }) })
+                or(byAssignment.values.map { flowSubset ->
+                    and(flowSubset.mapNotNull { it.getTypeStatement(variable) }) ?: return@computeStatement null
+                })
             } else {
                 // One input flow executes - one set of statements is true, others might be false.
-                or(flows.map { it.getTypeStatement(variable) })
+                or(flows.map { it.getTypeStatement(variable) ?: return@computeStatement null })
             }
-            if (statement.isNotEmpty) variable to statement.toPersistent() else null
+            if (statement?.isEmpty == false) variable to statement.toPersistent() else null
         }
 
         // If a variable was reassigned in one branch, it was reassigned at the join point.
@@ -213,13 +212,15 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         }
     }
 
-    override fun addTypeStatement(flow: PersistentFlow, statement: TypeStatement) {
+    override fun addTypeStatement(flow: PersistentFlow, statement: TypeStatement): TypeStatement? {
-        if (statement.exactType.isEmpty()) return
+        if (statement.exactType.isEmpty()) return null
         val variable = statement.variable
         val oldExactType = flow.approvedTypeStatements[variable]?.exactType
         val newExactType = oldExactType?.addAll(statement.exactType) ?: statement.exactType.toPersistentSet()
-        if (newExactType === oldExactType) return
+        if (newExactType === oldExactType) return null
-        flow.approvedTypeStatements = flow.approvedTypeStatements.put(variable, PersistentTypeStatement(variable, newExactType))
+        val newStatement = PersistentTypeStatement(variable, newExactType)
+        flow.approvedTypeStatements = flow.approvedTypeStatements.put(variable, newStatement)
+        return newStatement
     }
 
     override fun addImplication(flow: PersistentFlow, implication: Implication) {
@@ -283,7 +284,7 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
                 flow.logicStatements = flow.logicStatements.put(variable, stillUnknown)
             }
         }
-        return approvedTypeStatements.asMap().mapValues { and(it.value) }
+        return approvedTypeStatements.asMap().mapValues { and(it.value)!! }
     }
 
     override fun recordNewAssignment(flow: PersistentFlow, variable: RealVariable, index: Int) {