FIR: Adapt DataFlowAnalyzer to PersistentReceiverStackImpl

It will be used as main implementation in further commits

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

@@ -12,31 +12,37 @@ import kotlinx.collections.immutable.persistentMapOf
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitDispatchReceiverValue
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.name.Name
 
 class PersistentImplicitReceiverStack private constructor(
     private val stack: PersistentList<ImplicitReceiverValue<*>>,
     // This multi-map holds indexes of the stack ^
     private val indexesPerLabel: PersistentSetMultimap<Name, Int>,
-    private val indexesPerSymbol: PersistentMap<FirBasedSymbol<*>, Int>
+    private val indexesPerSymbol: PersistentMap<FirBasedSymbol<*>, Int>,
+    private val originalTypes: PersistentList<ConeKotlinType>,
 ) : ImplicitReceiverStack(), Iterable<ImplicitReceiverValue<*>> {
     val size: Int get() = stack.size
 
     constructor() : this(
         persistentListOf(),
         PersistentSetMultimap(),
-        persistentMapOf()
+        persistentMapOf(),
+        persistentListOf(),
     )
 
     fun add(name: Name?, value: ImplicitReceiverValue<*>): PersistentImplicitReceiverStack {
         val stack = stack.add(value)
+        val originalTypes = originalTypes.add(value.type)
         val index = stack.size - 1
         val indexesPerLabel = name?.let { indexesPerLabel.put(it, index) } ?: indexesPerLabel
         val indexesPerSymbol = indexesPerSymbol.put(value.boundSymbol, index)
+
         return PersistentImplicitReceiverStack(
             stack,
             indexesPerLabel,
-            indexesPerSymbol
+            indexesPerSymbol,
+            originalTypes
         )
     }
 
@@ -64,4 +70,19 @@ class PersistentImplicitReceiverStack private constructor(
     override operator fun iterator(): Iterator<ImplicitReceiverValue<*>> {
         return stack.iterator()
     }
+
+    // These methods are only used at org.jetbrains.kotlin.fir.resolve.dfa.FirDataFlowAnalyzer.Companion.createFirDataFlowAnalyzer
+    // No need to be extracted to an interface
+    fun getReceiverIndex(symbol: FirBasedSymbol<*>): Int? = indexesPerSymbol[symbol]
+
+    fun getOriginalType(index: Int): ConeKotlinType {
+        return originalTypes[index]
+    }
+
+    // This method is only used from DFA and it's in some sense breaks persistency contracts of the data structure
+    // But it's ok since DFA handles everything properly yet, but still may be it should be rewritten somehow
+    fun replaceReceiverType(index: Int, type: ConeKotlinType) {
+        assert(index >= 0 && index < stack.size)
+        stack[index].replaceType(type)
+    }
 }

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

@@ -14,9 +14,8 @@ import org.jetbrains.kotlin.fir.contracts.description.ConeConstantReference
 import org.jetbrains.kotlin.fir.contracts.description.ConeReturnsEffectDeclaration
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.references.FirResolvedCallableReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
-import org.jetbrains.kotlin.fir.resolve.ImplicitReceiverStackImpl
+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.*
@@ -26,7 +25,6 @@ import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirAbstractBod
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.CallableId
-import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.name.FqName
@@ -60,8 +58,8 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
             dataFlowAnalyzerContext: DataFlowAnalyzerContext<PersistentFlow>
         ): FirDataFlowAnalyzer<*> =
             object : FirDataFlowAnalyzer<PersistentFlow>(components, dataFlowAnalyzerContext) {
-                private val receiverStack: ImplicitReceiverStackImpl
+                private val receiverStack: PersistentImplicitReceiverStack
-                    get() = components.implicitReceiverStack as ImplicitReceiverStackImpl
+                    get() = components.implicitReceiverStack as PersistentImplicitReceiverStack
 
                 override val logicSystem: PersistentLogicSystem = object : PersistentLogicSystem(components.inferenceComponents.ctx) {
                     override fun processUpdatedReceiverVariable(flow: PersistentFlow, variable: RealVariable) {