FIR DFA: use element-wise join everywhere

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

@@ -649,8 +649,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     fun enterCatchClause(catch: FirCatch) {
         // NB: fork to isolate effects inside the catch clause
         // Otherwise, changes in the catch clause could affect the previous node: try main block.
-        // NB: element-wise join due to multiple incoming flows: try main enter and exit
+        graphBuilder.enterCatchClause(catch).mergeIncomingFlow(updateReceivers = true, shouldForkFlow = true)
-        graphBuilder.enterCatchClause(catch).mergeIncomingFlowElementwise(updateReceivers = true, shouldForkFlow = true)
     }
 
     fun exitCatchClause(catch: FirCatch) {
@@ -660,8 +659,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     fun enterFinallyBlock() {
         // NB: fork to isolate effects inside the finally block
         // Otherwise, changes in the finally block could affect the previous nodes: try main block and catch clauses.
-        // NB: element-wise join due to multiple incoming flows: try expression enter, try main exit, and catch exits
+        graphBuilder.enterFinallyBlock().mergeIncomingFlow(shouldForkFlow = true)
-        graphBuilder.enterFinallyBlock().mergeIncomingFlowElementwise(shouldForkFlow = true)
     }
 
     fun exitFinallyBlock(tryExpression: FirTryExpression) {
@@ -672,8 +670,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         val (tryExpressionExitNode, unionNode) = graphBuilder.exitTryExpression(callCompleted)
         // NB: fork to prevent effects after the try expression from being flown into the try expression
         // Otherwise, changes in any following nodes could affect the previous nodes, including try main block and finally block if any.
-        // NB: element-wise join due to multiple incoming flows: try main exit and catch exits (if no finally exists)
+        tryExpressionExitNode.mergeIncomingFlow(shouldForkFlow = true)
-        tryExpressionExitNode.mergeIncomingFlowElementwise(shouldForkFlow = true)
         unionNode?.let { unionFlowFromArguments(it) }
     }
 
@@ -1127,23 +1124,12 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     private fun <T : CFGNode<*>> T.mergeIncomingFlow(
         updateReceivers: Boolean = false,
         shouldForkFlow: Boolean = false
-    ): T = foldIncomingFlow(logicSystem::joinFlow, updateReceivers, shouldForkFlow)
-
-    private fun <T : CFGNode<*>> T.mergeIncomingFlowElementwise(
-        updateReceivers: Boolean = false,
-        shouldForkFlow: Boolean = false
-    ): T = foldIncomingFlow(logicSystem::elementwiseJoinFlow, updateReceivers, shouldForkFlow)
-
-    private inline fun <T : CFGNode<*>> T.foldIncomingFlow(
-        mergeOperation: (Collection<FLOW>) -> FLOW,
-        updateReceivers: Boolean = false,
-        shouldForkFlow: Boolean = false
     ): T = this.also { node ->
         val previousFlows = if (node.isDead)
             node.previousNodes.mapNotNull { runIf(!node.incomingEdges.getValue(it).kind.isBack) { it.flow } }
         else
             node.previousNodes.mapNotNull { prev -> prev.takeIf { node.incomingEdges.getValue(it).kind.usedInDfa }?.flow }
-        var flow = mergeOperation.invoke(previousFlows)
+        var flow = logicSystem.joinFlow(previousFlows)
         if (updateReceivers) {
             logicSystem.updateAllReceivers(flow)
         }

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

@@ -28,15 +28,9 @@ abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceCont
 
     abstract fun createEmptyFlow(): FLOW
     abstract fun forkFlow(flow: FLOW): FLOW
-
-    // Differential computation
     abstract fun joinFlow(flows: Collection<FLOW>): FLOW
     abstract fun unionFlow(flows: Collection<FLOW>): FLOW
 
-    // Comprehensive element-wise computation
-    abstract fun elementwiseJoinFlow(flows: Collection<FLOW>): FLOW
-    abstract fun elementwiseUnionFlow(flows: Collection<FLOW>): FLOW
-
     abstract fun addTypeStatement(flow: FLOW, statement: TypeStatement)
 
     abstract fun addImplication(flow: FLOW, implication: Implication)

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

@@ -109,11 +109,6 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         return foldFlow(
             flows,
             mergeOperation = { statements -> this.or(statements).takeIf { it.isNotEmpty } },
-            computeVariables = { computeVariablesDiffWithCommonFlow ->
-                flows.map {
-                    computeVariablesDiffWithCommonFlow(it).toList()
-                }.intersectSets().takeIf { it.isNotEmpty() }
-            }
         )
     }
 
@@ -121,18 +116,12 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         return foldFlow(
             flows,
             this::and,
-            computeVariables = { computeVariablesDiffWithCommonFlow ->
-                flows.flatMapTo(mutableSetOf()) {
-                    computeVariablesDiffWithCommonFlow(it)
-                }
-            }
         )
     }
 
     private inline fun foldFlow(
         flows: Collection<PersistentFlow>,
         mergeOperation: (Collection<TypeStatement>) -> MutableTypeStatement?,
-        computeVariables: (computeVariablesDiffWithCommonFlow: (PersistentFlow) -> Iterable<RealVariable>) -> Collection<RealVariable>?
     ): PersistentFlow {
         if (flows.isEmpty()) return createEmptyFlow()
         flows.singleOrNull()?.let { return it }
@@ -140,11 +129,11 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         val aliasedVariablesThatDontChangeAlias = computeAliasesThatDontChange(flows)
 
         val commonFlow = flows.reduce(::lowestCommonFlow)
-        val variables =
-            computeVariables { it.diffVariablesIterable(commonFlow, aliasedVariablesThatDontChangeAlias.keys) } ?: return commonFlow
 
+        val variables = flows.flatMap { it.approvedTypeStatements.keys }.toSet()
         for (variable in variables) {
-            val info = mergeOperation(flows.map { it.getApprovedTypeStatementsDiff(variable, commonFlow) }) ?: continue
+            val info = mergeOperation(flows.map { it.getApprovedTypeStatements(variable, commonFlow) }) ?: continue
+            removeAllAboutVariable(commonFlow, variable)
             commonFlow.addApprovedStatements(info)
         }
 
@@ -208,121 +197,12 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         flow.backwardsAliasMap = flow.backwardsAliasMap.put(original, variables - alias)
     }
 
