From 8e4056652004e84a5a500ba881fe904268a95fbe Mon Sep 17 00:00:00 2001
From: pyos <pyos@google.com>
Date: Thu, 3 Nov 2022 10:39:36 +0100
Subject: [PATCH] Minor: make a function that intersects N maps shorter

---
 .../fir/resolve/dfa/PersistentLogicSystem.kt  | 25 ++++++-------------
 1 file changed, 8 insertions(+), 17 deletions(-)

diff --git a/compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/PersistentLogicSystem.kt b/compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/PersistentLogicSystem.kt
index 0f992a9eb3b..e6bc6438270 100644
--- a/compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/PersistentLogicSystem.kt
+++ b/compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/PersistentLogicSystem.kt
@@ -132,9 +132,11 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         }
 
     private fun foldFlow(flows: Collection<PersistentFlow>, statements: Collection<TypeStatement>): PersistentFlow {
-        val aliasedVariablesThatDontChangeAlias = computeAliasesThatDontChange(flows)
-
         val commonFlow = flows.reduce(::lowestCommonFlow)
+        // Aliases that have occurred before branching should already be in `commonFlow`,
+        // but it might be useful to also add aliases that happen in all branches, e.g.
+        // the aliasing of `y` to `a.x` after `if (p) { y = a.x } else { y = a.x }`.
+        val commonAliases = computeCommonAliases(flows)
 
         for (flow in flows) {
             for ((variable, index) in flow.assignmentIndex) {
@@ -150,28 +152,17 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
             commonFlow.approvedTypeStatementsDiff += toReplace
         }
 
-        for ((alias, underlyingVariable) in aliasedVariablesThatDontChangeAlias) {
+        for ((alias, underlyingVariable) in commonAliases) {
             addLocalVariableAlias(commonFlow, alias, underlyingVariable)
         }
         return commonFlow
     }
 
-    private fun computeAliasesThatDontChange(flows: Collection<PersistentFlow>): MutableMap<RealVariable, RealVariableAndType> {
-        val flowsSize = flows.size
-        val aliasedVariablesThatDontChangeAlias = mutableMapOf<RealVariable, RealVariableAndType>()
-
-        flows.flatMapTo(mutableSetOf()) { it.directAliasMap.keys }.forEach { aliasedVariable ->
-            val originals = flows.map { it.directAliasMap[aliasedVariable] ?: return@forEach }
-            if (originals.size != flowsSize) return@forEach
-            val firstOriginal = originals.first()
-            if (originals.all { it == firstOriginal }) {
-                aliasedVariablesThatDontChangeAlias[aliasedVariable] = firstOriginal
-            }
+    private fun computeCommonAliases(flows: Collection<PersistentFlow>): MutableMap<RealVariable, RealVariableAndType> =
+        flows.first().directAliasMap.filterTo(mutableMapOf()) { (variable, alias) ->
+            flows.all { it.directAliasMap[variable] == alias }
         }
 
-        return aliasedVariablesThatDontChangeAlias
-    }
-
     override fun addLocalVariableAlias(flow: PersistentFlow, alias: RealVariable, underlyingVariable: RealVariableAndType) {
         removeLocalVariableAlias(flow, alias)
         flow.directAliasMap = flow.directAliasMap.put(alias, underlyingVariable)