FIR checker: refactor ControlFlowInfos whose key is EdgeLabel

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/CfaUtils.kt

@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.contracts.description.EventOccurrencesRange
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.*
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
-import java.lang.IllegalStateException
 
 abstract class EventOccurrencesRangeInfo<E : EventOccurrencesRangeInfo<E, K>, K : Any>(
     map: PersistentMap<K, EventOccurrencesRange> = persistentMapOf()
@@ -39,6 +38,8 @@ class PropertyInitializationInfo(
     override val constructor: (PersistentMap<FirPropertySymbol, EventOccurrencesRange>) -> PropertyInitializationInfo =
         ::PropertyInitializationInfo
 
+    override val empty: () -> PropertyInitializationInfo =
+        ::EMPTY
 }
 
 class LocalPropertyCollector private constructor() : ControlFlowGraphVisitorVoid() {
@@ -61,7 +62,7 @@ class LocalPropertyCollector private constructor() : ControlFlowGraphVisitorVoid
 
 class PathAwarePropertyInitializationInfo(
     map: PersistentMap<EdgeLabel, PropertyInitializationInfo> = persistentMapOf()
-) : ControlFlowInfo<PathAwarePropertyInitializationInfo, EdgeLabel, PropertyInitializationInfo>(map) {
+) : PathAwareControlFlowInfo<PathAwarePropertyInitializationInfo, PropertyInitializationInfo>(map) {
     companion object {
         val EMPTY = PathAwarePropertyInitializationInfo(persistentMapOf(NormalPath to PropertyInitializationInfo.EMPTY))
     }
@@ -69,70 +70,8 @@ class PathAwarePropertyInitializationInfo(
     override val constructor: (PersistentMap<EdgeLabel, PropertyInitializationInfo>) -> PathAwarePropertyInitializationInfo =
         ::PathAwarePropertyInitializationInfo
 
-    val infoAtNormalPath: PropertyInitializationInfo
+    override val empty: () -> PathAwarePropertyInitializationInfo =
-        get() = map[NormalPath] ?: PropertyInitializationInfo.EMPTY
+        ::EMPTY
-
-    val hasNormalPath: Boolean
-        get() = map.containsKey(NormalPath)
-
-    fun applyLabel(node: CFGNode<*>, label: EdgeLabel): PathAwarePropertyInitializationInfo {
-        if (label.isNormal) {
-            // Special case: when we exit the try expression, null label means a normal path.
-            // Filter out any info bound to non-null label
-            // One day, if we allow multiple edges between nodes with different labels, e.g., labeling all paths in try/catch/finally,
-            // instead of this kind of special handling, proxy enter/exit nodes per label are preferred.
-            if (node is TryExpressionExitNode) {
-                return if (hasNormalPath) {
-                    constructor(persistentMapOf(NormalPath to infoAtNormalPath))
-                } else {
-                    /* This means no info for normal path. */
-                    EMPTY
-                }
-            }
-            // In general, null label means no additional path info, hence return `this` as-is.
-            return this
-        }
-
-        val hasAbnormalLabels = map.keys.any { !it.isNormal }
-        return if (hasAbnormalLabels) {
-            // { |-> ... l1 |-> I1, l2 |-> I2, ... }
-            //   | l1         // path exit: if the given info has non-null labels, this acts like a filtering
-            // { |-> I1 }     // NB: remove the path info
-            if (map.keys.contains(label)) {
-                constructor(persistentMapOf(NormalPath to map[label]!!))
-            } else {
-                /* This means no info for the specific label. */
-                EMPTY
-            }
-        } else {
-            // { |-> ... }    // empty path info
-            //   | l1         // path entry
-            // { l1 -> ... }  // now, every info bound to the label
-            constructor(persistentMapOf(label to infoAtNormalPath))
-        }
-    }
-
-    override fun merge(other: PathAwarePropertyInitializationInfo): PathAwarePropertyInitializationInfo {
-        var resultMap = persistentMapOf<EdgeLabel, PropertyInitializationInfo>()
-        for (label in keys.union(other.keys)) {
-            // disjoint merging to preserve paths. i.e., merge the property initialization info if and only if both have the key.
-            // merge({ |-> I1 }, { |-> I2, l1 |-> I3 }
-            //   == { |-> merge(I1, I2), l1 |-> I3 }
-            val i1 = this[label]
-            val i2 = other[label]
-            resultMap = when {
-                i1 != null && i2 != null ->
-                    resultMap.put(label, i1.merge(i2))
-                i1 != null ->
-                    resultMap.put(label, i1)
-                i2 != null ->
-                    resultMap.put(label, i2)
-                else ->
-                    throw IllegalStateException()
-            }
-        }
-        return constructor(resultMap)
-    }
 }
 
 class PropertyInitializationInfoCollector(private val localProperties: Set<FirPropertySymbol>) :

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/ControlFlowInfo.kt

@@ -13,6 +13,8 @@ abstract class ControlFlowInfo<S : ControlFlowInfo<S, K, V>, K : Any, V : Any> p
 
     protected abstract val constructor: (PersistentMap<K, V>) -> S
 
+    protected abstract val empty: () -> S
+
     override fun equals(other: Any?): Boolean {
         return map == (other as? ControlFlowInfo<*, *, *>)?.map
     }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/FirCallsEffectAnalyzer.kt

@@ -38,7 +38,6 @@ import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.coneTypeSafe
 import org.jetbrains.kotlin.utils.addIfNotNull
-import java.lang.IllegalStateException
 import kotlin.contracts.ExperimentalContracts
 import kotlin.contracts.contract
 
@@ -213,11 +212,13 @@ object FirCallsEffectAnalyzer : FirControlFlowChecker() {
         override val constructor: (PersistentMap<FirBasedSymbol<*>, EventOccurrencesRange>) -> LambdaInvocationInfo =
             ::LambdaInvocationInfo
 
+        override val empty: () -> LambdaInvocationInfo =
+            ::EMPTY
     }
 
     class PathAwareLambdaInvocationInfo(
         map: PersistentMap<EdgeLabel, LambdaInvocationInfo> = persistentMapOf()
-    ) : ControlFlowInfo<PathAwareLambdaInvocationInfo, EdgeLabel, LambdaInvocationInfo>(map) {
+    ) : PathAwareControlFlowInfo<PathAwareLambdaInvocationInfo, LambdaInvocationInfo>(map) {
         companion object {
             val EMPTY = PathAwareLambdaInvocationInfo(persistentMapOf(NormalPath to LambdaInvocationInfo.EMPTY))
         }
@@ -225,70 +226,8 @@ object FirCallsEffectAnalyzer : FirControlFlowChecker() {
         override val constructor: (PersistentMap<EdgeLabel, LambdaInvocationInfo>) -> PathAwareLambdaInvocationInfo =
             ::PathAwareLambdaInvocationInfo
 
-        val infoAtNormalPath: LambdaInvocationInfo
+        override val empty: () -> PathAwareLambdaInvocationInfo =
-            get() = map[NormalPath] ?: LambdaInvocationInfo.EMPTY
+            ::EMPTY
-
-        val hasNormalPath: Boolean
-            get() = map.containsKey(NormalPath)
-
-        fun applyLabel(node: CFGNode<*>, label: EdgeLabel): PathAwareLambdaInvocationInfo {
-            if (label.isNormal) {
-                // Special case: when we exit the try expression, null label means a normal path.
-                // Filter out any info bound to non-null label
-                // One day, if we allow multiple edges between nodes with different labels, e.g., labeling all paths in try/catch/finally,
-                // instead of this kind of special handling, proxy enter/exit nodes per label are preferred.
-                if (node is TryExpressionExitNode) {
-                    return if (hasNormalPath) {
-                        constructor(persistentMapOf(NormalPath to infoAtNormalPath))
-                    } else {
-                        /* This means no info for normal path. */
-                        EMPTY
-                    }
-                }
-                // In general, null label means no additional path info, hence return `this` as-is.
-                return this
-            }
-
-            val hasAbnormalLabels = map.keys.any { !it.isNormal }
-            return if (hasAbnormalLabels) {
-                // { |-> ... l1 |-> I1, l2 |-> I2, ... }
-                //   | l1         // path exit: if the given info has non-null labels, this acts like a filtering
-                // { |-> I1 }     // NB: remove the path info
-                if (map.keys.contains(label)) {
-                    constructor(persistentMapOf(NormalPath to map[label]!!))
-                } else {
-                    /* This means no info for the specific label. */
-                    EMPTY
-                }
-            } else {
-                // { |-> ... }    // empty path info
-                //   | l1         // path entry
-                // { l1 -> ... }  // now, every info bound to the label
-                constructor(persistentMapOf(label to infoAtNormalPath))
-            }
-        }
-
-        override fun merge(other: PathAwareLambdaInvocationInfo): PathAwareLambdaInvocationInfo {
-            var resultMap = persistentMapOf<EdgeLabel, LambdaInvocationInfo>()
-            for (label in keys.union(other.keys)) {
-                // disjoint merging to preserve paths. i.e., merge the property initialization info if and only if both have the key.
-                // merge({ |-> I1 }, { |-> I2, l1 |-> I3 }
-                //   == { |-> merge(I1, I2), l1 |-> I3 }
-                val i1 = this[label]
-                val i2 = other[label]
-                resultMap = when {
-                    i1 != null && i2 != null ->
-                        resultMap.put(label, i1.merge(i2))
-                    i1 != null ->
-                        resultMap.put(label, i1)
-                    i2 != null ->
-                        resultMap.put(label, i2)
-                    else ->
-                        throw IllegalStateException()
-                }
-            }
-            return constructor(resultMap)
-        }
     }
 
     private class InvocationDataCollector(

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/PathAwareControlFlowInfo.kt

@@ -0,0 +1,84 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.analysis.cfa
+
+import kotlinx.collections.immutable.PersistentMap
+import kotlinx.collections.immutable.persistentMapOf
+import org.jetbrains.kotlin.fir.resolve.dfa.cfg.CFGNode
+import org.jetbrains.kotlin.fir.resolve.dfa.cfg.EdgeLabel
+import org.jetbrains.kotlin.fir.resolve.dfa.cfg.NormalPath
+import org.jetbrains.kotlin.fir.resolve.dfa.cfg.TryExpressionExitNode
+
+abstract class PathAwareControlFlowInfo<P : PathAwareControlFlowInfo<P, S>, S : ControlFlowInfo<S, *, *>>(
+    map: PersistentMap<EdgeLabel, S>,
+) : ControlFlowInfo<P, EdgeLabel, S>(map) {
+
+    internal val infoAtNormalPath: S
+        get() = map.getValue(NormalPath)
+
+    private val hasNormalPath: Boolean
+        get() = map.containsKey(NormalPath)
+
+    fun applyLabel(node: CFGNode<*>, label: EdgeLabel): P {
+        if (label.isNormal) {
+            // Special case: when we exit the try expression, null label means a normal path.
+            // Filter out any info bound to non-null label
+            // One day, if we allow multiple edges between nodes with different labels, e.g., labeling all paths in try/catch/finally,
+            // instead of this kind of special handling, proxy enter/exit nodes per label are preferred.
+            if (node is TryExpressionExitNode) {
+                return if (hasNormalPath) {
+                    constructor(persistentMapOf(NormalPath to infoAtNormalPath))
+                } else {
+                    /* This means no info for normal path. */
+                    empty()
+                }
+            }
+            // In general, null label means no additional path info, hence return `this` as-is.
+            @Suppress("UNCHECKED_CAST")
+            return this as P
+        }
+
+        val hasAbnormalLabels = map.keys.any { !it.isNormal }
+        return if (hasAbnormalLabels) {
+            // { |-> ... l1 |-> I1, l2 |-> I2, ... }
+            //   | l1         // path exit: if the given info has non-null labels, this acts like a filtering
+            // { |-> I1 }     // NB: remove the path info
+            if (map.keys.contains(label)) {
+                constructor(persistentMapOf(NormalPath to map[label]!!))
+            } else {
+                /* This means no info for the specific label. */
+                empty()
+            }
+        } else {
+            // { |-> ... }    // empty path info
+            //   | l1         // path entry
+            // { l1 -> ... }  // now, every info bound to the label
+            constructor(persistentMapOf(label to infoAtNormalPath))
+        }
+    }
+
+    override fun merge(other: P): P {
+        var resultMap = persistentMapOf<EdgeLabel, S>()
+        for (label in keys.union(other.keys)) {
+            // disjoint merging to preserve paths. i.e., merge the property initialization info if and only if both have the key.
+            // merge({ |-> I1 }, { |-> I2, l1 |-> I3 })
+            //   == { |-> merge(I1, I2), l1 |-> I3 }
+            val i1 = this[label]
+            val i2 = other[label]
+            resultMap = when {
+                i1 != null && i2 != null ->
+                    resultMap.put(label, i1.merge(i2))
+                i1 != null ->
+                    resultMap.put(label, i1)
+                i2 != null ->
+                    resultMap.put(label, i2)
+                else ->
+                    throw IllegalStateException()
+            }
+        }
+        return constructor(resultMap)
+    }
+}

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/extended/UnusedChecker.kt

@@ -106,6 +106,9 @@ object UnusedChecker : FirControlFlowChecker() {
         override val constructor: (PersistentMap<FirPropertySymbol, VariableStatus>) -> VariableStatusInfo =
             ::VariableStatusInfo
 
+        override val empty: () -> VariableStatusInfo =
+            ::EMPTY
+
         override fun merge(other: VariableStatusInfo): VariableStatusInfo {
             var result = this
             for (symbol in keys.union(other.keys)) {