Implement light classes for Kotlin scripts

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/IDELightClassGenerationSupport.kt

@@ -27,10 +27,7 @@ import org.jetbrains.kotlin.asJava.LightClassBuilder
 import org.jetbrains.kotlin.asJava.LightClassGenerationSupport
 import org.jetbrains.kotlin.asJava.builder.ClsWrapperStubPsiFactory
 import org.jetbrains.kotlin.asJava.builder.LightClassDataHolder
-import org.jetbrains.kotlin.asJava.classes.FakeLightClassForFileOfPackage
-import org.jetbrains.kotlin.asJava.classes.KtLightClass
-import org.jetbrains.kotlin.asJava.classes.KtLightClassForFacade
-import org.jetbrains.kotlin.asJava.classes.KtLightClassForSourceDeclaration
+import org.jetbrains.kotlin.asJava.classes.*
 import org.jetbrains.kotlin.asJava.finder.JavaElementFinder
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.fileClasses.JvmFileClassUtil
@@ -84,6 +81,14 @@ class IDELightClassGenerationSupport(private val project: Project) : LightClassG
         )
     }
 
+    override fun createDataHolderForScript(script: KtScript, builder: LightClassBuilder): LightClassDataHolder.ForScript {
+        return LazyLightClassDataHolder.ForScript(
+                builder,
+                exactContextProvider = { IDELightClassContexts.contextForScript(script) },
+                dummyContextProvider = { IDELightClassContexts.lightContextForScript(script) }
+        )
+    }
+
     override fun findClassOrObjectDeclarations(fqName: FqName, searchScope: GlobalSearchScope): Collection<KtClassOrObject> {
         return runReadAction {
             KotlinFullClassNameIndex.getInstance().get(fqName.asString(), project, sourceAndClassFiles(searchScope, project))
@@ -131,6 +136,8 @@ class IDELightClassGenerationSupport(private val project: Project) : LightClassG
         return null
     }
 
+    override fun getLightClassForScript(script: KtScript): KtLightClassForScript? = KtLightClassForScript.create(script)
+
     private fun withFakeLightClasses(
             lightClassForFacade: KtLightClassForFacade?,
             facadeFiles: List<KtFile>
@@ -155,6 +162,12 @@ class IDELightClassGenerationSupport(private val project: Project) : LightClassG
         }
     }
 
+    override fun getScriptClasses(scriptFqName: FqName, scope: GlobalSearchScope): Collection<PsiClass> {
+        return KotlinScriptFqnIndex.instance.get(scriptFqName.asString(), project, scope).mapNotNull {
+            getLightClassForScript(it)
+        }
+    }
+
     override fun getKotlinInternalClasses(fqName: FqName, scope: GlobalSearchScope): Collection<PsiClass> {
         if (fqName.isRoot) return emptyList()
 

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/lightClasses/IDELightClassContexts.kt

@@ -129,6 +129,21 @@ object IDELightClassContexts {
         return IDELightClassConstructionContext(resolveSession.bindingContext, resolveSession.moduleDescriptor, EXACT)
     }
 
+    fun contextForScript(script: KtScript): LightClassConstructionContext {
+        val resolutionFacade = script.getResolutionFacade()
+        val bindingContext = resolutionFacade.analyze(script)
+
+        val descriptor = bindingContext[BindingContext.SCRIPT, script]
+        if (descriptor == null) {
+            LOG.warn("No script descriptor in context for script: " + script.getElementTextWithContext())
+            return IDELightClassConstructionContext(bindingContext, resolutionFacade.moduleDescriptor, EXACT)
+        }
+
+        ForceResolveUtil.forceResolveAllContents(descriptor)
+
+        return IDELightClassConstructionContext(bindingContext, resolutionFacade.moduleDescriptor, EXACT)
+    }
+
     fun lightContextForClassOrObject(classOrObject: KtClassOrObject): LightClassConstructionContext? {
         if (!isDummyResolveApplicable(classOrObject)) return null
 
@@ -139,6 +154,26 @@ object IDELightClassContexts {
         return IDELightClassConstructionContext(resolveSession.bindingContext, resolveSession.moduleDescriptor, LIGHT)
     }
 
+    fun lightContextForFacade(files: List<KtFile>): LightClassConstructionContext {
+        val representativeFile = files.first()
+        val resolveSession = setupAdHocResolve(representativeFile.project, representativeFile.getResolutionFacade().moduleDescriptor, files)
+
+        forceResolvePackageDeclarations(files, resolveSession)
+
+        return IDELightClassConstructionContext(resolveSession.bindingContext, resolveSession.moduleDescriptor, LIGHT)
+    }
+
+    fun lightContextForScript(script: KtScript): LightClassConstructionContext {
+        val resolveSession = setupAdHocResolve(
+                script.project,
+                script.getResolutionFacade().moduleDescriptor,
+                listOf(script.containingKtFile))
+
+        ForceResolveUtil.forceResolveAllContents(resolveSession.resolveToDescriptor(script))
+
+        return IDELightClassConstructionContext(resolveSession.bindingContext, resolveSession.moduleDescriptor, LIGHT)
+    }
+
     private fun isDummyResolveApplicable(classOrObject: KtClassOrObject): Boolean {
         if (classOrObject.hasLightClassMatchingErrors) return false
 
@@ -193,15 +228,6 @@ object IDELightClassContexts {
         return result
     }
 
-    fun lightContextForFacade(files: List<KtFile>): LightClassConstructionContext {
-        val representativeFile = files.first()
-        val resolveSession = setupAdHocResolve(representativeFile.project, representativeFile.getResolutionFacade().moduleDescriptor, files)
-
-        forceResolvePackageDeclarations(files, resolveSession)
-
-        return IDELightClassConstructionContext(resolveSession.bindingContext, resolveSession.moduleDescriptor, LIGHT)
-    }
-
     fun forceResolvePackageDeclarations(files: Collection<KtFile>, session: ResolveSession) {
         for (file in files) {
             if (file.isScript()) continue

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/lightClasses/LazyLightClassDataHolder.kt

@@ -74,6 +74,10 @@ sealed class LazyLightClassDataHolder(
             builder: LightClassBuilder, exactContextProvider: ExactLightClassContextProvider, dummyContextProvider: DummyLightClassContextProvider
     ) : LazyLightClassDataHolder(builder, exactContextProvider, dummyContextProvider), LightClassDataHolder.ForFacade
 
+    class ForScript(
+            builder: LightClassBuilder, exactContextProvider: ExactLightClassContextProvider, dummyContextProvider: DummyLightClassContextProvider
+    ) : LazyLightClassDataHolder(builder, exactContextProvider, dummyContextProvider), LightClassDataHolder.ForScript
+
     private inner class LazyLightClassData(
             findDelegate: (LightClassBuilderResult) -> PsiClass
     ) : LightClassData {