[FIR] Cleanup FIR modules. Part 2 (dfa package)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirCallResolver.kt

@@ -16,7 +16,6 @@ import org.jetbrains.kotlin.fir.references.FirReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.references.FirSuperReference
 import org.jetbrains.kotlin.fir.references.builder.buildBackingFieldReference
-import org.jetbrains.kotlin.fir.references.builder.buildErrorNamedReference
 import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
 import org.jetbrains.kotlin.fir.references.impl.FirSimpleNamedReference
 import org.jetbrains.kotlin.fir.resolve.*
@@ -64,7 +63,7 @@ class FirCallResolver(
         @Suppress("NAME_SHADOWING")
         val functionCall = functionCall.transformExplicitReceiver(transformer, ResolutionMode.ContextIndependent)
             .also {
-                dataFlowAnalyzer.enterQualifiedAccessExpression(functionCall)
+                dataFlowAnalyzer.enterQualifiedAccessExpression()
                 functionCall.argumentList.transformArguments(transformer, ResolutionMode.ContextDependent)
             }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.fir.resolve.dfa
 
-import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.contracts.FirResolvedContractDescription
 import org.jetbrains.kotlin.fir.contracts.description.ConeBooleanConstantReference
@@ -17,7 +16,6 @@ import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.PersistentImplicitReceiverStack
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
-import org.jetbrains.kotlin.fir.resolve.defaultType
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.*
 import org.jetbrains.kotlin.fir.resolve.dfa.contracts.buildContractFir
 import org.jetbrains.kotlin.fir.resolve.dfa.contracts.createArgumentsMapping
@@ -30,6 +28,7 @@ import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.utils.addIfNotNull
+import org.jetbrains.kotlin.utils.addToStdlib.runIf
 
 class DataFlowAnalyzerContext<FLOW : Flow>(
     val graphBuilder: ControlFlowGraphBuilder,
@@ -599,9 +598,9 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         graphBuilder.exitCatchClause(catch).mergeIncomingFlow()
     }
 
-    fun enterFinallyBlock(tryExpression: FirTryExpression) {
+    fun enterFinallyBlock() {
         // TODO
-        graphBuilder.enterFinallyBlock(tryExpression).mergeIncomingFlow()
+        graphBuilder.enterFinallyBlock().mergeIncomingFlow()
     }
 
     fun exitFinallyBlock(tryExpression: FirTryExpression) {
@@ -609,9 +608,9 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         graphBuilder.exitFinallyBlock(tryExpression).mergeIncomingFlow()
     }
 
-    fun exitTryExpression(tryExpression: FirTryExpression, callCompleted: Boolean) {
+    fun exitTryExpression(callCompleted: Boolean) {
         // TODO
-        val (tryExpressionExitNode, unionNode) = graphBuilder.exitTryExpression(tryExpression, callCompleted)
+        val (tryExpressionExitNode, unionNode) = graphBuilder.exitTryExpression(callCompleted)
         tryExpressionExitNode.mergeIncomingFlow()
         unionNode?.let { unionFlowFromArguments(it) }
     }
@@ -619,7 +618,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     // ----------------------------------- Resolvable call -----------------------------------
 
     // Intentionally left empty for potential future needs (call sites are preserved)
-    fun enterQualifiedAccessExpression(qualifiedAccessExpression: FirQualifiedAccessExpression) {}
+    fun enterQualifiedAccessExpression() {}
 
     fun exitQualifiedAccessExpression(qualifiedAccessExpression: FirQualifiedAccessExpression) {
         graphBuilder.exitQualifiedAccessExpression(qualifiedAccessExpression).mergeIncomingFlow()
@@ -663,7 +662,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     }
 
     fun exitSafeCall(safeCall: FirSafeCallExpression) {
-        val node = graphBuilder.exitSafeCall(safeCall).mergeIncomingFlow()
+        val node = graphBuilder.exitSafeCall().mergeIncomingFlow()
         val previousFlow = node.previousFlow
 
         val variable = variableStorage.getOrCreateVariable(previousFlow, safeCall)
@@ -682,8 +681,8 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         graphBuilder.exitResolvedQualifierNode(resolvedQualifier).mergeIncomingFlow()
     }
 
-    fun enterCall(functionCall: FirCall) {
+    fun enterCall() {
-        graphBuilder.enterCall(functionCall)
+        graphBuilder.enterCall()
     }
 
     fun exitFunctionCall(functionCall: FirFunctionCall, callCompleted: Boolean) {
@@ -838,7 +837,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
 
         if (isAssignment) {
             if (initializer is FirConstExpression<*> && initializer.kind == FirConstKind.Null) return
-            flow.addTypeStatement(propertyVariable typeEq initializer.typeRef.coneTypeUnsafe<ConeKotlinType>())
+            flow.addTypeStatement(propertyVariable typeEq initializer.typeRef.coneTypeUnsafe())
         }
     }
 
@@ -894,6 +893,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         node: AbstractBinaryExitNode<*>,
         isAnd: Boolean
     ) {
+        @Suppress("UnnecessaryVariable")
         val bothEvaluated = isAnd
         val onlyLeftEvaluated = !bothEvaluated
 
@@ -1064,13 +1064,3 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     private val CFGNode<*>.previousFlow: FLOW
         get() = firstPreviousNode.flow
 }
-
-@DfaInternals
-fun FirElement.extractReturnType(): ConeKotlinType = when (this) {
-    is FirVariable<*> -> returnTypeRef.coneTypeUnsafe()
-    is FirSimpleFunction -> receiverTypeRef?.coneTypeUnsafe()
-    is FirAnonymousFunction -> receiverTypeRef?.coneTypeUnsafe()
-    is FirRegularClass -> defaultType()
-    is FirAnonymousObject -> typeRef.coneTypeUnsafe()
-    else -> null
-} ?: throw IllegalArgumentException("Unsupported fir: $this")

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/LogicSystem.kt

@@ -137,7 +137,7 @@ abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceCont
         val allTypes = types.flatMapTo(mutableSetOf()) { it }
         val commonTypes = allTypes.toMutableSet()
         types.forEach { commonTypes.retainAll(it) }
-        val differentTypes = types.mapNotNull { (it - commonTypes).takeIf { it.isNotEmpty() } }
+        val differentTypes = types.mapNotNull { typeSet -> (typeSet - commonTypes).takeIf { it.isNotEmpty() } }
         if (differentTypes.size == types.size) {
             context.commonSuperTypeOrNull(differentTypes.flatten())?.let { commonTypes += it }
         }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/PersistentLogicSystem.kt

@@ -248,6 +248,7 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         }
 
     private abstract class DiffIterable<T>(private val parentFlow: PersistentFlow, private var currentFlow: PersistentFlow) : Iterable<T> {
+        @Suppress("LeakingThis")
         private var currentIterator = extractIterator(currentFlow)
 
         abstract fun extractIterator(flow: PersistentFlow): Iterator<T>

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/CFGNode.kt

@@ -12,9 +12,9 @@ import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.resolve.dfa.controlFlowGraph
-import org.jetbrains.kotlin.fir.resolve.dfa.runIf
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
+import org.jetbrains.kotlin.utils.addToStdlib.runIf
 
 sealed class CFGNode<out E : FirElement>(val owner: ControlFlowGraph, val level: Int, private val id: Int) {
     companion object {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt

@@ -18,6 +18,7 @@ import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.types.isNothing
 import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitorVoid
+import org.jetbrains.kotlin.utils.addToStdlib.runIf
 import kotlin.random.Random
 
 @RequiresOptIn
@@ -57,7 +58,7 @@ class ControlFlowGraphBuilder {
     private val shouldPassFlowFromInplaceLambda: Stack<Boolean> = stackOf(true)
 
     private enum class Mode {
-        Function, DefaultArguments, TopLevel, Body, ClassInitializer, PropertyInitializer
+        Function, TopLevel, Body, ClassInitializer, PropertyInitializer
     }
 
     // ----------------------------------- Node caches -----------------------------------
@@ -773,7 +774,7 @@ class ControlFlowGraphBuilder {
         }
     }
 
-    fun enterFinallyBlock(tryExpression: FirTryExpression): FinallyBlockEnterNode {
+    fun enterFinallyBlock(): FinallyBlockEnterNode {
         val enterNode = finallyEnterNodes.pop()
         lastNodes.push(enterNode)
         return enterNode
@@ -787,7 +788,6 @@ class ControlFlowGraphBuilder {
     }
 
     fun exitTryExpression(
-        tryExpression: FirTryExpression,
         callCompleted: Boolean
     ): Pair<TryExpressionExitNode, UnionFunctionCallArgumentsNode?> {
         levelCounter--
@@ -816,7 +816,7 @@ class ControlFlowGraphBuilder {
         return createResolvedQualifierNode(resolvedQualifier).also(this::addNewSimpleNode)
     }
 
-    fun enterCall(call: FirCall) {
+    fun enterCall() {
         levelCounter++
     }
 
@@ -979,7 +979,7 @@ class ControlFlowGraphBuilder {
         return enterNode
     }
 
-    fun exitSafeCall(safeCall: FirSafeCallExpression): ExitSafeCallNode {
+    fun exitSafeCall(): ExitSafeCallNode {
         return exitSafeCallNodes.pop().also {
             addNewSimpleNode(it)
             it.updateDeadStatus()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/model.kt

@@ -14,6 +14,8 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
 import org.jetbrains.kotlin.fir.types.ConeClassErrorType
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import kotlin.contracts.ExperimentalContracts
+import kotlin.contracts.contract
 
 // --------------------------------------- Variables ---------------------------------------
 
@@ -159,6 +161,8 @@ abstract class TypeStatement : Statement<TypeStatement>() {
     }
 }
 
+operator fun TypeStatement.plus(other: TypeStatement?): TypeStatement = other?.let { this + other } ?: this
+
 class MutableTypeStatement(
     override val variable: RealVariable,
     override val exactType: MutableSet<ConeKotlinType> = linkedSetOf(),
@@ -205,6 +209,16 @@ fun Implication.invertCondition(): Implication = Implication(condition.invert(),
 typealias TypeStatements = Map<RealVariable, TypeStatement>
 typealias MutableTypeStatements = MutableMap<RealVariable, MutableTypeStatement>
 
+fun MutableTypeStatements.addStatement(variable: RealVariable, statement: TypeStatement) {
+    put(variable, statement.asMutableStatement()) { it.apply { this += statement } }
+}
+
+fun MutableTypeStatements.mergeTypeStatements(other: TypeStatements) {
+    other.forEach { (variable, info) ->
+        addStatement(variable, info)
+    }
+}
+
 // --------------------------------------- DSL ---------------------------------------
 
 infix fun DataFlowVariable.eq(constant: Boolean?): OperationStatement {
@@ -240,3 +254,23 @@ infix fun RealVariable.typeNotEq(type: ConeKotlinType): TypeStatement =
     } else {
         MutableTypeStatement(this)
     }
+
+// --------------------------------------- Utils ---------------------------------------
+
+@OptIn(ExperimentalContracts::class)
+fun DataFlowVariable.isSynthetic(): Boolean {
+    contract {
+        returns(true) implies (this@isSynthetic is SyntheticVariable)
+        returns(false) implies (this@isSynthetic is RealVariable)
+    }
+    return this is SyntheticVariable
+}
+
+@OptIn(ExperimentalContracts::class)
+fun DataFlowVariable.isReal(): Boolean {
+    contract {
+        returns(true) implies (this@isReal is RealVariable)
+        returns(false) implies (this@isReal is SyntheticVariable)
+    }
+    return this is RealVariable
+}

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/util.kt

@@ -25,36 +25,6 @@ import kotlin.contracts.ExperimentalContracts
 import kotlin.contracts.InvocationKind
 import kotlin.contracts.contract
 
-@OptIn(ExperimentalContracts::class)
-fun DataFlowVariable.isSynthetic(): Boolean {
-    contract {
-        returns(true) implies (this@isSynthetic is SyntheticVariable)
-        returns(false) implies (this@isSynthetic is RealVariable)
-    }
-    return this is SyntheticVariable
-}
-
-@OptIn(ExperimentalContracts::class)
-fun DataFlowVariable.isReal(): Boolean {
-    contract {
-        returns(true) implies (this@isReal is RealVariable)
-        returns(false) implies (this@isReal is SyntheticVariable)
-    }
-    return this is RealVariable
-}
-
-operator fun TypeStatement.plus(other: TypeStatement?): TypeStatement = other?.let { this + other } ?: this
-
-fun MutableTypeStatements.addStatement(variable: RealVariable, statement: TypeStatement) {
-    put(variable, statement.asMutableStatement()) { it.apply { this += statement } }
-}
-
-fun MutableTypeStatements.mergeTypeStatements(other: TypeStatements) {
-    other.forEach { (variable, info) ->
-        addStatement(variable, info)
-    }
-}
-
 @OptIn(ExperimentalContracts::class)
 internal inline fun <K, V> MutableMap<K, V>.put(key: K, value: V, remappingFunction: (existing: V) -> V) {
     contract {
@@ -69,7 +39,11 @@ internal inline fun <K, V> MutableMap<K, V>.put(key: K, value: V, remappingFunct
 }
 
 @OptIn(ExperimentalContracts::class)
-internal inline fun <K, V> PersistentMap<K, V>.put(key: K, valueProducer: () -> V, remappingFunction: (existing: V) -> V): PersistentMap<K, V> {
+internal inline fun <K, V> PersistentMap<K, V>.put(
+    key: K,
+    valueProducer: () -> V,
+    remappingFunction: (existing: V) -> V
+): PersistentMap<K, V> {
     contract {
         callsInPlace(remappingFunction, InvocationKind.AT_MOST_ONCE)
         callsInPlace(valueProducer, InvocationKind.AT_MOST_ONCE)
@@ -135,10 +109,3 @@ internal val FirResolvable.symbol: AbstractFirBasedSymbol<*>?
         is FirNamedReferenceWithCandidate -> reference.candidateSymbol
         else -> null
     }
-
-//val ConeKotlinType.isNothingOrNullableNothing: Boolean = when (this) {
-//    is ConeFlexibleType -> lowerBound.isNothingOrNullableNothing
-//    else -> false
-//}
-
-inline fun <R> runIf(condition: Boolean, block: () -> R): R? = if (condition) block() else null

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirControlFlowStatementsResolveTransformer.kt

@@ -162,13 +162,13 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
         }
 
         result = if (result.finallyBlock != null) {
-            result.also(dataFlowAnalyzer::enterFinallyBlock)
+            result.also { dataFlowAnalyzer.enterFinallyBlock() }
                 .transformFinallyBlock(transformer, ResolutionMode.ContextIndependent)
                 .also(dataFlowAnalyzer::exitFinallyBlock)
         } else {
             result
         }
-        dataFlowAnalyzer.exitTryExpression(result, callCompleted)
+        dataFlowAnalyzer.exitTryExpression(callCompleted)
         return result.compose()
     }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt

@@ -158,7 +158,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
         }
         when (result) {
             is FirQualifiedAccessExpression -> {
-                dataFlowAnalyzer.enterQualifiedAccessExpression(result)
+                dataFlowAnalyzer.enterQualifiedAccessExpression()
                 result = components.transformQualifiedAccessUsingSmartcastInfo(result)
                 dataFlowAnalyzer.exitQualifiedAccessExpression(result)
             }
@@ -206,7 +206,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
             storeTypeFromCallee(functionCall)
         }
         if (functionCall.calleeReference !is FirSimpleNamedReference) return functionCall.compose()
-        dataFlowAnalyzer.enterCall(functionCall)
+        dataFlowAnalyzer.enterCall()
         functionCall.annotations.forEach { it.accept(this, data) }
         functionCall.transform<FirFunctionCall, Nothing?>(InvocationKindTransformer, null)
         functionCall.transformTypeArguments(transformer, ResolutionMode.ContextIndependent)
@@ -700,7 +700,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
             })
         }
 
-        dataFlowAnalyzer.enterCall(delegatedConstructorCall)
+        dataFlowAnalyzer.enterCall()
         var callCompleted = true
         var result = delegatedConstructorCall
         try {