FIR DFA: element-wise join at merging points of try expression

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

@@ -649,7 +649,8 @@ 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.
-        graphBuilder.enterCatchClause(catch).mergeIncomingFlow(updateReceivers = true, shouldForkFlow = true)
+        // NB: element-wise join due to multiple incoming flows: try main enter and exit
+        graphBuilder.enterCatchClause(catch).mergeIncomingFlowElementwise(updateReceivers = true, shouldForkFlow = true)
     }
 
     fun exitCatchClause(catch: FirCatch) {
@@ -659,7 +660,8 @@ 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.
-        graphBuilder.enterFinallyBlock().mergeIncomingFlow(shouldForkFlow = true)
+        // NB: element-wise join due to multiple incoming flows: try expression enter, try main exit, and catch exits
+        graphBuilder.enterFinallyBlock().mergeIncomingFlowElementwise(shouldForkFlow = true)
     }
 
     fun exitFinallyBlock(tryExpression: FirTryExpression) {
@@ -670,7 +672,8 @@ 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.
-        tryExpressionExitNode.mergeIncomingFlow(shouldForkFlow = true)
+        // NB: element-wise join due to multiple incoming flows: try main exit and catch exits (if no finally exists)
+        tryExpressionExitNode.mergeIncomingFlowElementwise(shouldForkFlow = true)
         unionNode?.let { unionFlowFromArguments(it) }
     }
 
@@ -899,8 +902,11 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         }
 
         if (isAssignment) {
-            if (initializer is FirConstExpression<*> && initializer.kind == FirConstKind.Null) return
+            if (initializer is FirConstExpression<*> && initializer.kind == FirConstKind.Null) {
-            flow.addTypeStatement(propertyVariable typeEq initializer.typeRef.coneType)
+                flow.addTypeStatement(propertyVariable typeEq property.returnTypeRef.coneType.withNullability(ConeNullability.NULLABLE))
+            } else {
+                flow.addTypeStatement(propertyVariable typeEq initializer.typeRef.coneType)
+            }
         }
     }
 
@@ -1121,12 +1127,23 @@ 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 = logicSystem.joinFlow(previousFlows)
+        var flow = mergeOperation.invoke(previousFlows)
         if (updateReceivers) {
             logicSystem.updateAllReceivers(flow)
         }

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

@@ -28,9 +28,15 @@ 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

@@ -275,6 +275,75 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         }
     }
 
+    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 {
+                    result += it
+                }
+            } else {
+                while (flow != parentFlow) {
+                    flow.approvedTypeStatements[variable]?.let {
+                        result += it
+                    }
+                    flow = flow.previousFlow!!
+                }
+            }
+            // <<< comprehensive element-wise fold <<<
+        } else {
+            result.exactType.addIfNotNull(variableUnderAlias.originalType)
+            flow.approvedTypeStatements[variableUnderAlias.variable]?.let { result += it }
+        }
+        return result
+    }
+
     override fun addTypeStatement(flow: PersistentFlow, statement: TypeStatement) {
         if (statement.isEmpty) return
         with(flow) {

compiler/testData/diagnostics/testsWithStdLib/tryCatch/correctSmartcasts.fir.kt

@@ -16,7 +16,7 @@ fun test1(s: String?) {
             requireNotNull(s)
         }
         t2.<!INAPPLICABLE_CANDIDATE!>not<!>()
-        s.length
+        s.<!INAPPLICABLE_CANDIDATE!>length<!>
     }
 }
 
@@ -45,7 +45,7 @@ fun test3() {
         s = null
         return
     }
-    s.<!INAPPLICABLE_CANDIDATE!>length<!>
+    s.length
 }
 
 fun test4() {
@@ -61,7 +61,7 @@ fun test4() {
     catch (e: ExcB) {
 
     }
-    s.<!INAPPLICABLE_CANDIDATE!>length<!>
+    s.length
 }
 
 fun test5(s: String?) {
@@ -74,7 +74,7 @@ fun test5(s: String?) {
     catch (e: ExcB) {
 
     }
-    s.length
+    s.<!INAPPLICABLE_CANDIDATE!>length<!>
 }
 
 fun test6(s: String?) {

compiler/testData/diagnostics/testsWithStdLib/tryCatch/correctSmartcasts_after.fir.kt

@@ -17,7 +17,7 @@ fun test1(s: String?) {
             requireNotNull(s)
         }
         t2.<!INAPPLICABLE_CANDIDATE!>not<!>()
-        s.length
+        s.<!INAPPLICABLE_CANDIDATE!>length<!>
     }
 }
 
@@ -46,7 +46,7 @@ fun test3() {
         s = null
         return
     }
-    s.<!INAPPLICABLE_CANDIDATE!>length<!>
+    s.length
 }
 
 fun test4() {
@@ -62,7 +62,7 @@ fun test4() {
     catch (e: ExcB) {
 
     }
-    s.<!INAPPLICABLE_CANDIDATE!>length<!>
+    s.length
 }
 
 fun test5(s: String?) {
@@ -75,7 +75,7 @@ fun test5(s: String?) {
     catch (e: ExcB) {
 
     }
-    s.length
+    s.<!INAPPLICABLE_CANDIDATE!>length<!>
 }
 
 fun test6(s: String?) {

compiler/testData/diagnostics/testsWithStdLib/tryCatch/falsePositiveSmartcasts.fir.kt

@@ -101,6 +101,6 @@ fun test6(s1: String?, s2: String?) {
         requireNotNull(s2)
     }
     s.<!INAPPLICABLE_CANDIDATE!>length<!>
-    s1.length
+    s1.<!INAPPLICABLE_CANDIDATE!>length<!>
     s2.length
 }

compiler/testData/diagnostics/testsWithStdLib/tryCatch/falsePositiveSmartcasts_after.fir.kt

@@ -103,6 +103,6 @@ fun test6(s1: String?, s2: String?) {
         requireNotNull(s2)
     }
     s.<!INAPPLICABLE_CANDIDATE!>length<!>
-    s1.length
+    s1.<!INAPPLICABLE_CANDIDATE!>length<!>
     s2.length
 }