[LL API] Encapsulate on-air class provider logic in LLFirProvider

The old 'FirProviderInterceptor' is too invasive, ad-hoc, and only a
few places in the compiler are aware of it.

This fixes the test 'FirOnAirResolveTestGenerated#testOnAirTypesResolve'
which started failing after the designation path collection refactoring.
The newly appeared class existence check started to fail:

```
val declaration = firProvider.getFirClassifierByFqName(classId)
check(declaration != null)
```

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirModuleLazyDeclarationResolver.kt

@@ -12,8 +12,8 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.tryCollectDesignation
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionInvalidator
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirSession
-import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.LLFirFirProviderInterceptor
 import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.LLFirLazyTransformerExecutor
+import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.withSyntheticClasses
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkCanceled
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.getContainingFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.withFirEntry
@@ -163,8 +163,7 @@ internal class LLFirModuleLazyDeclarationResolver(val moduleComponents: LLFirMod
                 scopeSession = scopeSession,
                 phaseRunner = moduleComponents.globalResolveComponents.phaseRunner,
                 lockProvider = moduleComponents.globalResolveComponents.lockProvider,
-                towerDataContextCollector = null,
-                firProviderInterceptor = null,
+                towerDataContextCollector = null
             )
         }
     }
@@ -175,31 +174,30 @@ internal class LLFirModuleLazyDeclarationResolver(val moduleComponents: LLFirMod
         towerDataContextCollector: FirTowerDataContextCollector?,
     ) {
         resolveFileToImportsWithoutLock(designation.firFile)
-        var currentPhase = maxOf(designation.target.resolvePhase, FirResolvePhase.IMPORTS)
 
-        val scopeSession = ScopeSession()
+        fun runTransformation() {
+            val scopeSession = ScopeSession()
+            var currentPhase = maxOf(designation.target.resolvePhase, FirResolvePhase.IMPORTS)
 
-        val firProviderInterceptor = if (onAirCreatedDeclaration) {
-            LLFirFirProviderInterceptor.createForFirElement(
-                session = designation.firFile.moduleData.session,
-                firFile = designation.firFile,
-                element = designation.target
-            )
-        } else null
+            while (currentPhase < FirResolvePhase.BODY_RESOLVE) {
+                currentPhase = currentPhase.next
+                checkCanceled()
 
-        while (currentPhase < FirResolvePhase.BODY_RESOLVE) {
-            currentPhase = currentPhase.next
-            checkCanceled()
+                LLFirLazyTransformerExecutor.execute(
+                    phase = currentPhase,
+                    designation = designation,
+                    scopeSession = scopeSession,
+                    phaseRunner = moduleComponents.globalResolveComponents.phaseRunner,
+                    lockProvider = moduleComponents.globalResolveComponents.lockProvider,
+                    towerDataContextCollector = towerDataContextCollector
+                )
+            }
+        }
 
-            LLFirLazyTransformerExecutor.execute(
-                phase = currentPhase,
-                designation = designation,
-                scopeSession = scopeSession,
-                phaseRunner = moduleComponents.globalResolveComponents.phaseRunner,
-                lockProvider = moduleComponents.globalResolveComponents.lockProvider,
-                towerDataContextCollector = towerDataContextCollector,
-                firProviderInterceptor = firProviderInterceptor,
-            )
+        if (onAirCreatedDeclaration) {
+            withSyntheticClasses(designation, ::runTransformation)
+        } else {
+            runTransformation()
         }
     }
 }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirLibrarySessionProvider.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.providers
 
+import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.SyntheticFirClassProvider
 import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
