[FIR] Pass context data flow on code fragment analysis

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.resolve.dfa
 
+import kotlinx.collections.immutable.toPersistentSet
 import org.jetbrains.kotlin.contracts.description.LogicOperationKind
 import org.jetbrains.kotlin.contracts.description.canBeRevisited
 import org.jetbrains.kotlin.descriptors.Modality
@@ -238,7 +239,14 @@ abstract class FirDataFlowAnalyzer(
     // ----------------------------------- Code Fragment ------------------------------------------
 
     fun enterCodeFragment(codeFragment: FirCodeFragment) {
-        graphBuilder.enterCodeFragment(codeFragment).mergeIncomingFlow()
+        graphBuilder.enterCodeFragment(codeFragment).mergeIncomingFlow { flow ->
+            val realVariablesFromContext = codeFragment.codeFragmentContext?.variables.orEmpty()
+            for ((symbol, exactTypes) in realVariablesFromContext) {
+                val realVariable = variableStorage.getOrCreateIfReal(flow, symbol.fir) as? RealVariable ?: continue
+                val typeStatement = PersistentTypeStatement(realVariable, exactTypes.toPersistentSet())
+                flow.addTypeStatement(typeStatement)
+            }
+        }
     }
 
     fun exitCodeFragment(): ControlFlowGraph {

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

@@ -531,7 +531,8 @@ class BodyResolveContext(
     }
 
     fun <T> withScopesForCodeFragment(codeFragment: FirCodeFragment, holder: SessionHolder, f: () -> T): T {
-        val towerDataContext = codeFragment.towerDataContext ?: error("Context is not set for a code fragment")
+        val codeFragmentContext = codeFragment.codeFragmentContext ?: error("Context is not set for a code fragment")
+        val towerDataContext = codeFragmentContext.towerDataContext
 
         val fragmentImportTowerDataElements = computeImportingScopes(file, holder.session, holder.scopeSession)
             .map { it.asTowerDataElement(isLocal = false) }

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/DeclarationUtils.kt

@@ -10,12 +10,12 @@ import org.jetbrains.kotlin.fir.containingClassForLocalAttr
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.isLocal
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirConstructorSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.LookupTagInternals
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.name.ClassId
-import org.jetbrains.kotlin.util.OperatorNameConventions
 
 fun FirClassLikeDeclaration.getContainingDeclaration(session: FirSession): FirClassLikeDeclaration? {
     if (isLocal) {
@@ -88,6 +88,11 @@ var FirConstructor.originalConstructorIfTypeAlias: FirConstructor? by FirDeclara
 val FirConstructorSymbol.isTypeAliasedConstructor: Boolean
     get() = fir.originalConstructorIfTypeAlias != null
 
-private object CodeFragmentTowerDataContext : FirDeclarationDataKey()
+interface FirCodeFragmentContext {
+    val towerDataContext: FirTowerDataContext
+    val variables: Map<FirBasedSymbol<*>, Set<ConeKotlinType>>
+}
 
-var FirCodeFragment.towerDataContext: FirTowerDataContext? by FirDeclarationDataRegistry.data(CodeFragmentTowerDataContext)
+private object CodeFragmentContextDataKey : FirDeclarationDataKey()
+
+var FirCodeFragment.codeFragmentContext: FirCodeFragmentContext? by FirDeclarationDataRegistry.data(CodeFragmentContextDataKey)