FIR DFA: move eq/notEq null-to-type translation to LogicSystem

This makes the `returns() implies` checker slightly cleaner, and also
fixes the case that I've missed where in RHS of `x ?:` type of `x` was
not set to `Nothing?`.

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

@@ -79,12 +79,17 @@ abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceCont
         val intersected = types.map { ConeTypeIntersector.intersectTypes(context, it.toList()) }
         val unified = context.commonSuperTypeOrNull(intersected) ?: return null
         return when {
+            unified.isNullableAny -> null
             unified.isAcceptableForSmartcast() -> unified
             unified.canBeNull -> null
             else -> context.anyType()
         }
     }
 
+    protected operator fun MutableTypeStatement.plusAssign(other: TypeStatement) {
+        exactType += other.exactType
+    }
+
     protected fun and(statements: Collection<TypeStatement>): TypeStatement? =
         statements.singleOrNew { statements.flatMapTo(mutableSetOf()) { it.exactType } }
 

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

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.fir.resolve.dfa
 
-import com.google.common.collect.ArrayListMultimap
 import kotlinx.collections.immutable.*
 import org.jetbrains.kotlin.fir.types.ConeInferenceContext
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
@@ -254,26 +253,37 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         }
     }
 
+    private val nullableNothingType = context.session.builtinTypes.nullableNothingType.type
+    private val anyType = context.session.builtinTypes.anyType.type
+
     override fun approveOperationStatement(
         flow: PersistentFlow,
         approvedStatement: OperationStatement,
         removeApprovedOrImpossible: Boolean,
     ): TypeStatements {
-        val approvedTypeStatements: ArrayListMultimap<RealVariable, TypeStatement> = ArrayListMultimap.create()
+        val result = mutableMapOf<RealVariable, MutableTypeStatement>()
         val queue = LinkedList<OperationStatement>().apply { this += approvedStatement }
         val approved = mutableSetOf<OperationStatement>()
         while (queue.isNotEmpty()) {
-            val next: OperationStatement = queue.removeFirst()
+            val next = queue.removeFirst()
             // Defense from cycles in facts
-            if (!approved.add(next)) continue
+            if (!removeApprovedOrImpossible && !approved.add(next)) continue
+
+            val operation = next.operation
             val variable = next.variable
+            if (variable.isReal()) {
+                val impliedType = if (operation == Operation.EqNull) nullableNothingType else anyType
+                result.getOrPut(variable) { MutableTypeStatement(variable) }.exactType.add(impliedType)
+            }
+
             val statements = flow.logicStatements[variable] ?: continue
             val stillUnknown = statements.removeAll {
-                val knownValue = it.condition.operation.valueIfKnown(next.operation)
+                val knownValue = it.condition.operation.valueIfKnown(operation)
                 if (knownValue == true) {
                     when (val effect = it.effect) {
                         is OperationStatement -> queue += effect
-                        is TypeStatement -> approvedTypeStatements.put(effect.variable, effect)
+                        is TypeStatement ->
+                            result.getOrPut(effect.variable) { MutableTypeStatement(effect.variable) } += effect
                     }
                 }
                 removeApprovedOrImpossible && knownValue != null
@@ -284,7 +294,7 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
                 flow.logicStatements = flow.logicStatements.put(variable, stillUnknown)
             }
         }
-        return approvedTypeStatements.asMap().mapValues { and(it.value)!! }
+        return result
     }
 
     override fun recordNewAssignment(flow: PersistentFlow, variable: RealVariable, index: Int) {

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

@@ -54,7 +54,7 @@ enum class Operation {
     fun valueIfKnown(given: Operation): Boolean? = when (this) {
         EqTrue, EqFalse -> if (given == NotEqNull) null else given == this
         EqNull -> given == EqNull
-        NotEqNull -> given == NotEqNull
+        NotEqNull -> given != EqNull
     }
 
     override fun toString(): String = when (this) {