AA/LC: introduce an API to find script files

analysis/analysis-api-providers/src/org/jetbrains/kotlin/analysis/providers/KotlinDeclarationProvider.kt

@@ -33,12 +33,15 @@ public abstract class KotlinDeclarationProvider {
     public abstract fun getTopLevelCallableNamesInPackage(packageFqName: FqName): Set<Name>
 
     public abstract fun findFilesForFacadeByPackage(packageFqName: FqName): Collection<KtFile>
+
     public abstract fun findFilesForFacade(facadeFqName: FqName): Collection<KtFile>
 
     /**
      * Currently we want only classes from libraries ([org.jetbrains.kotlin.analysis.decompiler.psi.file.KtClsFile])
      */
     public abstract fun findInternalFilesForFacade(facadeFqName: FqName): Collection<KtFile>
+
+    public abstract fun findFilesForScript(scriptFqName: FqName): Collection<KtScript>
 }
 
 public abstract class KotlinDeclarationProviderFactory {

analysis/analysis-api-providers/src/org/jetbrains/kotlin/analysis/providers/impl/KotlinStaticDeclarationIndex.kt

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.psi.*
 internal class KotlinStaticDeclarationIndex {
     internal val facadeFileMap: MutableMap<FqName, MutableSet<KtFile>> = mutableMapOf()
     internal val multiFileClassPartMap: MutableMap<FqName, MutableSet<KtFile>> = mutableMapOf()
+    internal val scriptMap: MutableMap<FqName, MutableSet<KtScript>> = mutableMapOf()
     internal val classMap: MutableMap<FqName, MutableSet<KtClassOrObject>> = mutableMapOf()
     internal val typeAliasMap: MutableMap<FqName, MutableSet<KtTypeAlias>> = mutableMapOf()
     internal val topLevelFunctionMap: MutableMap<FqName, MutableSet<KtNamedFunction>> = mutableMapOf()

analysis/analysis-api-providers/src/org/jetbrains/kotlin/analysis/providers/impl/KotlinStaticDeclarationProvider.kt

@@ -66,12 +66,9 @@ public class KotlinStaticDeclarationProvider internal constructor(
             .mapNotNullTo(mutableSetOf()) { it.nameAsName }
     }
 
-    override fun findFilesForFacadeByPackage(packageFqName: FqName): Collection<KtFile> =
+    override fun findFilesForFacadeByPackage(packageFqName: FqName): Collection<KtFile> {
-        index.facadeFileMap[packageFqName]
+        return index.facadeFileMap[packageFqName].orEmpty().filter { it.virtualFile in scope }
-            ?.filter { ktFile ->
+    }
-                ktFile.virtualFile in scope
-            }
-            ?: emptyList()
 
     override fun findFilesForFacade(facadeFqName: FqName): Collection<KtFile> {
         if (facadeFqName.shortNameOrSpecial().isSpecial) return emptyList()
@@ -83,6 +80,10 @@ public class KotlinStaticDeclarationProvider internal constructor(
         return index.multiFileClassPartMap[facadeFqName].orEmpty().filter { it.virtualFile in scope }
     }
 
+    override fun findFilesForScript(scriptFqName: FqName): Collection<KtScript> {
+        return index.scriptMap[scriptFqName].orEmpty().filter { it.containingKtFile.virtualFile in scope }
+    }
+
     override fun getTopLevelProperties(callableId: CallableId): Collection<KtProperty> =
         index.topLevelPropertyMap[callableId.packageName]
             ?.filter { ktProperty ->
@@ -165,6 +166,7 @@ public class KotlinStaticDeclarationProviderFactory(
 
         override fun visitKtFile(file: KtFile) {
             addToFacadeFileMap(file)
+            file.script?.let { addToScriptMap(it) }
             super.visitKtFile(file)
         }
 
@@ -196,6 +198,12 @@ public class KotlinStaticDeclarationProviderFactory(
         }.add(file)
     }
 
+    private fun addToScriptMap(script: KtScript) {
+        index.scriptMap.computeIfAbsent(script.fqName) {
+            mutableSetOf()
+        }.add(script)
+    }
+
     private fun addToClassMap(classOrObject: KtClassOrObject) {
         classOrObject.getClassId()?.let { classId ->
             index.classMap.computeIfAbsent(classId.packageFqName) {

analysis/light-classes-base/src/org/jetbrains/kotlin/asJava/KotlinAsJavaSupport.kt

@@ -45,6 +45,8 @@ abstract class KotlinAsJavaSupport {
 
     abstract fun findFilesForFacadeByPackage(packageFqName: FqName, searchScope: GlobalSearchScope): Collection<KtFile>
 
+    abstract fun findFilesForScript(scriptFqName: FqName, searchScope: GlobalSearchScope): Collection<KtScript>
+
     abstract fun findClassOrObjectDeclarations(fqName: FqName, searchScope: GlobalSearchScope): Collection<KtClassOrObject>
 
     abstract fun packageExists(fqName: FqName, scope: GlobalSearchScope): Boolean

analysis/light-classes-base/src/org/jetbrains/kotlin/asJava/KotlinAsJavaSupportBase.kt

@@ -207,9 +207,7 @@ abstract class KotlinAsJavaSupportBase<TModule>(protected val project: Project)
             return emptyList()
         }
 
-        return findFilesForPackage(scriptFqName.parent(), scope).mapNotNull { file ->
+        return findFilesForScript(scriptFqName, scope).mapNotNull { getLightClassForScript(it) }
-            file.script?.takeIf { it.fqName == scriptFqName }?.let { getLightClassForScript(it) }
-        }
     }
 }
 

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

@@ -61,6 +61,10 @@ private constructor(
         return providers.flatMapTo(mutableListOf()) { it.findInternalFilesForFacade(facadeFqName) }
     }
 
+    override fun findFilesForScript(scriptFqName: FqName): Collection<KtScript> {
+        return providers.flatMapTo(mutableListOf()) { it.findFilesForScript(scriptFqName) }
+    }
+
     companion object {
         fun create(providers: List<KotlinDeclarationProvider>): KotlinDeclarationProvider {
             return when (providers.size) {

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

@@ -24,4 +24,5 @@ internal object EmptyKotlinDeclarationProvider : KotlinDeclarationProvider() {
     override fun findFilesForFacadeByPackage(packageFqName: FqName) = emptyList<KtFile>()
     override fun findFilesForFacade(facadeFqName: FqName) = emptyList<KtFile>()
     override fun findInternalFilesForFacade(facadeFqName: FqName) = emptyList<KtFile>()
+    override fun findFilesForScript(scriptFqName: FqName) = emptyList<KtScript>()
 }

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

@@ -130,6 +130,9 @@ internal class FileBasedKotlinDeclarationProvider(val kotlinFile: KtFile) : Kotl
 
     override fun findInternalFilesForFacade(facadeFqName: FqName): Collection<KtFile> = emptyList()
 
+    override fun findFilesForScript(scriptFqName: FqName): Collection<KtScript> =
+        listOfNotNull(kotlinFile.script?.takeIf { it.fqName == scriptFqName })
+
     private inline fun <reified T : KtCallableDeclaration> getTopLevelCallables(callableId: CallableId): Collection<T> {
         require(callableId.classId == null)
         return getTopLevelDeclarations(callableId.packageName, callableId.callableName)

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

@@ -24,10 +24,7 @@ import org.jetbrains.kotlin.analysis.providers.KotlinPackageProvider
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSessionComponent
 import org.jetbrains.kotlin.load.kotlin.PackagePartClassUtils
-import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.*
-import org.jetbrains.kotlin.name.ClassId
-import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.platform.TargetPlatform
 import org.jetbrains.kotlin.psi.*
 import java.util.concurrent.ConcurrentHashMap
@@ -181,6 +178,15 @@ private class LLFirResolveExtensionToolDeclarationProvider(
         return emptyList()
     }
 
+    override fun findFilesForScript(scriptFqName: FqName): Collection<KtScript> = forbidAnalysis {
+        if (scriptFqName.isRoot) return emptyList()
+        val packageFqName = scriptFqName.parent()
+        return getDeclarationProvidersByPackage(packageFqName) { file ->
+            scriptFqName.shortName() == NameUtils.getScriptNameForFile(file.getFileName())
+        }
+            .mapNotNullTo(mutableListOf()) { it.kotlinFile.script }
+    }
+
     private inline fun getDeclarationProvidersByPackage(
         packageFqName: FqName,
         crossinline filter: (KtResolveExtensionFile) -> Boolean

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/SymbolKotlinAsJavaSupport.kt

@@ -55,6 +55,10 @@ class SymbolKotlinAsJavaSupport(project: Project) : KotlinAsJavaSupportBase<KtMo
         return project.createDeclarationProvider(searchScope).findFilesForFacadeByPackage(packageFqName)
     }
 
+    override fun findFilesForScript(scriptFqName: FqName, searchScope: GlobalSearchScope): Collection<KtScript> {
+        return project.createDeclarationProvider(searchScope).findFilesForScript(scriptFqName)
+    }
+
     private fun FqName.toClassIdSequence(): Sequence<ClassId> {
         var currentName = shortNameOrSpecial()
         if (currentName.isSpecial) return emptySequence()

compiler/cli/cli-base/src/org/jetbrains/kotlin/cli/jvm/compiler/CliKotlinAsJavaSupport.kt

@@ -112,6 +112,13 @@ class CliKotlinAsJavaSupport(project: Project, private val traceHolder: CliTrace
         }.orEmpty()
     }
 
+    override fun findFilesForScript(scriptFqName: FqName, searchScope: GlobalSearchScope): Collection<KtScript> {
+        return findFilesForPackage(scriptFqName.parent(), searchScope)
+            .mapNotNull { file ->
+                file.script?.takeIf { it.fqName == scriptFqName }
+            }
+    }
+
     override fun createFacadeForSyntheticFile(file: KtFile): KtLightClassForFacade = error("Should not be called")
     override fun declarationLocation(file: KtFile): DeclarationLocation = DeclarationLocation.ProjectSources
     override fun createInstanceOfDecompiledLightClass(classOrObject: KtClassOrObject): KtLightClass = error("Should not be called")