[FIR] Refactor DataFlowVariables

Make inheritors of DataFlowVariable public, replace flags in DFV with types

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

@@ -7,76 +7,90 @@ package org.jetbrains.kotlin.fir.resolve.dfa
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import kotlin.contracts.ExperimentalContracts
+import kotlin.contracts.contract
 
 /*
  * isSynthetic = false for variables that represents actual variables in fir
  * isSynthetic = true for complex expressions (like when expression)
  */
-sealed class DataFlowVariable(val variableIndexForDebug: Int, val fir: FirElement) {
+sealed class DataFlowVariable(private val variableIndexForDebug: Int, val fir: FirElement) {
-    abstract val isSynthetic: Boolean
-    abstract val aliasedVariable: DataFlowVariable
-    abstract val isThisReference: Boolean
-
     final override fun toString(): String {
         return "d$variableIndexForDebug"
     }
 }
 
-private class RealDataFlowVariable(index: Int, fir: FirElement, override val isThisReference: Boolean) : DataFlowVariable(index, fir) {
+open class RealDataFlowVariable(index: Int, fir: FirElement, val isThisReference: Boolean) : DataFlowVariable(index, fir)
-    override val isSynthetic: Boolean get() = false
 
-    override val aliasedVariable: DataFlowVariable get() = this
+class SyntheticDataFlowVariable(index: Int, fir: FirElement) : DataFlowVariable(index, fir)
+
+class AliasedDataFlowVariable(
+    index: Int,
+    fir: FirElement,
+    var delegate: RealDataFlowVariable
+) : RealDataFlowVariable(index, fir, delegate.isThisReference)
+
+// -------------------------------------------------------------------------------------------------------------------------
+
+@UseExperimental(ExperimentalContracts::class)
+fun DataFlowVariable.isSynthetic(): Boolean {
+    contract {
+        returns(true) implies (this@isSynthetic is SyntheticDataFlowVariable)
+    }
+    return this is SyntheticDataFlowVariable
 }
 
-private class SyntheticDataFlowVariable(index: Int, fir: FirElement) : DataFlowVariable(index, fir) {
+@UseExperimental(ExperimentalContracts::class)
-    override val isSynthetic: Boolean get() = true
+fun DataFlowVariable.isAliasVariable(): Boolean {
-
+    contract {
-    override val aliasedVariable: DataFlowVariable get() = this
+        returns(true) implies (this@isAliasVariable is AliasedDataFlowVariable)
-
+    }
-    override val isThisReference: Boolean get() = false
+    return this is AliasedDataFlowVariable
 }
 
-private class AliasedDataFlowVariable(index: Int, fir: FirElement, var delegate: DataFlowVariable) : DataFlowVariable(index, fir) {
+val RealDataFlowVariable.variableUnderAlias: RealDataFlowVariable
-    override val isSynthetic: Boolean get() = delegate.isSynthetic
+    get() {
-
+        var variable = this
-    override val aliasedVariable: DataFlowVariable get() = delegate.aliasedVariable
+        while (variable.isAliasVariable()) {
-
+            variable = variable.delegate
-    override val isThisReference: Boolean get() = false
+        }
-}
+        return variable
+    }
 
+// -------------------------------------------------------------------------------------------------------------------------
 
 class DataFlowVariableStorage {
     private val fir2DfiMap: MutableMap<FirElement, DataFlowVariable> = mutableMapOf()
     private var debugIndexCounter: Int = 1
 
-    fun getOrCreateNewRealVariable(symbol: FirBasedSymbol<*>): DataFlowVariable {
+    fun getOrCreateNewRealVariable(symbol: FirBasedSymbol<*>): RealDataFlowVariable {
         return getOrCreateNewRealVariableImpl(symbol, false)
     }
 
-    fun getOrCreateNewThisRealVariable(symbol: FirBasedSymbol<*>): DataFlowVariable {
+    fun getOrCreateNewThisRealVariable(symbol: FirBasedSymbol<*>): RealDataFlowVariable {
         return getOrCreateNewRealVariableImpl(symbol, true)
     }
 
-    private fun getOrCreateNewRealVariableImpl(symbol: FirBasedSymbol<*>, isThisReference: Boolean): DataFlowVariable {
+    private fun getOrCreateNewRealVariableImpl(symbol: FirBasedSymbol<*>, isThisReference: Boolean): RealDataFlowVariable {
         val fir = symbol.fir
-        get(fir)?.let { return it }
+        get(fir)?.let { return it as RealDataFlowVariable }
         return RealDataFlowVariable(debugIndexCounter++, fir, isThisReference).also { storeVariable(it, fir) }
     }
 
-    fun getOrCreateNewSyntheticVariable(fir: FirElement): DataFlowVariable {
+    fun getOrCreateNewSyntheticVariable(fir: FirElement): SyntheticDataFlowVariable {
-        get(fir)?.let { return it }
+        get(fir)?.let { return it as SyntheticDataFlowVariable }
         return SyntheticDataFlowVariable(debugIndexCounter++, fir).also { storeVariable(it, fir) }
     }
 
-    fun createAliasVariable(symbol: FirBasedSymbol<*>, variable: DataFlowVariable) {
+    fun createAliasVariable(symbol: FirBasedSymbol<*>, variable: RealDataFlowVariable) {
         createAliasVariable(symbol.fir, variable)
     }
 
-    private fun createAliasVariable(fir: FirElement, variable: DataFlowVariable) {
+    private fun createAliasVariable(fir: FirElement, variable: RealDataFlowVariable) {
         AliasedDataFlowVariable(debugIndexCounter++, fir, variable).also { storeVariable(it, fir) }
     }
 
-    fun rebindAliasVariable(aliasVariable: DataFlowVariable, newVariable: DataFlowVariable) {
+    fun rebindAliasVariable(aliasVariable: RealDataFlowVariable, newVariable: RealDataFlowVariable) {
         val fir = removeVariable(aliasVariable)
         requireNotNull(fir)
         createAliasVariable(fir, newVariable)
@@ -105,8 +119,8 @@ class DataFlowVariableStorage {
         return fir2DfiMap[firElement]
     }
 
-    operator fun get(symbol: FirBasedSymbol<*>): DataFlowVariable? {
+    operator fun get(symbol: FirBasedSymbol<*>): RealDataFlowVariable? {
-        return fir2DfiMap[symbol.fir]
+        return fir2DfiMap[symbol.fir] as RealDataFlowVariable?
     }
 
     fun reset() {

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

@@ -23,7 +23,10 @@ import org.jetbrains.kotlin.fir.symbols.FirSymbolOwner
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
-import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.coneTypeSafe
+import org.jetbrains.kotlin.fir.types.coneTypeUnsafe
+import org.jetbrains.kotlin.fir.types.isMarkedNullable
 import org.jetbrains.kotlin.ir.expressions.IrConstKind
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -52,7 +55,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
          * DataFlowAnalyzer holds variables only for declarations that have some smartcast (or can have)
          * If there is no useful information there is no data flow variable also
          */
-        val variable = qualifiedAccessExpression.variable?.aliasedVariable ?: return null
+        val variable = qualifiedAccessExpression.realVariable?.variableUnderAlias ?: return null
         return graphBuilder.lastNode.flow.approvedFacts(variable)?.exactType ?: return null
     }
 
@@ -109,7 +112,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
 
         if (typeOperatorCall.operation !in FirOperation.TYPES) return
         val type = typeOperatorCall.conversionTypeRef.coneTypeSafe<ConeKotlinType>() ?: return
-        val operandVariable = getOrCreateRealVariable(typeOperatorCall.argument)?.aliasedVariable ?: return
+        val operandVariable = getOrCreateRealVariable(typeOperatorCall.argument)?.variableUnderAlias ?: return
 
         var flow = node.flow
         when (typeOperatorCall.operation) {
@@ -457,7 +460,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
          *   }
          */
         variableStorage[initializer]?.let { initializerVariable ->
-            assert(initializerVariable.isSynthetic)
+            assert(initializerVariable.isSynthetic())
             val realVariable = getOrCreateRealVariable(variable.symbol)
             node.flow = node.flow.copyNotApprovedFacts(initializerVariable, realVariable)
         }
@@ -471,7 +474,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
     fun exitVariableAssignment(assignment: FirVariableAssignment) {
         graphBuilder.exitVariableAssignment(assignment).mergeIncomingFlow()
         val lhsVariable = variableStorage[assignment.resolvedSymbol ?: return] ?: return
-        val rhsVariable = variableStorage[assignment.rValue.resolvedSymbol ?: return]?.takeIf { !it.isSynthetic } ?: return
+        val rhsVariable = variableStorage[assignment.rValue.resolvedSymbol ?: return]?.takeIf { !it.isSynthetic() } ?: return
         variableStorage.rebindAliasVariable(lhsVariable, rhsVariable)
     }
 
@@ -598,7 +601,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
 
     // -------------------------------------------------------------------------------------------------------------------------
 
-    private fun approveFact(variable: DataFlowVariable, condition: Condition, flow: Flow): MutableMap<DataFlowVariable, FirDataFlowInfo>? =
+    private fun approveFact(variable: DataFlowVariable, condition: Condition, flow: Flow): MutableMap<RealDataFlowVariable, FirDataFlowInfo>? =
         logicSystem.approveFact(variable, condition, flow)
 
     private fun FirBinaryLogicExpression.getVariables(): Pair<DataFlowVariable, DataFlowVariable> =
@@ -628,9 +631,9 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
     }
 
     // -------------------------------- get or create variable --------------------------------
-    private fun getOrCreateSyntheticVariable(fir: FirElement): DataFlowVariable = variableStorage.getOrCreateNewSyntheticVariable(fir)
+    private fun getOrCreateSyntheticVariable(fir: FirElement): SyntheticDataFlowVariable = variableStorage.getOrCreateNewSyntheticVariable(fir)
 
-    private fun getOrCreateRealVariable(fir: FirElement): DataFlowVariable? {
+    private fun getOrCreateRealVariable(fir: FirElement): RealDataFlowVariable? {
         if (fir is FirThisReceiverExpressionImpl) {
             return variableStorage.getOrCreateNewThisRealVariable(fir.calleeReference.boundSymbol ?: return null)
         }
@@ -638,8 +641,8 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
         return variableStorage.getOrCreateNewRealVariable(symbol)
     }
 
-    private fun getOrCreateRealVariable(symbol: FirBasedSymbol<*>): DataFlowVariable =
+    private fun getOrCreateRealVariable(symbol: FirBasedSymbol<*>): RealDataFlowVariable =
-        variableStorage.getOrCreateNewRealVariable(symbol).aliasedVariable
+        variableStorage.getOrCreateNewRealVariable(symbol).variableUnderAlias
 
     private fun getOrCreateVariable(fir: FirElement): DataFlowVariable {
         val symbol = fir.resolvedSymbol
@@ -651,7 +654,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
 
     // -------------------------------- get variable --------------------------------
 
-    private val FirElement.variable: DataFlowVariable?
+    private val FirElement.realVariable: RealDataFlowVariable?
         get() {
             val symbol: FirBasedSymbol<*> = if (this is FirThisReceiverExpressionImpl) {
                 calleeReference.boundSymbol
@@ -661,8 +664,8 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
             return variableStorage[symbol]
         }
 
-    private fun getRealVariablesForSafeCallChain(call: FirExpression): Collection<DataFlowVariable> {
+    private fun getRealVariablesForSafeCallChain(call: FirExpression): Collection<RealDataFlowVariable> {
-        val result = mutableListOf<DataFlowVariable>()
+        val result = mutableListOf<RealDataFlowVariable>()
 
         fun collect(call: FirExpression) {
             when (call) {
@@ -713,7 +716,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
 
     }
 
-    private fun addApprovedFact(flow: Flow, variable: DataFlowVariable, info: FirDataFlowInfo): Flow {
+    private fun addApprovedFact(flow: Flow, variable: RealDataFlowVariable, info: FirDataFlowInfo): Flow {
         return flow.addApprovedFact(variable, info).also {
             if (variable.isThisReference) {
                 updateReceiverType(flow, variable)
@@ -727,7 +730,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
     }
 
     private fun Flow.removeSyntheticVariable(variable: DataFlowVariable): Flow {
-        if (!variable.isSynthetic) return this
+        if (!variable.isSynthetic()) return this
         return removeVariable(variable)
     }
 }

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

@@ -10,7 +10,7 @@ import org.jetbrains.kotlin.fir.types.render
 
 data class ConditionalFirDataFlowInfo(
     val condition: Condition,
-    val variable: DataFlowVariable,
+    val variable: RealDataFlowVariable,
     val info: FirDataFlowInfo
 ) {
     fun invert(): ConditionalFirDataFlowInfo {

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

@@ -7,18 +7,27 @@ package org.jetbrains.kotlin.fir.resolve.dfa
 
 import com.google.common.collect.HashMultimap
 
-class Flow(
+class Flow private constructor(
-    val approvedInfos: MutableMap<DataFlowVariable, FirDataFlowInfo> = mutableMapOf(),
+    val approvedInfos: MutableMap<RealDataFlowVariable, FirDataFlowInfo>,
-    val conditionalInfos: HashMultimap<DataFlowVariable, ConditionalFirDataFlowInfo> = HashMultimap.create(),
+    val conditionalInfos: HashMultimap<DataFlowVariable, ConditionalFirDataFlowInfo>,
     private var state: State = State.Building
 ) {
+    constructor(
+        approvedInfos: MutableMap<RealDataFlowVariable, FirDataFlowInfo> = mutableMapOf(),
+        conditionalInfos: HashMultimap<DataFlowVariable, ConditionalFirDataFlowInfo> = HashMultimap.create()
+    ) : this(approvedInfos, conditionalInfos, State.Building)
+
+    companion object {
+        val EMPTY = Flow(mutableMapOf(), HashMultimap.create(), State.Frozen)
+    }
+
     private val isFrozen: Boolean get() = state == State.Frozen
 
     fun freeze() {
         state = State.Frozen
     }
 
-    fun addApprovedFact(variable: DataFlowVariable, info: FirDataFlowInfo): Flow {
+    fun addApprovedFact(variable: RealDataFlowVariable, info: FirDataFlowInfo): Flow {
         if (isFrozen) return copyForBuilding().addApprovedFact(variable, info)
         approvedInfos.compute(variable) { _, existingInfo ->
             if (existingInfo == null) info
@@ -38,9 +47,11 @@ class Flow(
         to: DataFlowVariable,
         transform: ((ConditionalFirDataFlowInfo) -> ConditionalFirDataFlowInfo)? = null
     ): Flow {
-        if (isFrozen)
+        if (isFrozen) {
             return copyForBuilding().copyNotApprovedFacts(from, to, transform)
-        var facts = if (from.isSynthetic) {
+        }
+
+        var facts = if (from.isSynthetic()) {
             conditionalInfos.removeAll(from)
         } else {
             conditionalInfos[from]
@@ -52,7 +63,7 @@ class Flow(
         return this
     }
 
-    fun approvedFacts(variable: DataFlowVariable): FirDataFlowInfo? {
+    fun approvedFacts(variable: RealDataFlowVariable): FirDataFlowInfo? {
         return approvedInfos[variable]
     }
 
@@ -63,11 +74,7 @@ class Flow(
         return this
     }
 
-    companion object {
+    private enum class State {
-        val EMPTY = Flow(mutableMapOf(), HashMultimap.create(), State.Frozen)
-    }
-
-    enum class State {
         Building, Frozen
     }
 

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

@@ -14,7 +14,7 @@ class LogicSystem(private val context: DataFlowInferenceContext) {
         storages.singleOrNull()?.let {
             return it
         }
-        val approvedFacts = mutableMapOf<DataFlowVariable, FirDataFlowInfo>()
+        val approvedFacts = mutableMapOf<RealDataFlowVariable, FirDataFlowInfo>()
         storages.map { it.approvedInfos.keys }
             .intersectSets()
             .forEach { variable ->
@@ -39,13 +39,13 @@ class LogicSystem(private val context: DataFlowInferenceContext) {
     }
 
     fun andForVerifiedFacts(
-        left: Map<DataFlowVariable, FirDataFlowInfo>?,
+        left: Map<RealDataFlowVariable, FirDataFlowInfo>?,
-        right: Map<DataFlowVariable, FirDataFlowInfo>?
+        right: Map<RealDataFlowVariable, FirDataFlowInfo>?
-    ): Map<DataFlowVariable, FirDataFlowInfo>? {
+    ): Map<RealDataFlowVariable, FirDataFlowInfo>? {
         if (left.isNullOrEmpty()) return right
         if (right.isNullOrEmpty()) return left
 
-        val map = mutableMapOf<DataFlowVariable, FirDataFlowInfo>()
+        val map = mutableMapOf<RealDataFlowVariable, FirDataFlowInfo>()
         for (variable in left.keys.union(right.keys)) {
             val leftInfo = left[variable]
             val rightInfo = right[variable]
@@ -55,11 +55,11 @@ class LogicSystem(private val context: DataFlowInferenceContext) {
     }
 
     fun orForVerifiedFacts(
-        left: Map<DataFlowVariable, FirDataFlowInfo>?,
+        left: Map<RealDataFlowVariable, FirDataFlowInfo>?,
-        right: Map<DataFlowVariable, FirDataFlowInfo>?
+        right: Map<RealDataFlowVariable, FirDataFlowInfo>?
-    ): Map<DataFlowVariable, FirDataFlowInfo>? {
+    ): Map<RealDataFlowVariable, FirDataFlowInfo>? {
         if (left.isNullOrEmpty() || right.isNullOrEmpty()) return null
-        val map = mutableMapOf<DataFlowVariable, FirDataFlowInfo>()
+        val map = mutableMapOf<RealDataFlowVariable, FirDataFlowInfo>()
         for (variable in left.keys.intersect(right.keys)) {
             val leftInfo = left[variable]!!
             val rightInfo = right[variable]!!
@@ -68,20 +68,20 @@ class LogicSystem(private val context: DataFlowInferenceContext) {
         return map
     }
 
-    fun approveFactsInsideFlow(variable: DataFlowVariable, condition: Condition, flow: Flow): Pair<Flow, Collection<DataFlowVariable>> {
+    fun approveFactsInsideFlow(variable: DataFlowVariable, condition: Condition, flow: Flow): Pair<Flow, Collection<RealDataFlowVariable>> {
         val notApprovedFacts: Set<ConditionalFirDataFlowInfo> = flow.conditionalInfos[variable]
         if (notApprovedFacts.isEmpty()) {
             return flow to emptyList()
         }
         @Suppress("NAME_SHADOWING")
         val flow = flow.copyForBuilding()
-        val newFacts = HashMultimap.create<DataFlowVariable, FirDataFlowInfo>()
+        val newFacts = HashMultimap.create<RealDataFlowVariable, FirDataFlowInfo>()
         notApprovedFacts.forEach {
             if (it.condition == condition) {
                 newFacts.put(it.variable, it.info)
             }
         }
-        val updatedReceivers = mutableSetOf<DataFlowVariable>()
+        val updatedReceivers = mutableSetOf<RealDataFlowVariable>()
 
         newFacts.asMap().forEach { (variable, infos) ->
             @Suppress("NAME_SHADOWING")
@@ -97,12 +97,12 @@ class LogicSystem(private val context: DataFlowInferenceContext) {
         return flow to updatedReceivers
     }
 
-    fun approveFact(variable: DataFlowVariable, condition: Condition, flow: Flow): MutableMap<DataFlowVariable, FirDataFlowInfo> {
+    fun approveFact(variable: DataFlowVariable, condition: Condition, flow: Flow): MutableMap<RealDataFlowVariable, FirDataFlowInfo> {
         val notApprovedFacts: Set<ConditionalFirDataFlowInfo> = flow.conditionalInfos[variable]
         if (notApprovedFacts.isEmpty()) {
             return mutableMapOf()
         }
-        val newFacts = HashMultimap.create<DataFlowVariable, FirDataFlowInfo>()
+        val newFacts = HashMultimap.create<RealDataFlowVariable, FirDataFlowInfo>()
         notApprovedFacts.forEach {
             if (it.condition == condition) {
                 newFacts.put(it.variable, it.info)