Rewrite DelegatingDataFlowInfo using immutable collections from javaslang

It may be useful because data-flow info is mostly the same for
neighbouring expressions and changes between them should be minor.

Thus, it should help to avoid making copies of huge maps and reusing
them instead

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/smartcasts/DataFlowInfo.kt

@@ -16,9 +16,10 @@
 
 package org.jetbrains.kotlin.resolve.calls.smartcasts
 
-import com.google.common.collect.SetMultimap
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.util.javaslang.ImmutableMap
+import org.jetbrains.kotlin.util.javaslang.ImmutableSet
 
 /**
  * This interface is intended to provide and edit information about value nullabilities and possible types.
@@ -26,9 +27,9 @@ import org.jetbrains.kotlin.types.KotlinType
  */
 interface DataFlowInfo {
 
-    val completeNullabilityInfo: Map<DataFlowValue, Nullability>
+    val completeNullabilityInfo: ImmutableMap<DataFlowValue, Nullability>
 
-    val completeTypeInfo: SetMultimap<DataFlowValue, KotlinType>
+    val completeTypeInfo: ImmutableMap<DataFlowValue, ImmutableSet<KotlinType>>
 
     /**
      * Returns collected nullability for the given value, NOT taking its stability into account.

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/smartcasts/DelegatingDataFlowInfo.kt

@@ -18,6 +18,7 @@ package org.jetbrains.kotlin.resolve.calls.smartcasts
 
 import com.google.common.collect.LinkedHashMultimap
 import com.google.common.collect.SetMultimap
+import javaslang.Tuple2
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.LanguageVersionSettings
@@ -26,25 +27,35 @@ import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.checker.NewCapturedTypeConstructor
 import org.jetbrains.kotlin.types.typeUtil.contains
 import org.jetbrains.kotlin.types.typeUtil.isSubtypeOf
+import org.jetbrains.kotlin.util.javaslang.*
 import org.jetbrains.kotlin.utils.newLinkedHashSetWithExpectedSize
 import java.util.*
 
+private typealias ImmutableMultimap<K, V> = ImmutableMap<K, ImmutableSet<V>>
+
+private fun <K, V> ImmutableMultimap<K, V>.put(key: K, value: V): ImmutableMultimap<K, V> {
+    val oldSet = this[key].getOrElse(ImmutableLinkedHashSet.empty<V>())
+    if (oldSet.contains(value)) return this
+
+    return put(key, oldSet.add(value))
+}
+
 internal class DelegatingDataFlowInfo private constructor(
-    override val completeNullabilityInfo: Map<DataFlowValue, Nullability>,
+    override val completeNullabilityInfo: ImmutableMap<DataFlowValue, Nullability>,
-    override val completeTypeInfo: SetMultimap<DataFlowValue, KotlinType>
+    override val completeTypeInfo: ImmutableMultimap<DataFlowValue, KotlinType>
 ) : DataFlowInfo {
 
-    constructor() : this(emptyMap(), newTypeInfo())
+    constructor() : this(EMPTY_NULLABILITY_INFO, EMPTY_TYPE_INFO)
 
     override fun getCollectedNullability(key: DataFlowValue) = getNullability(key, false)
 
     override fun getStableNullability(key: DataFlowValue) = getNullability(key, true)
 
-    private fun getNullability(key: DataFlowValue, stableOnly: Boolean) =
+    private fun getNullability(key: DataFlowValue, stableOnly: Boolean): Nullability =
         if (stableOnly && !key.isStable) {
             key.immanentNullability
         } else {
-            completeNullabilityInfo.getOrDefault(key, key.immanentNullability)
+            completeNullabilityInfo[key].getOrElse(key.immanentNullability)
         }
 
     private fun putNullabilityAndTypeInfo(
@@ -52,14 +63,14 @@ internal class DelegatingDataFlowInfo private constructor(
         value: DataFlowValue,
         nullability: Nullability,
         languageVersionSettings: LanguageVersionSettings,
-        typeInfo: SetMultimap<DataFlowValue, KotlinType>? = null,
+        newTypeInfoBuilder: SetMultimap<DataFlowValue, KotlinType>? = null,
         affectReceiver: Boolean = true,
         // TODO: remove me in version 1.3! I'm very dirty hack!
         // In normal circumstances this should be always true
         recordUnstable: Boolean = true
     ): Boolean {
         if (value.isStable || recordUnstable) {
-            map.put(value, nullability)
+            map[value] = nullability
         }
 
         val identifierInfo = value.identifierInfo
@@ -72,7 +83,7 @@ internal class DelegatingDataFlowInfo private constructor(
                         val receiverValue = DataFlowValue(identifierInfo.receiverInfo, receiverType)
                         putNullabilityAndTypeInfo(
                             map, receiverValue, nullability,
-                            languageVersionSettings, typeInfo, recordUnstable = recordUnstable
+                            languageVersionSettings, newTypeInfoBuilder, recordUnstable = recordUnstable
                         )
                     }
                 }
@@ -85,17 +96,17 @@ internal class DelegatingDataFlowInfo private constructor(
                         val subjectValue = DataFlowValue(identifierInfo.subjectInfo, subjectType)
                         putNullabilityAndTypeInfo(
                             map, subjectValue, nullability,
-                            languageVersionSettings, typeInfo, recordUnstable = false
+                            languageVersionSettings, newTypeInfoBuilder, recordUnstable = false
                         )
                         if (subjectValue.isStable) {
-                            typeInfo?.put(subjectValue, targetType)
+                            newTypeInfoBuilder?.put(subjectValue, targetType)
                         }
                     }
                 }
                 is IdentifierInfo.Variable -> identifierInfo.bound?.let {
                     putNullabilityAndTypeInfo(
                         map, it, nullability,
-                        languageVersionSettings, typeInfo, recordUnstable = recordUnstable
+                        languageVersionSettings, newTypeInfoBuilder, recordUnstable = recordUnstable
                     )
                 }
             }
@@ -112,12 +123,12 @@ internal class DelegatingDataFlowInfo private constructor(
         enrichWithNotNull: Boolean,
         languageVersionSettings: LanguageVersionSettings
     ): Set<KotlinType> {
-        val types = completeTypeInfo[key]
+        val types = completeTypeInfo[key].getOrElse(ImmutableLinkedHashSet.empty())
         if (!enrichWithNotNull || getCollectedNullability(key).canBeNull()) {
-            return types
+            return types.toJavaSet()
         }
 
-        val enrichedTypes = newLinkedHashSetWithExpectedSize<KotlinType>(types.size + 1)
+        val enrichedTypes = newLinkedHashSetWithExpectedSize<KotlinType>(types.size() + 1)
         val originalType = key.type
         types.mapTo(enrichedTypes) { type -> type.makeReallyNotNullIfNeeded(languageVersionSettings) }
         if (originalType.canBeDefinitelyNotNullOrNotNull(languageVersionSettings)) {
@@ -164,7 +175,6 @@ internal class DelegatingDataFlowInfo private constructor(
         val nullabilityOfB = getStableNullability(b)
         putNullabilityAndTypeInfo(nullability, a, nullabilityOfB, languageVersionSettings, affectReceiver = false)
 
-        val newTypeInfo = newTypeInfo()
         var typesForB = getStableTypes(b, languageVersionSettings)
         // Own type of B must be recorded separately, e.g. for a constant
         // But if its type is the same as A, there is no reason to do it
@@ -173,9 +183,8 @@ internal class DelegatingDataFlowInfo private constructor(
         if (!b.type.isError && a.type != b.type) {
             typesForB += b.type
         }
-        newTypeInfo.putAll(a, typesForB)
 
-        return create(this, nullability, if (newTypeInfo.isEmpty) EMPTY_TYPE_INFO else newTypeInfo, a)
+        return create(this, nullability, listOf(Tuple2(a, typesForB)), a)
     }
 
     override fun equate(
@@ -185,40 +194,40 @@ internal class DelegatingDataFlowInfo private constructor(
         val nullabilityOfA = getStableNullability(a)
         val nullabilityOfB = getStableNullability(b)
 
-        val newTypeInfo = newTypeInfo()
+        val newTypeInfoBuilder = newTypeInfoBuilder()
         var changed =
             putNullabilityAndTypeInfo(
                 resultNullabilityInfo,
                 a,
                 nullabilityOfA.refine(nullabilityOfB),
                 languageVersionSettings,
-                newTypeInfo
+                newTypeInfoBuilder
             ) or
                     putNullabilityAndTypeInfo(
                         resultNullabilityInfo,
                         b,
                         nullabilityOfB.refine(nullabilityOfA),
                         languageVersionSettings,
-                        newTypeInfo
+                        newTypeInfoBuilder
                     )
 
         // NB: == has no guarantees of type equality, see KT-11280 for the example
         if (identityEquals || !nullabilityOfA.canBeNonNull() || !nullabilityOfB.canBeNonNull()) {
-            newTypeInfo.putAll(a, getStableTypes(b, false, languageVersionSettings))
+            newTypeInfoBuilder.putAll(a, getStableTypes(b, false, languageVersionSettings))
-            newTypeInfo.putAll(b, getStableTypes(a, false, languageVersionSettings))
+            newTypeInfoBuilder.putAll(b, getStableTypes(a, false, languageVersionSettings))
             if (a.type != b.type) {
                 // To avoid recording base types of own type
                 if (!a.type.isSubtypeOf(b.type)) {
-                    newTypeInfo.put(a, b.type)
+                    newTypeInfoBuilder.put(a, b.type)
                 }
                 if (!b.type.isSubtypeOf(a.type)) {
-                    newTypeInfo.put(b, a.type)
+                    newTypeInfoBuilder.put(b, a.type)
                 }
             }
-            changed = changed or !newTypeInfo.isEmpty
+            changed = changed or !newTypeInfoBuilder.isEmpty
         }
 
-        return if (changed) create(this, resultNullabilityInfo, if (newTypeInfo.isEmpty) EMPTY_TYPE_INFO else newTypeInfo) else this
+        return if (changed) create(this, resultNullabilityInfo, newTypeInfoBuilder) else this
     }
 
     override fun disequate(
@@ -228,24 +237,24 @@ internal class DelegatingDataFlowInfo private constructor(
         val nullabilityOfA = getStableNullability(a)
         val nullabilityOfB = getStableNullability(b)
 
-        val newTypeInfo = newTypeInfo()
+        val newTypeInfoBuilder = newTypeInfoBuilder()
         val changed =
             putNullabilityAndTypeInfo(
                 resultNullabilityInfo,
                 a,
                 nullabilityOfA.refine(nullabilityOfB.invert()),
                 languageVersionSettings,
-                newTypeInfo
+                newTypeInfoBuilder
             ) or
                     putNullabilityAndTypeInfo(
                         resultNullabilityInfo,
                         b,
                         nullabilityOfB.refine(nullabilityOfA.invert()),
                         languageVersionSettings,
-                        newTypeInfo
+                        newTypeInfoBuilder
                     )
 
-        return if (changed) create(this, resultNullabilityInfo, if (newTypeInfo.isEmpty) EMPTY_TYPE_INFO else newTypeInfo) else this
+        return if (changed) create(this, resultNullabilityInfo, newTypeInfoBuilder) else this
 
     }
 
@@ -255,13 +264,17 @@ internal class DelegatingDataFlowInfo private constructor(
         if (value.type == type) return this
         if (getCollectedTypes(value, languageVersionSettings).contains(type)) return this
         if (!value.type.isFlexible() && value.type.isSubtypeOf(type)) return this
-        val newTypeInfo = newTypeInfo()
+
-        newTypeInfo.put(value, type)
         val nullabilityInfo = hashMapOf<DataFlowValue, Nullability>()
         if (!type.isMarkedNullable) {
             putNullabilityAndTypeInfo(nullabilityInfo, value, NOT_NULL, languageVersionSettings)
         }
-        return create(this, if (type.isMarkedNullable) emptyMap() else nullabilityInfo, newTypeInfo)
+
+        return create(
+            this,
+            nullabilityInfo,
+            listOf(Tuple2(value, listOf(type)))
+        )
     }
 
     override fun and(other: DataFlowInfo): DataFlowInfo {
@@ -280,24 +293,27 @@ internal class DelegatingDataFlowInfo private constructor(
             }
         }
 
-        val myTypeInfo = completeTypeInfo
         val otherTypeInfo = other.completeTypeInfo
-        if (resultNullabilityInfo.isEmpty() && containsAll(myTypeInfo, otherTypeInfo)) {
-            return this
-        }
 
         return create(this, resultNullabilityInfo, otherTypeInfo)
     }
 
-    private fun Set<KotlinType>.containsNothing() = any { KotlinBuiltIns.isNothing(it) }
+    private fun ImmutableSet<KotlinType>?.containsNothing() = this?.any { KotlinBuiltIns.isNothing(it) } ?: false
 
-    private fun Set<KotlinType>.intersectConsideringNothing(other: Set<KotlinType>) =
+    private fun ImmutableSet<KotlinType>?.intersectConsideringNothing(other: ImmutableSet<KotlinType>?) =
         when {
             other.containsNothing() -> this
             this.containsNothing() -> other
             else -> this.intersect(other)
         }
 
+    private fun ImmutableSet<KotlinType>?.intersect(other: ImmutableSet<KotlinType>?): ImmutableSet<KotlinType> =
+        when {
+            this == null -> other ?: ImmutableLinkedHashSet.empty()
+            other == null -> this
+            else -> this.intersect(other)
+        }
+
     override fun or(other: DataFlowInfo): DataFlowInfo {
         if (other === DataFlowInfo.EMPTY) return DataFlowInfo.EMPTY
         if (this === DataFlowInfo.EMPTY) return DataFlowInfo.EMPTY
@@ -313,62 +329,77 @@ internal class DelegatingDataFlowInfo private constructor(
 
         val myTypeInfo = completeTypeInfo
         val otherTypeInfo = other.completeTypeInfo
-        val newTypeInfo = newTypeInfo()
+        val newTypeInfoBuilder = newTypeInfoBuilder()
 
         for (key in myTypeInfo.keySet()) {
             if (key in otherTypeInfo.keySet()) {
-                newTypeInfo.putAll(key, myTypeInfo[key].intersectConsideringNothing(otherTypeInfo[key]))
+                newTypeInfoBuilder.putAll(
+                    key,
+                    myTypeInfo[key].getOrNull().intersectConsideringNothing(otherTypeInfo[key].getOrNull())
+                        ?: ImmutableLinkedHashSet.empty()
+                )
             }
         }
-        return create(null, resultNullabilityInfo, newTypeInfo)
+        return create(null, resultNullabilityInfo, newTypeInfoBuilder)
     }
 
     override fun toString() = if (completeTypeInfo.isEmpty && completeNullabilityInfo.isEmpty()) "EMPTY" else "Non-trivial DataFlowInfo"
 
     companion object {
-        private val EMPTY_TYPE_INFO = newTypeInfo()
+        private val EMPTY_NULLABILITY_INFO: ImmutableMap<DataFlowValue, Nullability> =
+            ImmutableHashMap.empty()
 
-        private fun containsAll(first: SetMultimap<DataFlowValue, KotlinType>, second: SetMultimap<DataFlowValue, KotlinType>) =
+        private val EMPTY_TYPE_INFO: ImmutableMultimap<DataFlowValue, KotlinType> =
-            first.entries().containsAll(second.entries())
+            ImmutableHashMap.empty()
 
-        fun newTypeInfo(): SetMultimap<DataFlowValue, KotlinType> = LinkedHashMultimap.create<DataFlowValue, KotlinType>()
+        private fun newTypeInfoBuilder(): SetMultimap<DataFlowValue, KotlinType> =
+            LinkedHashMultimap.create()
+
+        private fun create(
+            parent: DataFlowInfo?,
+            updatedNullabilityInfo: Map<DataFlowValue, Nullability>,
+            updatedTypeInfo: SetMultimap<DataFlowValue, KotlinType>
+        ): DataFlowInfo =
+            create(
+                parent,
+                updatedNullabilityInfo,
+                updatedTypeInfo.asMap().entries.map { Tuple2(it.key, it.value) }
+            )
 
         private fun create(
             parent: DataFlowInfo?,
             updatedNullabilityInfo: Map<DataFlowValue, Nullability>,
             // NB: typeInfo must be mutable here!
-            updatedTypeInfo: SetMultimap<DataFlowValue, KotlinType>,
+            updatedTypeInfo: Iterable<Tuple2<DataFlowValue, out Iterable<KotlinType>>>,
             valueToClearPreviousTypeInfo: DataFlowValue? = null
         ): DataFlowInfo {
-            val toDelete = newTypeInfo()
+            if (updatedNullabilityInfo.isEmpty() && updatedTypeInfo.none() && valueToClearPreviousTypeInfo == null) {
-            for (value in updatedTypeInfo.keys()) {
+                return parent ?: DataFlowInfo.EMPTY
-                for (type in updatedTypeInfo[value]) {
+            }
-                    // Remove original type (see also KT-10666)
+
-                    if (value.type == type || type.contains { it.constructor is NewCapturedTypeConstructor }) {
+            val resultingNullabilityInfo =
-                        toDelete.put(value, type)
+                updatedNullabilityInfo.entries.fold(
-                    }
+                    parent?.completeNullabilityInfo ?: EMPTY_NULLABILITY_INFO
+                ) { result, (dataFlowValue, nullability) ->
+                    result.put(dataFlowValue, nullability)
+                }
+
+            var resultingTypeInfo = parent?.completeTypeInfo ?: EMPTY_TYPE_INFO
+
+            valueToClearPreviousTypeInfo?.let {
+                resultingTypeInfo = resultingTypeInfo.put(it, ImmutableLinkedHashSet.empty())
+            }
+
+            for ((value, types) in updatedTypeInfo) {
+                for (type in types) {
+                    if (value.type == type || type.contains { it.constructor is NewCapturedTypeConstructor }) continue
+                    resultingTypeInfo = resultingTypeInfo.put(value, type)
                 }
             }
-            for ((value, type) in toDelete.entries()) {
-                updatedTypeInfo.remove(value, type)
-            }
 
-            if (updatedNullabilityInfo.isEmpty() && updatedTypeInfo.isEmpty && valueToClearPreviousTypeInfo == null) {
+            if (resultingNullabilityInfo.isEmpty && resultingTypeInfo.isEmpty) return DataFlowInfo.EMPTY
-                return parent ?: DataFlowInfoFactory.EMPTY
+            if (resultingNullabilityInfo === parent?.completeNullabilityInfo && resultingTypeInfo === parent.completeTypeInfo) {
-            }
+                return parent
-
-            val resultingNullabilityInfo = parent?.completeNullabilityInfo?.toMutableMap() ?: hashMapOf()
-            resultingNullabilityInfo.putAll(updatedNullabilityInfo)
-
-            val resultingTypeInfo =
-                LinkedHashMultimap.create<DataFlowValue, KotlinType>(
-                    parent?.completeTypeInfo ?: newTypeInfo()
-                )
-
-            valueToClearPreviousTypeInfo?.let(resultingTypeInfo::removeAll)
-
-            for ((value, types) in updatedTypeInfo.asMap().entries) {
-                resultingTypeInfo.putAll(value, types)
             }
 
             return DelegatingDataFlowInfo(resultingNullabilityInfo, resultingTypeInfo)

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/PreliminaryLoopVisitor.kt

@@ -38,7 +38,7 @@ class PreliminaryLoopVisitor private constructor() : AssignedVariablesSearcher()
         var resultFlowInfo = dataFlowInfo
         val nullabilityMap = resultFlowInfo.completeNullabilityInfo
         val valueSetToClear = LinkedHashSet<DataFlowValue>()
-        for (value in nullabilityMap.keys) {
+        for (value in nullabilityMap.keySet()) {
             // Only stable variables are under interest here
             val identifierInfo = value.identifierInfo
             if (value.kind == DataFlowValue.Kind.STABLE_VARIABLE && identifierInfo is IdentifierInfo.Variable) {

idea/idea-core/src/org/jetbrains/kotlin/idea/core/SmartCastCalculator.kt

@@ -27,7 +27,6 @@ import org.jetbrains.kotlin.idea.util.getImplicitReceiversWithInstance
 import org.jetbrains.kotlin.idea.util.getResolutionScope
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.resolve.BindingContext
-import org.jetbrains.kotlin.resolve.bindingContextUtil.getDataFlowInfoAfter
 import org.jetbrains.kotlin.resolve.bindingContextUtil.getDataFlowInfoBefore
 import org.jetbrains.kotlin.resolve.calls.smartcasts.*
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope
@@ -35,6 +34,8 @@ import org.jetbrains.kotlin.resolve.scopes.receivers.ImplicitReceiver
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.typeUtil.isSubtypeOf
 import org.jetbrains.kotlin.types.typeUtil.makeNotNullable
+import org.jetbrains.kotlin.util.javaslang.component1
+import org.jetbrains.kotlin.util.javaslang.component2
 import java.util.*
 
 class SmartCastCalculator(
@@ -119,10 +120,10 @@ class SmartCastCalculator(
 
         val entityToInfo = HashMap<Any, SmartCastInfo>()
 
-        for ((dataFlowValue, types) in dataFlowInfo.completeTypeInfo.asMap().entries) {
+        for ((dataFlowValue, types) in dataFlowInfo.completeTypeInfo) {
             val entity = dataFlowValueToEntity.invoke(dataFlowValue)
             if (entity != null) {
-                entityToInfo[entity] = SmartCastInfo(types, false)
+                entityToInfo[entity] = SmartCastInfo(types.toJavaList(), false)
             }
         }
 

idea/tests/org/jetbrains/kotlin/DataFlowValueRenderingTest.kt

@@ -66,7 +66,7 @@ abstract class AbstractDataFlowValueRenderingTest: KotlinLightCodeInsightFixture
         val expression = element.getStrictParentOfType<KtExpression>()!!
         val info = expression.analyze().getDataFlowInfoAfter(expression)
 
-        val allValues = (info.completeTypeInfo.keySet() + info.completeNullabilityInfo.keys).toSet()
+        val allValues = (info.completeTypeInfo.keySet() + info.completeNullabilityInfo.keySet()).toSet()
         val actual = allValues.mapNotNull { it.render() }.sorted().joinToString("\n")
 
         KotlinTestUtils.assertEqualsToFile(File(FileUtil.getNameWithoutExtension(fileName) + ".txt"), actual)