Highlight scripts according to ScriptLocation annotation in ScriptDefinition

idea/idea-analysis/src/org/jetbrains/kotlin/idea/highlighter/KotlinHighlightingUtil.kt

@@ -16,13 +16,19 @@
 
 package org.jetbrains.kotlin.idea.highlighter
 
+import com.intellij.ide.projectView.impl.ProjectRootsUtil.isInTestSource
+import com.intellij.ide.scratch.ScratchFileService
+import com.intellij.openapi.roots.ProjectRootManager
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.caches.project.NotUnderContentRootModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.core.script.scriptDependencies
 import org.jetbrains.kotlin.idea.util.ProjectRootsUtil
+import org.jetbrains.kotlin.idea.util.ProjectRootsUtil.isInContent
 import org.jetbrains.kotlin.psi.KtCodeFragment
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.script.getScriptDefinition
+import kotlin.script.experimental.location.ScriptExpectedLocation
 
 object KotlinHighlightingUtil {
     fun shouldHighlight(psiElement: PsiElement): Boolean {
@@ -33,7 +39,7 @@ object KotlinHighlightingUtil {
         }
 
         if (ktFile.isScript()) {
-            return ktFile.virtualFile.scriptDependencies != null
+            return shouldHighlightScript(ktFile)
         }
 
         return ProjectRootsUtil.isInProjectOrLibraryContent(ktFile) && ktFile.getModuleInfo() !is NotUnderContentRootModuleInfo
@@ -49,9 +55,33 @@ object KotlinHighlightingUtil {
         }
 
         if (ktFile.isScript()) {
-            return ktFile.virtualFile.scriptDependencies != null
+            return shouldHighlightScript(ktFile)
         }
 
         return ProjectRootsUtil.isInProjectSource(ktFile)
     }
+
+
+    @Suppress("DEPRECATION")
+    private fun shouldHighlightScript(ktFile: KtFile): Boolean {
+        if (ScratchFileService.isInScratchRoot(ktFile.virtualFile)) return true
+        if (ktFile.virtualFile.scriptDependencies == null) return false
+
+        val scriptScope = getScriptDefinition(ktFile)?.scriptExpectedLocations ?: return false
+        return when {
+            scriptScope.contains(ScriptExpectedLocation.Everywhere) -> true
+            scriptScope.contains(ScriptExpectedLocation.Project)
+                    && ProjectRootManager.getInstance(ktFile.project).fileIndex.isInContent(ktFile.virtualFile) -> true
+            scriptScope.contains(ScriptExpectedLocation.TestsOnly)
+                    && isInTestSource(ktFile) -> true
+            else -> return isInContent(
+                ktFile,
+                scriptScope.contains(ScriptExpectedLocation.SourcesOnly),
+                scriptScope.contains(ScriptExpectedLocation.Libraries),
+                scriptScope.contains(ScriptExpectedLocation.Libraries),
+                scriptScope.contains(ScriptExpectedLocation.Libraries)
+            )
+        }
+
+    }
 }

idea/idea-gradle/src/org/jetbrains/kotlin/idea/core/script/GradleScriptTemplateProvider.kt

@@ -33,6 +33,7 @@ import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtScript
 import org.jetbrains.kotlin.script.KotlinScriptDefinition
+import org.jetbrains.kotlin.script.KotlinScriptDefinitionFromAnnotatedTemplate
 import org.jetbrains.plugins.gradle.config.GradleSettingsListenerAdapter
 import org.jetbrains.plugins.gradle.service.execution.GradleExecutionHelper
 import org.jetbrains.plugins.gradle.settings.DistributionType
@@ -50,6 +51,7 @@ import kotlin.script.experimental.dependencies.DependenciesResolver
 import kotlin.script.experimental.dependencies.DependenciesResolver.ResolveResult
 import kotlin.script.experimental.dependencies.ScriptDependencies
 import kotlin.script.experimental.dependencies.ScriptReport
+import kotlin.script.experimental.location.ScriptExpectedLocation
 import kotlin.script.templates.standard.ScriptTemplateWithArgs
 
 class GradleScriptDefinitionsContributor(private val project: Project) : ScriptDefinitionContributor {
@@ -206,7 +208,10 @@ class GradleScriptDefinitionsContributor(private val project: Project) : ScriptD
             templateClasspath,
             createEnvironment(gradleExeSettings),
             additionalResolverClasspath(gradleLibDir)
-        )
+        ).map {
+            // Expand scope for old gradle script definition
+            if (!it.scriptExpectedLocations.contains(ScriptExpectedLocation.Project)) GradleKotlinScriptDefinitionFromAnnotatedTemplate(it) else it
+        }
     }
 
     fun reloadIfNeccessary() {
@@ -250,6 +255,13 @@ internal class GradleSyncState {
     var isSyncInProgress: Boolean = false
 }
 
+class GradleKotlinScriptDefinitionFromAnnotatedTemplate(
+    base: KotlinScriptDefinitionFromAnnotatedTemplate
+) : KotlinScriptDefinitionFromAnnotatedTemplate(base.template, base.environment, base.templateClasspath) {
+    override val scriptExpectedLocations: List<ScriptExpectedLocation>
+        get() = listOf(ScriptExpectedLocation.Project)
+}
+
 class ReloadGradleTemplatesOnSync : ExternalSystemTaskNotificationListenerAdapter() {
     companion object {
         internal val gradleState = GradleSyncState()

idea/testData/script/definition/highlighting/acceptedAnnotations/template/template.kt

@@ -4,6 +4,8 @@ import java.io.File
 import kotlin.script.dependencies.*
 import kotlin.script.experimental.dependencies.*
 import kotlin.script.templates.ScriptTemplateDefinition
+import kotlin.script.experimental.location.ScriptExpectedLocation
+import kotlin.script.experimental.location.ScriptExpectedLocations
 
 class TestDependenciesResolver : DependenciesResolver {
     @kotlin.script.templates.AcceptedAnnotations(Anno::class)
@@ -26,5 +28,6 @@ annotation class Anno
 @Target(AnnotationTarget.FILE)
 annotation class Anno2
 
+@ScriptExpectedLocations([ScriptExpectedLocation.Everywhere])
 @ScriptTemplateDefinition(TestDependenciesResolver::class, scriptFilePattern = "script.kts")
 class Template

idea/tests/org/jetbrains/kotlin/idea/script/TestScriptDefinitionContributors.kt

@@ -25,6 +25,8 @@ import kotlin.script.experimental.dependencies.AsyncDependenciesResolver
 import kotlin.script.experimental.dependencies.DependenciesResolver
 import kotlin.script.experimental.dependencies.ScriptDependencies
 import kotlin.script.experimental.dependencies.asSuccess
+import kotlin.script.experimental.location.ScriptExpectedLocation
+import kotlin.script.experimental.location.ScriptExpectedLocations
 import kotlin.script.templates.ScriptTemplateDefinition
 
 
@@ -64,6 +66,7 @@ class FromTextDependenciesResolver : AsyncDependenciesResolver {
     }
 }
 
+@ScriptExpectedLocations([ScriptExpectedLocation.Everywhere])
 @Suppress("unused")
 @ScriptTemplateDefinition(FromTextDependenciesResolver::class, scriptFilePattern = "script.kts")
 class Template