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> =
-        index.facadeFileMap[packageFqName]
-            ?.filter { ktFile ->
-                ktFile.virtualFile in scope
-            }
-            ?: emptyList()
+    override fun findFilesForFacadeByPackage(packageFqName: FqName): Collection<KtFile> {
+        return index.facadeFileMap[packageFqName].orEmpty().filter { it.virtualFile in scope }
+    }
 
     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) {