[Analysis API] add support for generated declarations to KtFirPackageScope

^KT-58262

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/KtFirAnalysisSession.kt

@@ -18,9 +18,15 @@ import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LowLevelFirApiFacadeForResolveOnAir
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.CompositeKotlinDeclarationProvider
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.CompositeKotlinPackageProvider
+import org.jetbrains.kotlin.analysis.low.level.api.fir.resolve.extensions.llResolveExtensionTool
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.analysis.project.structure.allDirectDependencies
 import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProvider
+import org.jetbrains.kotlin.analysis.providers.KotlinPackageProvider
 import org.jetbrains.kotlin.analysis.providers.createDeclarationProvider
+import org.jetbrains.kotlin.analysis.providers.createPackageProvider
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.moduleData
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
@@ -29,6 +35,7 @@ import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.platform.TargetPlatform
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.utils.addIfNotNull
 
 @OptIn(KtAnalysisApiInternals::class)
 @Suppress("AnalysisApiMissingLifetimeCheck")
@@ -142,7 +149,32 @@ private constructor(
     internal val targetPlatform: TargetPlatform get() = useSiteSession.moduleData.platform
 
     val useSiteAnalysisScope: GlobalSearchScope = analysisScopeProviderImpl.getAnalysisScope()
-    val useSiteScopeDeclarationProvider: KotlinDeclarationProvider = project.createDeclarationProvider(useSiteAnalysisScope)
+
+    val useSiteScopeDeclarationProvider: KotlinDeclarationProvider
+    val useSitePackageProvider: KotlinPackageProvider
+
+    init {
+        val extensionTools = buildList {
+            addIfNotNull(useSiteSession.llResolveExtensionTool)
+            useSiteModule.allDirectDependencies().mapNotNullTo(this) { dependency ->
+                firResolveSession.getSessionFor(dependency).llResolveExtensionTool
+            }
+        }
+        useSiteScopeDeclarationProvider = CompositeKotlinDeclarationProvider.create(
+            buildList {
+                add(project.createDeclarationProvider(useSiteAnalysisScope))
+                extensionTools.mapTo(this) { it.declarationProvider }
+            }
+        )
+
+        useSitePackageProvider = CompositeKotlinPackageProvider.create(
+            buildList {
+                add(project.createPackageProvider(useSiteAnalysisScope))
+                extensionTools.mapTo(this) { it.packageProvider }
+            }
+        )
+    }
+
 
     fun getScopeSessionFor(session: FirSession): ScopeSession = withValidityAssertion { firResolveSession.getScopeSessionFor(session) }
 

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirScopeProvider.kt

@@ -252,6 +252,7 @@ internal class KtFirScopeProvider(
             builder,
             analysisSession.useSiteAnalysisScope,
             analysisSession.useSiteScopeDeclarationProvider,
+            analysisSession.useSitePackageProvider,
             analysisSession.targetPlatform
         )
     }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/scopes/KtFirPackageScope.kt

@@ -18,7 +18,7 @@ import org.jetbrains.kotlin.analysis.api.symbols.KtClassifierSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtConstructorSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtPackageSymbol
 import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProvider
-import org.jetbrains.kotlin.analysis.providers.createPackageProvider
+import org.jetbrains.kotlin.analysis.providers.KotlinPackageProvider
 import org.jetbrains.kotlin.fir.extensions.FirExtensionService
 import org.jetbrains.kotlin.fir.extensions.declarationGenerators
 import org.jetbrains.kotlin.fir.extensions.extensionService
@@ -33,12 +33,11 @@ internal class KtFirPackageScope(
     private val builder: KtSymbolByFirBuilder,
     private val searchScope: GlobalSearchScope,
     private val declarationProvider: KotlinDeclarationProvider,
+    private val packageProvider: KotlinPackageProvider,
     private val targetPlatform: TargetPlatform,
 ) : KtScope {
     override val token: KtLifetimeToken get() = builder.token
 
-    private val packageProvider = project.createPackageProvider(searchScope)
-
     private val firScope: FirPackageMemberScope by lazy(LazyThreadSafetyMode.PUBLICATION) {
         FirPackageMemberScope(fqName, builder.rootSession)
     }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/project/structure/CompositeKotlinDeclarationProvider.kt

@@ -12,7 +12,7 @@ import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 
-internal class CompositeKotlinDeclarationProvider
+class CompositeKotlinDeclarationProvider
 private constructor(
     private val providers: List<KotlinDeclarationProvider>
 ) : KotlinDeclarationProvider() {

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/project/structure/CompositeKotlinPackageProvider.kt

@@ -10,7 +10,7 @@ import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.platform.TargetPlatform
 
-internal class CompositeKotlinPackageProvider
+class CompositeKotlinPackageProvider
 private constructor(
     private val providers: List<KotlinPackageProvider>
 ) : KotlinPackageProvider() {

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/resolve/extensions/LLFirResolveExtensionTool.kt

@@ -34,14 +34,14 @@ import java.util.concurrent.ConcurrentHashMap
  *
  * Caches generated [KtResolveExtensionFile]s, creates [KotlinDeclarationProvider], [KotlinPackageProvider], [LLFirSymbolProviderNameCache] needed for the [org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider].
  */
-internal abstract class LLFirResolveExtensionTool : FirSessionComponent {
+abstract class LLFirResolveExtensionTool : FirSessionComponent {
     abstract val modificationTrackers: List<ModificationTracker>
     abstract val declarationProvider: KotlinDeclarationProvider
     abstract val packageProvider: KotlinPackageProvider
-    abstract val symbolNameCache: LLFirSymbolProviderNameCache
+    internal abstract val symbolNameCache: LLFirSymbolProviderNameCache
 }
 
-internal val FirSession.llResolveExtensionTool: LLFirResolveExtensionTool? by FirSession.nullableSessionComponentAccessor()
+val FirSession.llResolveExtensionTool: LLFirResolveExtensionTool? by FirSession.nullableSessionComponentAccessor()
 
 internal class LLFirNonEmptyResolveExtensionTool(
     session: LLFirSession,