FIR CFA: remove a redundant wrapper around PersistentMap

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

@@ -85,7 +85,6 @@ object FirCallsEffectAnalyzer : FirControlFlowChecker() {
 
         val invocationData = graph.collectDataForNode(
             TraverseDirection.Forward,
-            PathAwareLambdaInvocationInfo.EMPTY,
             InvocationDataCollector(functionalTypeEffects.keys.filterTo(mutableSetOf()) { it !in leakedSymbols })
         )
 
@@ -224,22 +223,15 @@ object FirCallsEffectAnalyzer : FirControlFlowChecker() {
             ::LambdaInvocationInfo
     }
 
-    class PathAwareLambdaInvocationInfo(
-        map: PersistentMap<EdgeLabel, LambdaInvocationInfo> = persistentMapOf()
-    ) : PathAwareControlFlowInfo<PathAwareLambdaInvocationInfo, LambdaInvocationInfo>(map) {
-        companion object {
-            val EMPTY = PathAwareLambdaInvocationInfo(persistentMapOf(NormalPath to LambdaInvocationInfo.EMPTY))
-        }
-
-        override val constructor: (PersistentMap<EdgeLabel, LambdaInvocationInfo>) -> PathAwareLambdaInvocationInfo =
-            ::PathAwareLambdaInvocationInfo
-    }
-
     private class InvocationDataCollector(
         val functionalTypeSymbols: Set<FirBasedSymbol<*>>
-    ) : PathAwareControlFlowGraphVisitor<PathAwareLambdaInvocationInfo>() {
+    ) : PathAwareControlFlowGraphVisitor<LambdaInvocationInfo>() {
+        companion object {
+            val EMPTY: PathAwareLambdaInvocationInfo = persistentMapOf(NormalPath to LambdaInvocationInfo.EMPTY)
+        }
+
         override val emptyInfo: PathAwareLambdaInvocationInfo
-            get() = PathAwareLambdaInvocationInfo.EMPTY
+            get() = EMPTY
 
         override fun visitFunctionCallNode(
             node: FunctionCallNode,
@@ -286,9 +278,7 @@ object FirCallsEffectAnalyzer : FirControlFlowChecker() {
             range: EventOccurrencesRange
         ): PathAwareLambdaInvocationInfo {
             val symbol = referenceToSymbol(reference)
-            return if (symbol != null) {
+            return if (symbol != null) addRange(this, symbol, range) else this
-                addRange(this, symbol, range, ::PathAwareLambdaInvocationInfo)
-            } else this
         }
     }
 
@@ -329,3 +319,5 @@ object FirCallsEffectAnalyzer : FirControlFlowChecker() {
         else -> null
     }
 }
+
+private typealias PathAwareLambdaInvocationInfo = PathAwareControlFlowInfo<FirCallsEffectAnalyzer.LambdaInvocationInfo>

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

@@ -33,41 +33,41 @@ fun ControlFlowGraph.traverse(
 
 // ---------------------- Path-sensitive data collection -----------------------
 
-fun <I : PathAwareControlFlowInfo<I, *>> ControlFlowGraph.collectDataForNode(
+fun <I : ControlFlowInfo<I, *, *>> ControlFlowGraph.collectDataForNode(
     direction: TraverseDirection,
-    initialInfo: I,
     visitor: PathAwareControlFlowGraphVisitor<I>,
     visitSubGraphs: Boolean = true
-): Map<CFGNode<*>, I> {
+): Map<CFGNode<*>, PathAwareControlFlowInfo<I>> {
-    val nodeMap = LinkedHashMap<CFGNode<*>, I>()
+    val nodeMap = LinkedHashMap<CFGNode<*>, PathAwareControlFlowInfo<I>>()
     val startNode = getEnterNode(direction)
-    nodeMap[startNode] = initialInfo
+    nodeMap[startNode] = visitor.emptyInfo
 
     var shouldContinue: Boolean
     do {
-        shouldContinue = collectDataForNodeInternal(direction, initialInfo, visitor, nodeMap, visitSubGraphs)
+        shouldContinue = collectDataForNodeInternal(direction, visitor, nodeMap, visitSubGraphs)
     } while (shouldContinue)
 
     return nodeMap
 }
 
-private fun <I : PathAwareControlFlowInfo<I, *>> ControlFlowGraph.collectDataForNodeInternal(
+private fun <I : ControlFlowInfo<I, *, *>> ControlFlowGraph.collectDataForNodeInternal(
     direction: TraverseDirection,
-    initialInfo: I,
     visitor: PathAwareControlFlowGraphVisitor<I>,
-    nodeMap: MutableMap<CFGNode<*>, I>,
+    nodeMap: MutableMap<CFGNode<*>, PathAwareControlFlowInfo<I>>,
     visitSubGraphs: Boolean = true
 ): Boolean {
     var changed = false
     val nodes = getNodesInOrder(direction)
     for (node in nodes) {
         if (visitSubGraphs && direction == TraverseDirection.Backward && node is CFGNodeWithSubgraphs<*>) {
-            node.subGraphs.forEach { changed = changed or it.collectDataForNodeInternal(direction, initialInfo, visitor, nodeMap) }
+            node.subGraphs.forEach { changed = changed or it.collectDataForNodeInternal(direction, visitor, nodeMap) }
         }
         val previousNodes = when (direction) {
             TraverseDirection.Forward -> node.previousCfgNodes
             TraverseDirection.Backward -> node.followingCfgNodes
         }
+        // TODO: if data for previousNodes hasn't changed, then should be no need to recompute data for this one
+        val union = node is UnionNodeMarker
         val previousData =
             previousNodes.mapNotNull {
                 val k = when (direction) {
@@ -76,20 +76,16 @@ private fun <I : PathAwareControlFlowInfo<I, *>> ControlFlowGraph.collectDataFor
                 }
                 val v = nodeMap[it] ?: return@mapNotNull null
                 visitor.visitEdge(it, node, k, v)
-            }
+            }.reduceOrNull { a, b -> a.join(b, union) }
-        val reduced = if (node is UnionNodeMarker)
-            previousData.reduceOrNull { a, b -> a.plus(b) }
-        else
-            previousData.reduceOrNull { a, b -> a.merge(b) }
         val data = nodeMap[node]
-        val newData = node.accept(visitor, reduced ?: visitor.emptyInfo)
+        val newData = node.accept(visitor, previousData ?: visitor.emptyInfo)
         val hasChanged = newData != data
         changed = changed or hasChanged
         if (hasChanged) {
             nodeMap[node] = newData
         }
         if (visitSubGraphs && direction == TraverseDirection.Forward && node is CFGNodeWithSubgraphs<*>) {
-            node.subGraphs.forEach { changed = changed or it.collectDataForNodeInternal(direction, initialInfo, visitor, nodeMap) }
+            node.subGraphs.forEach { changed = changed or it.collectDataForNodeInternal(direction, visitor, nodeMap) }
         }
     }
     return changed

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

@@ -5,18 +5,83 @@
 
 package org.jetbrains.kotlin.fir.analysis.cfa.util
 
-import org.jetbrains.kotlin.fir.resolve.dfa.cfg.CFGNode
+import kotlinx.collections.immutable.PersistentMap
-import org.jetbrains.kotlin.fir.resolve.dfa.cfg.ControlFlowGraphVisitor
+import kotlinx.collections.immutable.mutate
-import org.jetbrains.kotlin.fir.resolve.dfa.cfg.Edge
+import kotlinx.collections.immutable.persistentMapOf
-import org.jetbrains.kotlin.fir.resolve.dfa.cfg.UnionNodeMarker
+import org.jetbrains.kotlin.fir.resolve.dfa.cfg.*
 
-abstract class PathAwareControlFlowGraphVisitor<P : PathAwareControlFlowInfo<P, *>> : ControlFlowGraphVisitor<P, P>() {
+typealias PathAwareControlFlowInfo<I> = PersistentMap<EdgeLabel, I>
-    abstract val emptyInfo: P
 
-    open fun visitEdge(from: CFGNode<*>, to: CFGNode<*>, metadata: Edge, data: P): P =
+fun <I : ControlFlowInfo<I, *, *>> PathAwareControlFlowInfo<I>.join(
-        data.applyLabel(to, metadata.label) ?: emptyInfo
+    other: PathAwareControlFlowInfo<I>,
-
+    union: Boolean
-    override fun visitNode(node: CFGNode<*>, data: P): P = data
+): PathAwareControlFlowInfo<I> = mutate {
-
+    for ((label, rightValue) in other) {
-    override fun <T> visitUnionNode(node: T, data: P): P where T : CFGNode<*>, T : UnionNodeMarker = data
+        // 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 }
+        it[label] = this[label]?.let { leftValue ->
+            if (union) leftValue.plus(rightValue) else leftValue.merge(rightValue)
+        } ?: rightValue
+    }
 }
+
+abstract class PathAwareControlFlowGraphVisitor<I : ControlFlowInfo<I, *, *>> :
+    ControlFlowGraphVisitor<PathAwareControlFlowInfo<I>, PathAwareControlFlowInfo<I>>() {
+
+    abstract val emptyInfo: PathAwareControlFlowInfo<I>
+
+    open fun visitEdge(from: CFGNode<*>, to: CFGNode<*>, metadata: Edge, data: PathAwareControlFlowInfo<I>): PathAwareControlFlowInfo<I> {
+        val label = metadata.label
+        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 (to is TryExpressionExitNode) {
+                val infoAtNormalPath = data[NormalPath] ?: return emptyInfo
+                return persistentMapOf(NormalPath to infoAtNormalPath)
+            }
+            // In general, null label means no additional path info, hence return `this` as-is.
+            return data
+        }
+        return if (data.keys.any { !it.isNormal }) {
+            // { |-> ..., 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 label, except for uncaught exception path
+            val info = data[label] ?: return emptyInfo
+            if (label == UncaughtExceptionPath) {
+                // Special case: uncaught exception path, which still represents an uncaught exception path
+                // Target node is most likely fun/init exit, and we should keep info separated.
+                persistentMapOf(label to info)
+            } else {
+                // { |-> I }
+                //   | l1       // e.g., enter to proxy1 with l1
+                // { l1 -> I }
+                //   ...
+                // { |-> ..., l1 -> I', ... }
+                //   | l1       // e.g., exit proxy1 with l1
+                // { l1 -> I' }
+                persistentMapOf(NormalPath to info)
+            }
+        } else {
+            // { |-> ... }    // empty path info
+            //   | l1         // path entry
+            // { l1 -> ... }  // now, every info bound to the label
+            persistentMapOf(label to data.infoAtNormalPath)
+        }
+    }
+
+    override fun visitNode(
+        node: CFGNode<*>,
+        data: PathAwareControlFlowInfo<I>
+    ): PathAwareControlFlowInfo<I> = data
+
+    override fun <T> visitUnionNode(
+        node: T,
+        data: PathAwareControlFlowInfo<I>
+    ): PathAwareControlFlowInfo<I> where T : CFGNode<*>, T : UnionNodeMarker = data
+}
+
+internal val <I> PathAwareControlFlowInfo<I>.infoAtNormalPath: I
+    get() = getValue(NormalPath)

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

@@ -1,96 +0,0 @@
-/*
- * Copyright 2010-2021 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.util
-
-import kotlinx.collections.immutable.PersistentMap
-import kotlinx.collections.immutable.persistentMapOf
-import org.jetbrains.kotlin.fir.resolve.dfa.cfg.*
-
-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)
-
-    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) {
-                val infoAtNormalPath = map[NormalPath]
-                return if (infoAtNormalPath != null) {
-                    constructor(persistentMapOf(NormalPath to infoAtNormalPath))
-                } else {
-                    /* This means no info for normal path. */
-                    null
-                }
-            }
-            // 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 label, except for uncaught exception path
-            if (map.keys.contains(label)) {
-                if (label == UncaughtExceptionPath) {
-                    // Special case: uncaught exception path, which still represents an uncaught exception path
-                    // Target node is most likely fun/init exit, and we should keep info separated.
-                    constructor(persistentMapOf(label to map[label]!!))
-                } else {
-                    // { |-> I }
-                    //   | l1       // e.g., enter to proxy1 with l1
-                    // { l1 -> I }
-                    //   ...
-                    // { |-> ..., l1 -> I', ... }
-                    //   | l1       // e.g., exit proxy1 with l1
-                    // { l1 -> I' }
-                    constructor(persistentMapOf(NormalPath to map[label]!!))
-                }
-            } else {
-                /* This means no info for the specific label. */
-                null
-            }
-        } else {
-            // { |-> ... }    // empty path info
-            //   | l1         // path entry
-            // { l1 -> ... }  // now, every info bound to the label
-            constructor(persistentMapOf(label to infoAtNormalPath))
-        }
-    }
-
-    private fun join(other: P, union: Boolean): 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, if (union) i1.plus(i2) else i1.merge(i2))
-                i1 != null ->
-                    resultMap.put(label, i1)
-                i2 != null ->
-                    resultMap.put(label, i2)
-                else ->
-                    throw IllegalStateException()
-            }
-        }
-        return constructor(resultMap)
-    }
-
-    override fun merge(other: P): P = join(other, union = false)
-
-    override fun plus(other: P): P = join(other, union = true)
-}

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

