[FIR] Copy implications from previous data flows when there is only one

Postponed lambdas introduce a host of challenges in data-flow analysis.
While inheriting type statements was disabled while these challenges are
being considered, we must still copy type statements from non-postponed
lambda edges. It seems the same is true for implications.

Implications are copied from previous flows when there is only a single
previous flow. That is because it never seemed to be required based on
test results. However, a recent test case revealed that copying is
required when there are multiple previous flows, but only one flow is
from a non-postponed-lambda node.

Combining implications from multiple non-postponed-lambda nodes did not
have an impact on test results, so until such a test case can be
created, the overhead of calculating common implications from multiple
flows will be avoided.

^KT-63351 Fixed

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

@@ -70,7 +70,7 @@ class PersistentFlow internal constructor(
 class MutableFlow internal constructor(
     private val previousFlow: PersistentFlow?,
     internal val approvedTypeStatements: PersistentMap.Builder<RealVariable, PersistentTypeStatement>,
-    internal val logicStatements: PersistentMap.Builder<DataFlowVariable, PersistentList<Implication>>,
+    internal val implications: PersistentMap.Builder<DataFlowVariable, PersistentList<Implication>>,
     internal val assignmentIndex: PersistentMap.Builder<RealVariable, Int>,
     internal val directAliasMap: PersistentMap.Builder<RealVariable, RealVariable>,
     internal val backwardsAliasMap: PersistentMap.Builder<RealVariable, PersistentSet<RealVariable>>,
@@ -96,7 +96,7 @@ class MutableFlow internal constructor(
     fun freeze(): PersistentFlow = PersistentFlow(
         previousFlow,
         approvedTypeStatements.build(),
-        logicStatements.build(),
+        implications.build(),
         assignmentIndex.build(),
         directAliasMap.build(),
         backwardsAliasMap.build(),

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

@@ -28,7 +28,8 @@ abstract class LogicSystem(private val context: ConeInferenceContext) {
      *
      * @param flows All [PersistentFlow]s which flow into the join flow. These will determine assignments and variable aliases for the
      * resulting join flow.
-     * @param statementFlows A *subset* of [flows] used to determine what [TypeStatement]s will be copied to the join flow.
+     * @param statementFlows A *subset* of [flows] used to determine what [TypeStatement]s and [Implication]s will be copied to the joined
+     * flow.
      * @param union Determines if [TypeStatement]s from different flows should be combined with union or intersection logic.
      */
     fun joinFlow(flows: Collection<PersistentFlow>, statementFlows: Collection<PersistentFlow>, union: Boolean): MutableFlow {
@@ -48,7 +49,7 @@ abstract class LogicSystem(private val context: ConeInferenceContext) {
             result.copyCommonAliases(flows)
         }
         result.copyStatements(statementFlows, commonFlow, union)
-        // TODO: compute common implications?
+        result.copyImplications(statementFlows)
         return result
     }
 
@@ -77,7 +78,7 @@ abstract class LogicSystem(private val context: ConeInferenceContext) {
             (effect.isEmpty || flow.approvedTypeStatements[effect.variable]?.exactType?.containsAll(effect.exactType) == true)
         ) return
         val variable = implication.condition.variable
-        flow.logicStatements[variable] = flow.logicStatements[variable]?.add(implication) ?: persistentListOf(implication)
+        flow.implications[variable] = flow.implications[variable]?.add(implication) ?: persistentListOf(implication)
     }
 
     fun translateVariableFromConditionInStatements(
@@ -87,12 +88,12 @@ abstract class LogicSystem(private val context: ConeInferenceContext) {
         transform: (Implication) -> Implication? = { it }
     ) {
         val statements = if (originalVariable.isSynthetic())
-            flow.logicStatements.remove(originalVariable)
+            flow.implications.remove(originalVariable)
         else
-            flow.logicStatements[originalVariable]
+            flow.implications[originalVariable]
         if (statements.isNullOrEmpty()) return
-        val existing = flow.logicStatements[newVariable] ?: persistentListOf()
+        val existing = flow.implications[newVariable] ?: persistentListOf()
-        flow.logicStatements[newVariable] = statements.mapNotNullTo(existing.builder()) {
+        flow.implications[newVariable] = statements.mapNotNullTo(existing.builder()) {
             // TODO: rethink this API - technically it permits constructing invalid flows
             //  (transform can replace the variable in the condition with anything)
             transform(OperationStatement(newVariable, it.condition.operation) implies it.effect)
@@ -108,7 +109,7 @@ abstract class LogicSystem(private val context: ConeInferenceContext) {
         statement: OperationStatement,
         removeApprovedOrImpossible: Boolean,
     ): TypeStatements {
-        return approveOperationStatement(flow.logicStatements, statement, removeApprovedOrImpossible)
+        return approveOperationStatement(flow.implications, statement, removeApprovedOrImpossible)
     }
 
     fun recordNewAssignment(flow: MutableFlow, variable: RealVariable, index: Int) {
@@ -201,13 +202,21 @@ abstract class LogicSystem(private val context: ConeInferenceContext) {
         }
     }
 
+    private fun MutableFlow.copyImplications(flows: Collection<PersistentFlow>) {
+        when (flows.size) {
+            0 -> {}
+            1 -> implications += flows.first().implications
+            else -> {} // TODO, KT-65293: compute common implications?
+        }
+    }
+
     private fun MutableFlow.replaceVariable(variable: RealVariable, replacement: RealVariable?) {
         val original = directAliasMap.remove(variable)
         if (original != null) {
             // All statements should've been made about whatever variable this is an alias to. There is nothing to replace.
             if (AbstractTypeChecker.RUN_SLOW_ASSERTIONS) {
                 assert(variable !in backwardsAliasMap)
-                assert(variable !in logicStatements)
+                assert(variable !in implications)
                 assert(variable !in approvedTypeStatements)
             }
             // `variable.dependentVariables` is not separated by flow, so it may be non-empty if aliasing of this variable
@@ -230,7 +239,7 @@ abstract class LogicSystem(private val context: ConeInferenceContext) {
                     variableStorage.copyRealVariableWithRemapping(dependent, variable, it)
                 })
             }
-            logicStatements.replaceVariable(variable, replacementOrNext)
+            implications.replaceVariable(variable, replacementOrNext)
             approvedTypeStatements.replaceVariable(variable, replacementOrNext)
             if (aliases != null && replacementOrNext != null) {
                 directAliasMap -= replacementOrNext

compiler/testData/diagnostics/tests/smartCasts/safecalls/longChain.fir.kt

@@ -1,6 +1,14 @@
-fun calc(x: List<String>?) {
+// WITH_STDLIB
+// ISSUE: KT-63351
+
+fun test1(x: List<String>?) {
     // x should be non-null in arguments list, despite of a chain
     x?.subList(0, x.size)?.
        subList(0, x.size)?.
        get(x.size)
 }
+
+fun test2(x: List<String>?) {
+    x?.filter { true }!!.size
+    x.size
+}

compiler/testData/diagnostics/tests/smartCasts/safecalls/longChain.kt

@@ -1,6 +1,14 @@
-fun calc(x: List<String>?) {
+// WITH_STDLIB
+// ISSUE: KT-63351
+
+fun test1(x: List<String>?) {
     // x should be non-null in arguments list, despite of a chain
     x?.subList(0, <!DEBUG_INFO_SMARTCAST!>x<!>.size)?.
        subList(0, <!DEBUG_INFO_SMARTCAST!>x<!>.size)?.
        get(<!DEBUG_INFO_SMARTCAST!>x<!>.size)
 }
+
+fun test2(x: List<String>?) {
+    x?.filter { true }!!.size
+    <!DEBUG_INFO_SMARTCAST!>x<!>.size
+}

compiler/testData/diagnostics/tests/smartCasts/safecalls/longChain.txt

@@ -1,3 +1,4 @@
 package
 
-public fun calc(/*0*/ x: kotlin.collections.List<kotlin.String>?): kotlin.Unit
+public fun test1(/*0*/ x: kotlin.collections.List<kotlin.String>?): kotlin.Unit
+public fun test2(/*0*/ x: kotlin.collections.List<kotlin.String>?): kotlin.Unit