-    @OptIn(DfaInternals::class)
-    private fun PersistentFlow.getApprovedTypeStatementsDiff(variable: RealVariable, parentFlow: PersistentFlow): MutableTypeStatement {
-        var flow = this
-        val result = MutableTypeStatement(variable)
-        val variableUnderAlias = directAliasMap[variable]
-        if (variableUnderAlias == null) {
-            while (flow != parentFlow) {
-                flow.approvedTypeStatementsDiff[variable]?.let {
-                    result += it
-                }
-                flow = flow.previousFlow!!
-            }
-        } else {
-            result.exactType.addIfNotNull(variableUnderAlias.originalType)
-            flow.approvedTypeStatements[variableUnderAlias.variable]?.let { result += it }
-        }
-        return result
-    }
-
-    /**
-     * This is an iterable over real variable that has known facts in flow range
-     *   from [this] to [parentFlow]
-     */
-    private fun PersistentFlow.diffVariablesIterable(
-        parentFlow: PersistentFlow,
-        aliasedVariablesThatDontChangeAlias: Set<RealVariable>
-    ): Iterable<RealVariable> =
-        object : DiffIterable<RealVariable>(parentFlow, this) {
-            override fun extractIterator(flow: PersistentFlow): Iterator<RealVariable> {
-                val variablesWithNewInfo = flow.approvedTypeStatementsDiff.keys
-                val updatedVariables = ArrayList(variablesWithNewInfo)
-                updatedVariables += flow.updatedAliasDiff
-                variablesWithNewInfo.flatMapTo(updatedVariables) { variableWithNewInfo ->
-                    flow.backwardsAliasMap[variableWithNewInfo]?.filter { it !in aliasedVariablesThatDontChangeAlias } ?: emptyList()
-                }
-                return updatedVariables.iterator()
-            }
-        }
-
-    private abstract class DiffIterable<T>(private val parentFlow: PersistentFlow, private var currentFlow: PersistentFlow) : Iterable<T> {
-        @Suppress("LeakingThis")
-        private var currentIterator = extractIterator(currentFlow)
-
-        abstract fun extractIterator(flow: PersistentFlow): Iterator<T>
-
-        override fun iterator(): Iterator<T> {
-            return object : Iterator<T> {
-                override fun hasNext(): Boolean {
-                    if (currentIterator.hasNext()) return true
-                    while (currentFlow != parentFlow) {
-                        currentFlow = currentFlow.previousFlow!!
-                        currentIterator = extractIterator(currentFlow)
-                        if (currentIterator.hasNext()) return true
-                    }
-                    return false
-                }
-
-                override fun next(): T {
-                    if (!hasNext()) {
-                        throw NoSuchElementException()
-                    }
-                    return currentIterator.next()
-                }
-            }
-        }
-    }
-
-    override fun elementwiseJoinFlow(flows: Collection<PersistentFlow>): PersistentFlow {
-        return elementwiseFoldFlow(
-            flows,
-            mergeOperation = { statements -> this.or(statements).takeIf { it.isNotEmpty } }
-        )
-    }
-
-    override fun elementwiseUnionFlow(flows: Collection<PersistentFlow>): PersistentFlow {
-        return elementwiseFoldFlow(
-            flows,
-            mergeOperation = this::and
-        )
-    }
-
-    private inline fun elementwiseFoldFlow(
-        flows: Collection<PersistentFlow>,
-        mergeOperation: (Collection<TypeStatement>) -> MutableTypeStatement?,
-    ): PersistentFlow {
-        if (flows.isEmpty()) return createEmptyFlow()
-        flows.singleOrNull()?.let { return it }
-
-        val aliasedVariablesThatDontChangeAlias = computeAliasesThatDontChange(flows)
-
-        val commonFlow = flows.reduce(::lowestCommonFlow)
-
-        // >>> comprehensive element-wise fold >>>
-        val variables = flows.flatMap { it.approvedTypeStatements.keys }.toSet()
-        for (variable in variables) {
-            val info = mergeOperation(flows.map { it.getApprovedTypeStatements(variable, commonFlow) }) ?: continue
-            removeAllAboutVariable(commonFlow, variable)
-            commonFlow.addApprovedStatements(info)
-        }
-        // <<< comprehensive element-wise fold <<<
-
-        commonFlow.addVariableAliases(aliasedVariablesThatDontChangeAlias)
-
-        updateAllReceivers(commonFlow)
-
-        return commonFlow
-    }
-
     @OptIn(DfaInternals::class)
     private fun PersistentFlow.getApprovedTypeStatements(variable: RealVariable, parentFlow: PersistentFlow): MutableTypeStatement {
         var flow = this
         val result = MutableTypeStatement(variable)
         val variableUnderAlias = directAliasMap[variable]
         if (variableUnderAlias == null) {
-            // >>> comprehensive element-wise fold >>>
             // get approved type statement even though the starting flow == parent flow
             if (flow == parentFlow) {
                 flow.approvedTypeStatements[variable]?.let {
@@ -336,7 +216,6 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
                     flow = flow.previousFlow!!
                 }
             }
-            // <<< comprehensive element-wise fold <<<
         } else {
             result.exactType.addIfNotNull(variableUnderAlias.originalType)
             flow.approvedTypeStatements[variableUnderAlias.variable]?.let { result += it }

compiler/testData/diagnostics/tests/smartCasts/loops/whileWithAssertInConditionAndBreakAfter.fir.kt

@@ -8,7 +8,7 @@ fun foo() {
 
     }
     // TODO: this testdata fixates undesired behavior (it should be an unsafe call)
-    x.length // 'x' is unsoundly smartcasted here
+    x.<!INAPPLICABLE_CANDIDATE!>length<!> // 'x' is unsoundly smartcasted here
 }
 
 fun bar() {
@@ -19,5 +19,5 @@ fun bar() {
 
     }
     // TODO: this testdata fixates undesired behavior (it should be an unsafe call)
-    x.size // 'x' is unsoundly smartcasted here
+    x.<!INAPPLICABLE_CANDIDATE!>size<!> // 'x' is unsoundly smartcasted here
 }

compiler/testData/diagnostics/tests/smartCasts/loops/whileWithAssertInConditionAndBreakBefore.fir.kt

@@ -7,7 +7,7 @@ fun foo() {
         break
 
     }
-    x.length // 'x' is unsoundly smartcasted here
+    x.<!INAPPLICABLE_CANDIDATE!>length<!> // 'x' is unsoundly smartcasted here
 }
 
 fun bar() {
@@ -17,5 +17,5 @@ fun bar() {
         break
 
     }
-    x.size // 'x' is unsoundly smartcasted here
+    x.<!INAPPLICABLE_CANDIDATE!>size<!> // 'x' is unsoundly smartcasted here
 }