@@ -8,8 +8,6 @@ package org.jetbrains.kotlin.fir.analysis.cfa.util
 import kotlinx.collections.immutable.PersistentMap
 import kotlinx.collections.immutable.persistentMapOf
 import org.jetbrains.kotlin.contracts.description.EventOccurrencesRange
-import org.jetbrains.kotlin.fir.resolve.dfa.cfg.EdgeLabel
-import org.jetbrains.kotlin.fir.resolve.dfa.cfg.NormalPath
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
 
 abstract class EventOccurrencesRangeInfo<E : EventOccurrencesRangeInfo<E, K>, K : Any>(
@@ -51,13 +49,4 @@ class PropertyInitializationInfo(
         ::PropertyInitializationInfo
 }
 
-class PathAwarePropertyInitializationInfo(
+typealias PathAwarePropertyInitializationInfo = PathAwareControlFlowInfo<PropertyInitializationInfo>
-    map: PersistentMap<EdgeLabel, PropertyInitializationInfo> = persistentMapOf()
-) : PathAwareControlFlowInfo<PathAwarePropertyInitializationInfo, PropertyInitializationInfo>(map) {
-    companion object {
-        val EMPTY = PathAwarePropertyInitializationInfo(persistentMapOf(NormalPath to PropertyInitializationInfo.EMPTY))
-    }
-
-    override val constructor: (PersistentMap<EdgeLabel, PropertyInitializationInfo>) -> PathAwarePropertyInitializationInfo =
-        ::PathAwarePropertyInitializationInfo
-}

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

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.fir.analysis.cfa.util
 
-import kotlinx.collections.immutable.PersistentMap
 import kotlinx.collections.immutable.persistentMapOf
 import org.jetbrains.kotlin.contracts.description.EventOccurrencesRange
 import org.jetbrains.kotlin.fir.references.toResolvedPropertySymbol
@@ -25,9 +24,13 @@ class PropertyInitializationInfoData(properties: Set<FirPropertySymbol>, graph:
 class PropertyInitializationInfoCollector(
     private val localProperties: Set<FirPropertySymbol>,
     private val declaredVariableCollector: DeclaredVariableCollector = DeclaredVariableCollector(),
-) : PathAwareControlFlowGraphVisitor<PathAwarePropertyInitializationInfo>() {
+) : PathAwareControlFlowGraphVisitor<PropertyInitializationInfo>() {
+    companion object {
+        val EMPTY: PathAwarePropertyInitializationInfo = persistentMapOf(NormalPath to PropertyInitializationInfo.EMPTY)
+    }
+
     override val emptyInfo: PathAwarePropertyInitializationInfo
-        get() = PathAwarePropertyInitializationInfo.EMPTY
+        get() = EMPTY
 
     override fun visitVariableAssignmentNode(
         node: VariableAssignmentNode,
@@ -55,11 +58,7 @@ class PropertyInitializationInfoCollector(
     }
 
     fun getData(graph: ControlFlowGraph) =
-        graph.collectDataForNode(
+        graph.collectDataForNode(TraverseDirection.Forward, this)
-            TraverseDirection.Forward,
-            PathAwarePropertyInitializationInfo.EMPTY,
-            this
-        )
 
     private fun processVariableWithAssignment(
         dataForNode: PathAwarePropertyInitializationInfo,
@@ -68,9 +67,9 @@ class PropertyInitializationInfoCollector(
     ): PathAwarePropertyInitializationInfo {
         assert(dataForNode.keys.isNotEmpty())
         return if (overwriteRange)
-            overwriteRange(dataForNode, symbol, EventOccurrencesRange.ZERO, ::PathAwarePropertyInitializationInfo)
+            overwriteRange(dataForNode, symbol, EventOccurrencesRange.ZERO)
         else
-            addRange(dataForNode, symbol, EventOccurrencesRange.EXACTLY_ONCE, ::PathAwarePropertyInitializationInfo)
+            addRange(dataForNode, symbol, EventOccurrencesRange.EXACTLY_ONCE)
     }
 
     // --------------------------------------------------
@@ -92,7 +91,7 @@ class PropertyInitializationInfoCollector(
             else -> return result
         }
         return declaredVariableSymbolsInCapturedScope.fold(data) { filteredData, variableSymbol ->
-            removeRange(filteredData, variableSymbol, ::PathAwarePropertyInitializationInfo)
+            removeRange(filteredData, variableSymbol)
         }
     }
 
@@ -110,36 +109,33 @@ class PropertyInitializationInfoCollector(
     }
 }
 
-internal fun <P : PathAwareControlFlowInfo<P, S>, S : ControlFlowInfo<S, K, EventOccurrencesRange>, K : Any> addRange(
+internal fun <S : ControlFlowInfo<S, K, EventOccurrencesRange>, K : Any> addRange(
-    pathAwareInfo: P,
+    pathAwareInfo: PathAwareControlFlowInfo<S>,
     key: K,
     range: EventOccurrencesRange,
-    constructor: (PersistentMap<EdgeLabel, S>) -> P
+): PathAwareControlFlowInfo<S> {
-): P {
     // before: { |-> { p1 |-> PI1 }, l1 |-> { p2 |-> PI2 } }
     // after (if key is p1):
     //   { |-> { p1 |-> PI1 + r }, l1 |-> { p1 |-> r, p2 |-> PI2 } }
-    return updateRange(pathAwareInfo, key, { existingKind -> existingKind + range }, constructor)
+    return updateRange(pathAwareInfo, key) { existingKind -> existingKind + range }
 }
 
-private fun <P : PathAwareControlFlowInfo<P, S>, S : ControlFlowInfo<S, K, EventOccurrencesRange>, K : Any> overwriteRange(
+private fun <S : ControlFlowInfo<S, K, EventOccurrencesRange>, K : Any> overwriteRange(
-    pathAwareInfo: P,
+    pathAwareInfo: PathAwareControlFlowInfo<S>,
     key: K,
     range: EventOccurrencesRange,
-    constructor: (PersistentMap<EdgeLabel, S>) -> P
+): PathAwareControlFlowInfo<S> {
-): P {
     // before: { |-> { p1 |-> PI1 }, l1 |-> { p2 |-> PI2 } }
     // after (if key is p1):
     //   { |-> { p1 |-> r }, l1 |-> { p1 |-> r, p2 |-> PI2 } }
-    return updateRange(pathAwareInfo, key, { range }, constructor)
+    return updateRange(pathAwareInfo, key) { range }
 }
 
-private inline fun <P : PathAwareControlFlowInfo<P, S>, S : ControlFlowInfo<S, K, EventOccurrencesRange>, K : Any> updateRange(
+private inline fun <S : ControlFlowInfo<S, K, EventOccurrencesRange>, K : Any> updateRange(
-    pathAwareInfo: P,
+    pathAwareInfo: PathAwareControlFlowInfo<S>,
     key: K,
     computeNewRange: (EventOccurrencesRange) -> EventOccurrencesRange,
-    constructor: (PersistentMap<EdgeLabel, S>) -> P
+): PathAwareControlFlowInfo<S> {
-): P {
     var resultMap = persistentMapOf<EdgeLabel, S>()
     // before: { |-> { p1 |-> PI1 }, l1 |-> { p2 |-> PI2 } }
     for ((label, dataPerLabel) in pathAwareInfo) {
@@ -149,14 +145,13 @@ private inline fun <P : PathAwareControlFlowInfo<P, S>, S : ControlFlowInfo<S, K
     }
     // after (if key is p1):
     //   { |-> { p1 |-> computeNewRange(PI1) }, l1 |-> { p1 |-> r, p2 |-> PI2 } }
-    return constructor(resultMap)
+    return resultMap
 }
 
-private fun <P : PathAwareControlFlowInfo<P, S>, S : ControlFlowInfo<S, K, EventOccurrencesRange>, K : Any> removeRange(
+private fun <S : ControlFlowInfo<S, K, EventOccurrencesRange>, K : Any> removeRange(
-    pathAwareInfo: P,
+    pathAwareInfo: PathAwareControlFlowInfo<S>,
     key: K,
-    constructor: (PersistentMap<EdgeLabel, S>) -> P
+): PathAwareControlFlowInfo<S> {
-): P {
     var resultMap = persistentMapOf<EdgeLabel, S>()
     // before: { |-> { p1 |-> PI1 }, l1 |-> { p2 |-> PI2 } }
     for ((label, dataPerLabel) in pathAwareInfo) {
@@ -164,5 +159,5 @@ private fun <P : PathAwareControlFlowInfo<P, S>, S : ControlFlowInfo<S, K, Event
     }
     // after (if key is p1):
     //   { |-> { }, l1 |-> { p2 |-> PI2 } }
-    return constructor(resultMap)
+    return resultMap
 }

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

@@ -93,7 +93,7 @@ object CanBeValChecker : AbstractFirPropertyInitializationChecker() {
 
             val currentCharacteristic = propertiesCharacteristics.getOrDefault(symbol, EventOccurrencesRange.ZERO)
             val info = data.getValue(node)
-            propertiesCharacteristics[symbol] = currentCharacteristic.or(info.infoAtNormalPath[symbol] ?: EventOccurrencesRange.ZERO)
+            propertiesCharacteristics[symbol] = info.values.fold(currentCharacteristic) { a, b -> b[symbol]?.let(a::or) ?: a }
         }
 
         override fun visitVariableDeclarationNode(node: VariableDeclarationNode) {

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.analysis.checkers.extended
 
 import kotlinx.collections.immutable.PersistentMap
+import kotlinx.collections.immutable.mutate
 import kotlinx.collections.immutable.persistentMapOf
 import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.KtNodeTypes
@@ -151,26 +152,19 @@ object UnusedChecker : AbstractFirPropertyInitializationChecker() {
             merge(other) // TODO: not sure
     }
 
-    class PathAwareVariableStatusInfo(
-        map: PersistentMap<EdgeLabel, VariableStatusInfo> = persistentMapOf()
-    ) : PathAwareControlFlowInfo<PathAwareVariableStatusInfo, VariableStatusInfo>(map) {
-        companion object {
-            val EMPTY = PathAwareVariableStatusInfo(persistentMapOf(NormalPath to VariableStatusInfo.EMPTY))
-        }
-
-        override val constructor: (PersistentMap<EdgeLabel, VariableStatusInfo>) -> PathAwareVariableStatusInfo =
-            ::PathAwareVariableStatusInfo
-    }
-
     private class ValueWritesWithoutReading(
         private val session: FirSession,
         private val localProperties: Set<FirPropertySymbol>
-    ) : PathAwareControlFlowGraphVisitor<PathAwareVariableStatusInfo>() {
+    ) : PathAwareControlFlowGraphVisitor<VariableStatusInfo>() {
+        companion object {
+            val EMPTY: PathAwareVariableStatusInfo = persistentMapOf(NormalPath to VariableStatusInfo.EMPTY)
+        }
+
         override val emptyInfo: PathAwareVariableStatusInfo
-            get() = PathAwareVariableStatusInfo.EMPTY
+            get() = EMPTY
 
         fun getData(graph: ControlFlowGraph): Map<CFGNode<*>, PathAwareVariableStatusInfo> {
-            return graph.collectDataForNode(TraverseDirection.Backward, PathAwareVariableStatusInfo.EMPTY, this)
+            return graph.collectDataForNode(TraverseDirection.Backward, this)
         }
 
         private fun PathAwareVariableStatusInfo.withAnnotationsFrom(node: CFGNode<*>): PathAwareVariableStatusInfo =
@@ -306,21 +300,13 @@ object UnusedChecker : AbstractFirPropertyInitializationChecker() {
             pathAwareInfo: PathAwareVariableStatusInfo,
             vararg symbols: FirPropertySymbol,
             updater: (VariableStatus?) -> VariableStatus?,
-        ): PathAwareVariableStatusInfo {
+        ): PathAwareVariableStatusInfo = pathAwareInfo.mutate {
-            var resultMap = persistentMapOf<EdgeLabel, VariableStatusInfo>()
-            var changed = false
             for ((label, dataPerLabel) in pathAwareInfo) {
                 for (symbol in symbols) {
-                    val v = updater.invoke(dataPerLabel[symbol])
+                    val v = updater.invoke(dataPerLabel[symbol]) ?: continue
-                    if (v != null) {
+                    it[label] = dataPerLabel.put(symbol, v)
-                        resultMap = resultMap.put(label, dataPerLabel.put(symbol, v))
-                        changed = true
-                    } else {
-                        resultMap = resultMap.put(label, dataPerLabel)
-                    }
                 }
             }
-            return if (changed) PathAwareVariableStatusInfo(resultMap) else pathAwareInfo
         }
     }
 
@@ -330,3 +316,5 @@ object UnusedChecker : AbstractFirPropertyInitializationChecker() {
             return fir.initializer?.source?.kind == KtFakeSourceElementKind.DesugaredForLoop
         }
 }
+
+private typealias PathAwareVariableStatusInfo = PathAwareControlFlowInfo<UnusedChecker.VariableStatusInfo>