@@ -17,8 +18,10 @@ import org.jetbrains.kotlin.name.Name
 internal class LLFirLibrarySessionProvider(
     override val symbolProvider: FirSymbolProvider
 ) : FirProvider() {
-    override fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration? =
-        symbolProvider.getClassLikeSymbolByClassId(classId)?.fir
+    override fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration? {
+        return SyntheticFirClassProvider.getInstance(symbolProvider.session).getFirClassifierByFqName(classId)
+            ?: symbolProvider.getClassLikeSymbolByClassId(classId)?.fir
+    }
 
     override fun getFirClassifierContainerFile(fqName: ClassId): FirFile = shouldNotBeCalled()
 

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirProvider.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.providers
 
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
+import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.SyntheticFirClassProvider
 import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProvider
 import org.jetbrains.kotlin.analysis.providers.KotlinPackageProvider
 import org.jetbrains.kotlin.fir.FirSession
@@ -44,8 +45,10 @@ internal class LLFirProvider(
 
     override val isPhasedFirAllowed: Boolean get() = true
 
-    override fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration? =
-        providerHelper.getFirClassifierByFqName(classId)
+    override fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration? {
+        return SyntheticFirClassProvider.getInstance(session).getFirClassifierByFqName(classId)
+            ?: providerHelper.getFirClassifierByFqName(classId)
+    }
 
     override fun getFirClassifierContainerFile(fqName: ClassId): FirFile {
         return getFirClassifierContainerFileIfAny(fqName)
@@ -53,8 +56,8 @@ internal class LLFirProvider(
     }
 
     override fun getFirClassifierContainerFileIfAny(fqName: ClassId): FirFile? {
-        val fir = getFirClassifierByFqName(fqName) ?: return null // Necessary to ensure cacheProvider contains this classifier
-        return moduleComponents.cache.getContainerFirFile(fir)
+        return SyntheticFirClassProvider.getInstance(session).getFirClassifierContainerFileIfAny(fqName)
+            ?: getFirClassifierByFqName(fqName)?.let { moduleComponents.cache.getContainerFirFile(it) }
     }
 
     override fun getFirClassifierContainerFile(symbol: FirClassLikeSymbol<*>): FirFile {
@@ -62,9 +65,10 @@ internal class LLFirProvider(
             ?: error("Couldn't find container for ${symbol.classId}")
     }
 
-    override fun getFirClassifierContainerFileIfAny(symbol: FirClassLikeSymbol<*>): FirFile? =
-        moduleComponents.cache.getContainerFirFile(symbol.fir)
-
+    override fun getFirClassifierContainerFileIfAny(symbol: FirClassLikeSymbol<*>): FirFile? {
+        return SyntheticFirClassProvider.getInstance(session).getFirClassifierContainerFileIfAny(symbol.classId)
+            ?: moduleComponents.cache.getContainerFirFile(symbol.fir)
+    }
 
     override fun getFirCallableContainerFile(symbol: FirCallableSymbol<*>): FirFile? {
         symbol.fir.originalForSubstitutionOverride?.symbol?.let { originalSymbol ->

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirDesignatedBodyResolveTransformer.kt

@@ -9,7 +9,6 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptor
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyResolveComputationSession
@@ -28,8 +27,7 @@ internal class LLFirDesignatedBodyResolveTransformer(
     private val designation: FirDesignationWithFile,
     session: FirSession,
     scopeSession: ScopeSession,
-    towerDataContextCollector: FirTowerDataContextCollector?,
-    firProviderInterceptor: FirProviderInterceptor?,
+    towerDataContextCollector: FirTowerDataContextCollector?
 ) : LLFirLazyTransformer, FirBodyResolveTransformer(
     session,
     phase = FirResolvePhase.BODY_RESOLVE,
@@ -40,8 +38,7 @@ internal class LLFirDesignatedBodyResolveTransformer(
         ImplicitBodyResolveComputationSession(),
         ::LLFirEnsureBasedTransformerForReturnTypeCalculator
     ),
-    firTowerDataContextCollector = towerDataContextCollector,
-    firProviderInterceptor = firProviderInterceptor,
+    firTowerDataContextCollector = towerDataContextCollector
 ) {
     private val ideDeclarationTransformer = LLFirDeclarationTransformer(designation)
 

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirDesignatedSupertypeResolverTransformer.kt

@@ -35,8 +35,7 @@ internal class LLFirDesignatedSupertypeResolverTransformer(
     private val designation: FirDesignationWithFile,
     private val session: FirSession,
     private val scopeSession: ScopeSession,
-    private val lockProvider: LLFirLockProvider,
-    private val firProviderInterceptor: FirProviderInterceptor?,
+    private val lockProvider: LLFirLockProvider
 ) : LLFirLazyTransformer {
 
     private val supertypeComputationSession = SupertypeComputationSession()
@@ -47,8 +46,7 @@ internal class LLFirDesignatedSupertypeResolverTransformer(
             supertypeComputationSession = supertypeComputationSession,
             scopeSession = scopeSession,
             scopeForLocalClass = null,
-            localClassesNavigationInfo = null,
-            firProviderInterceptor = firProviderInterceptor,
+            localClassesNavigationInfo = null
         ) {
         val declarationTransformer = LLFirDeclarationTransformer(classDesignation)
 

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirFirProviderInterceptor.kt

@@ -1,49 +0,0 @@
-/*
- * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- */
-
-package org.jetbrains.kotlin.analysis.low.level.api.fir.transformers
-
-import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.resolve.providers.firProvider
-import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptor
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
-import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
-import org.jetbrains.kotlin.name.ClassId
-
-internal class LLFirFirProviderInterceptor private constructor(
-    private val firFile: FirFile,
-    private val session: FirSession,
-    private val symbolSet: Set<FirClassLikeSymbol<*>>,
-    private val classIdToElementMap: Map<ClassId, FirClassLikeDeclaration>
-) : FirProviderInterceptor {
-
-    override fun getFirClassifierContainerFileIfAny(symbol: FirClassLikeSymbol<*>): FirFile? =
-        if (symbolSet.contains(symbol)) firFile else session.firProvider.getFirClassifierContainerFileIfAny(symbol)
-
-    override fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration? =
-        classIdToElementMap[classId] ?: session.firProvider.getFirClassifierByFqName(classId)
-
-    companion object {
-        fun createForFirElement(session: FirSession, firFile: FirFile, element: FirElement): FirProviderInterceptor {
-            val nodeInfoCollector = object : FirVisitorVoid() {
-                val symbolSet = mutableSetOf<FirClassLikeSymbol<*>>()
-                val classIdToElementMap = mutableMapOf<ClassId, FirClassLikeDeclaration>()
-                override fun visitElement(element: FirElement) {
-                    if (element is FirClassLikeDeclaration) {
-                        symbolSet.add(element.symbol)
-                        classIdToElementMap[element.symbol.classId] = element
-                    }
-                    element.acceptChildren(this)
-                }
-            }
-            nodeInfoCollector.visitElement(element)
-
-            return LLFirFirProviderInterceptor(firFile, session, nodeInfoCollector.symbolSet, nodeInfoCollector.classIdToElementMap)
-        }
-    }
-}

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirLazyTransformerExecutor.kt

@@ -10,7 +10,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignationWithFil
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirLockProvider
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptor
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
 
 internal class LLFirLazyTransformerExecutor {
@@ -21,8 +20,7 @@ internal class LLFirLazyTransformerExecutor {
             scopeSession: ScopeSession,
             phaseRunner: LLFirPhaseRunner,
             lockProvider: LLFirLockProvider,
-            towerDataContextCollector: FirTowerDataContextCollector?,
-            firProviderInterceptor: FirProviderInterceptor?,
+            towerDataContextCollector: FirTowerDataContextCollector?
         ) {
 
             val lazyTransformer = LazyTransformerFactory.createLazyTransformer(
@@ -30,8 +28,7 @@ internal class LLFirLazyTransformerExecutor {
                 designation,
                 scopeSession,
                 lockProvider,
-                towerDataContextCollector,
-                firProviderInterceptor,
+                towerDataContextCollector
             )
             lazyTransformer.transformDeclaration(phaseRunner)
         }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LazyTransformerFactory.kt

@@ -9,7 +9,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignationWithFil
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirLockProvider
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptor
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
 
 internal object LazyTransformerFactory {
@@ -18,8 +17,7 @@ internal object LazyTransformerFactory {
         designation: FirDesignationWithFile,
         scopeSession: ScopeSession,
         lockProvider: LLFirLockProvider,
-        towerDataContextCollector: FirTowerDataContextCollector?,
-        firProviderInterceptor: FirProviderInterceptor?,
+        towerDataContextCollector: FirTowerDataContextCollector?
     ): LLFirLazyTransformer = when (phase) {
         FirResolvePhase.COMPANION_GENERATION -> LLFirDesignatedGeneratedCompanionObjectResolveTransformer(
             designation = designation,
@@ -30,8 +28,7 @@ internal object LazyTransformerFactory {
             designation = designation,
             session = designation.firFile.moduleData.session,
             scopeSession = scopeSession,
-            lockProvider = lockProvider,
-            firProviderInterceptor = firProviderInterceptor,
+            lockProvider = lockProvider
         )
         FirResolvePhase.TYPES -> LLFirDesignatedTypeResolverTransformer(
             designation,
@@ -73,8 +70,7 @@ internal object LazyTransformerFactory {
             designation,
             designation.firFile.moduleData.session,
             scopeSession,
-            towerDataContextCollector,
-            firProviderInterceptor,
+            towerDataContextCollector
         )
         FirResolvePhase.EXPECT_ACTUAL_MATCHING -> LLFirDesignatedExpectActualMatcherTransformer(
             designation,

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/SyntheticFirClassProvider.kt

@@ -0,0 +1,105 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.analysis.low.level.api.fir.transformers
+
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignationWithFile
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
+import org.jetbrains.kotlin.name.ClassId
+import java.util.WeakHashMap
+
+internal interface SyntheticFirClassProvider {
+    fun getFirClassifierContainerFileIfAny(classId: ClassId): FirFile?
+    fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration?
+
+    companion object {
+        fun getInstance(session: FirSession): SyntheticFirClassProvider {
+            return providersForThread.get()[session] ?: EmptySyntheticFirClassProvider
+        }
+    }
+}
+
+internal fun withSyntheticClasses(designation: FirDesignationWithFile, block: () -> Unit) {
+    val firSession = designation.firFile.moduleData.session
+
+    val compoundProvider = providersForThread.get()
+        .getOrPut(firSession) { CompoundSyntheticFirClassProvider() }
+
+    val provider = StaticSyntheticFirClassProvider.create(designation)
+
+    try {
+        compoundProvider.push(provider)
+        block()
+    } finally {
+        compoundProvider.pop(provider)
+    }
+}
+
+private class StaticSyntheticFirClassProvider private constructor(
+    private val firFile: FirFile,
+    private val classes: Map<ClassId, FirClassLikeDeclaration>
+) : SyntheticFirClassProvider {
+    override fun getFirClassifierContainerFileIfAny(classId: ClassId): FirFile? {
+        return if (classId in classes) firFile else null
+    }
+
+    override fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration? {
+        return classes[classId]
+    }
+
+    companion object {
+        fun create(designation: FirDesignationWithFile): SyntheticFirClassProvider {
+            val firFile = designation.firFile
+            val firElement = designation.target
+
+            val nodeInfoCollector = object : FirVisitorVoid() {
+                val classes = mutableMapOf<ClassId, FirClassLikeDeclaration>()
+                override fun visitElement(element: FirElement) {
+                    if (element is FirClassLikeDeclaration) {
+                        classes[element.symbol.classId] = element
+                    }
+                    element.acceptChildren(this)
+                }
+            }
+
+            nodeInfoCollector.visitElement(firElement)
+            return StaticSyntheticFirClassProvider(firFile, nodeInfoCollector.classes)
+        }
+    }
+}
+
+private class CompoundSyntheticFirClassProvider : SyntheticFirClassProvider {
+    private val providers = ArrayDeque<SyntheticFirClassProvider>()
+
+    fun push(provider: SyntheticFirClassProvider) {
+        providers.addFirst(provider)
+    }
+
+    fun pop(provider: SyntheticFirClassProvider) {
+        assert(providers.removeFirst() === provider)
+    }
+
+    override fun getFirClassifierContainerFileIfAny(classId: ClassId): FirFile? {
+        return providers.asReversed()
+            .firstNotNullOfOrNull { it.getFirClassifierContainerFileIfAny(classId) }
+    }
+
+    override fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration? {
+        return providers.asReversed()
+            .firstNotNullOfOrNull { it.getFirClassifierByFqName(classId) }
+    }
+}
+
+private object EmptySyntheticFirClassProvider : SyntheticFirClassProvider {
+    override fun getFirClassifierContainerFileIfAny(classId: ClassId) = null
+    override fun getFirClassifierByFqName(classId: ClassId) = null
+}
+
+private val providersForThread: ThreadLocal<WeakHashMap<FirSession, CompoundSyntheticFirClassProvider>> =
+    ThreadLocal.withInitial { WeakHashMap() }