FIR checker: make calls effect analyzer path-sensitive

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

@@ -38,6 +38,7 @@ 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
 
@@ -80,26 +81,43 @@ object FirCallsEffectAnalyzer : FirControlFlowChecker() {
             }
         }
 
-        val invocationData = graph.collectDataForNode(
+        val invocationData = graph.collectPathAwareDataForNode(
             TraverseDirection.Forward,
-            LambdaInvocationInfo.EMPTY,
+            PathAwareLambdaInvocationInfo.EMPTY,
             InvocationDataCollector(functionalTypeEffects.keys.filterTo(mutableSetOf()) { it !in leakedSymbols })
         )
 
         for ((symbol, effectDeclaration) in functionalTypeEffects) {
             graph.exitNode.previousCfgNodes.forEach { node ->
                 val requiredRange = effectDeclaration.kind
-                val foundRange = invocationData.getValue(node)[symbol] ?: EventOccurrencesRange.ZERO
+                val info = invocationData.getValue(node)
-
+                for (label in info.keys) {
-                if (foundRange !in requiredRange) {
+                    if (investigate(info.getValue(label), symbol, requiredRange, function, reporter)) {
-                    function.contractDescription.source?.let {
+                        // To avoid duplicate reports, stop investigating remaining paths once reported.
-                        reporter.report(FirErrors.WRONG_INVOCATION_KIND.on(it, symbol, requiredRange, foundRange))
+                        break
                     }
                 }
             }
         }
     }
 
+    private fun investigate(
+        info: LambdaInvocationInfo,
+        symbol: AbstractFirBasedSymbol<*>,
+        requiredRange: EventOccurrencesRange,
+        function: FirContractDescriptionOwner,
+        reporter: DiagnosticReporter
+    ): Boolean {
+        val foundRange = info[symbol] ?: EventOccurrencesRange.ZERO
+        if (foundRange !in requiredRange) {
+            function.contractDescription.source?.let {
+                reporter.report(FirErrors.WRONG_INVOCATION_KIND.on(it, symbol, requiredRange, foundRange))
+                return true
+            }
+        }
+        return false
+    }
+
     private class IllegalScopeContext(
         private val functionalTypeSymbols: Set<AbstractFirBasedSymbol<*>>,
         private val leakedSymbols: MutableMap<AbstractFirBasedSymbol<*>, MutableList<FirSourceElement>>,
@@ -185,7 +203,7 @@ object FirCallsEffectAnalyzer : FirControlFlowChecker() {
         }
     }
 
-    private class LambdaInvocationInfo(
+    class LambdaInvocationInfo(
         map: PersistentMap<FirBasedSymbol<*>, EventOccurrencesRange> = persistentMapOf(),
     ) : ControlFlowInfo<LambdaInvocationInfo, FirBasedSymbol<*>, EventOccurrencesRange>(map) {
 
@@ -207,19 +225,99 @@ object FirCallsEffectAnalyzer : FirControlFlowChecker() {
         }
     }
 
+    class PathAwareLambdaInvocationInfo(
+        map: PersistentMap<EdgeLabel, LambdaInvocationInfo> = persistentMapOf()
+    ) : ControlFlowInfo<PathAwareLambdaInvocationInfo, EdgeLabel, LambdaInvocationInfo>(map) {
+        companion object {
+            val EMPTY = PathAwareLambdaInvocationInfo(persistentMapOf(NormalPath to LambdaInvocationInfo.EMPTY))
+        }
+
+        override val constructor: (PersistentMap<EdgeLabel, LambdaInvocationInfo>) -> PathAwareLambdaInvocationInfo =
+            ::PathAwareLambdaInvocationInfo
+
+        val infoAtNormalPath: LambdaInvocationInfo
+            get() = map[NormalPath] ?: LambdaInvocationInfo.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))
+            }
+        }
+
+        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(
         val functionalTypeSymbols: Set<AbstractFirBasedSymbol<*>>
-    ) : ControlFlowGraphVisitor<LambdaInvocationInfo, Collection<LambdaInvocationInfo>>() {
+    ) : ControlFlowGraphVisitor<PathAwareLambdaInvocationInfo, Collection<Pair<EdgeLabel, PathAwareLambdaInvocationInfo>>>() {
 
-        override fun visitNode(node: CFGNode<*>, data: Collection<LambdaInvocationInfo>): LambdaInvocationInfo {
+        override fun visitNode(
-            if (data.isEmpty()) return LambdaInvocationInfo.EMPTY
+            node: CFGNode<*>,
-            return data.reduce(LambdaInvocationInfo::merge)
+            data: Collection<Pair<EdgeLabel, PathAwareLambdaInvocationInfo>>
+        ): PathAwareLambdaInvocationInfo {
+            if (data.isEmpty()) return PathAwareLambdaInvocationInfo.EMPTY
+            return data.map { (label, info) -> info.applyLabel(node, label) }
+                .reduce(PathAwareLambdaInvocationInfo::merge)
         }
 
         override fun visitFunctionCallNode(
             node: FunctionCallNode,
-            data: Collection<LambdaInvocationInfo>
+            data: Collection<Pair<EdgeLabel, PathAwareLambdaInvocationInfo>>
-        ): LambdaInvocationInfo {
+        ): PathAwareLambdaInvocationInfo {
             var dataForNode = visitNode(node, data)
 
             val functionSymbol = node.fir.toResolvedCallableSymbol() as? FirFunctionSymbol<*>?
@@ -249,22 +347,30 @@ object FirCallsEffectAnalyzer : FirControlFlowChecker() {
             return reference != null && referenceToSymbol(reference) in functionalTypeSymbols
         }
 
-        private inline fun LambdaInvocationInfo.checkReference(
+        private inline fun PathAwareLambdaInvocationInfo.checkReference(
             reference: FirReference?,
             rangeGetter: () -> EventOccurrencesRange
-        ): LambdaInvocationInfo {
+        ): PathAwareLambdaInvocationInfo {
             return if (collectDataForReference(reference)) addInvocationInfo(reference, rangeGetter()) else this
         }
 
-        private fun LambdaInvocationInfo.addInvocationInfo(
+        private fun PathAwareLambdaInvocationInfo.addInvocationInfo(
             reference: FirReference,
             range: EventOccurrencesRange
-        ): LambdaInvocationInfo {
+        ): PathAwareLambdaInvocationInfo {
             val symbol = referenceToSymbol(reference)
             return if (symbol != null) {
-                val existingKind = this[symbol] ?: EventOccurrencesRange.ZERO
+                var resultMap = persistentMapOf<EdgeLabel, LambdaInvocationInfo>()
-                val kind = existingKind + range
+                // before: { |-> { p1 |-> PI1 }, l1 |-> { p2 |-> PI2 }
-                this.put(symbol, kind)
+                for (label in this.keys) {
+                    val dataPerLabel = this[label]!!
+                    val existingKind = dataPerLabel[symbol] ?: EventOccurrencesRange.ZERO
+                    val kind = existingKind + range
+                    resultMap = resultMap.put(label, dataPerLabel.put(symbol, kind))
+                }
+                // after (if symbol is p1):
+                //   { |-> { p1 |-> PI1 + r }, l1 |-> { p1 |-> r, p2 |-> PI2 }
+                PathAwareLambdaInvocationInfo(resultMap)
             } else this
         }
     }