Support @AcceptedAnnotations for different kinds of resolvers

compiler/frontend.script/src/org/jetbrains/kotlin/script/AsyncDependencyResolverWrapper.kt

@@ -23,7 +23,9 @@ import kotlin.script.dependencies.ScriptContents
 import kotlin.script.dependencies.experimental.AsyncDependenciesResolver
 
 // wraps AsyncDependenciesResolver to provide implementation for synchronous DependenciesResolver::resolve
-class AsyncDependencyResolverWrapper(private val delegate: AsyncDependenciesResolver): AsyncDependenciesResolver {
+class AsyncDependencyResolverWrapper(
+        override val delegate: AsyncDependenciesResolver
+): AsyncDependenciesResolver, DependencyResolverWrapper<AsyncDependenciesResolver> {
 
     override fun resolve(
             scriptContents: ScriptContents, environment: Environment

compiler/frontend.script/src/org/jetbrains/kotlin/script/KotlinScriptDefinitionFromAnnotatedTemplate.kt

@@ -99,22 +99,31 @@ open class KotlinScriptDefinitionFromAnnotatedTemplate(
     override val acceptedAnnotations: List<KClass<out Annotation>> by lazy {
 
         fun sameSignature(left: KFunction<*>, right: KFunction<*>): Boolean =
+                left.name == right.name &&
                 left.parameters.size == right.parameters.size &&
                 left.parameters.zip(right.parameters).all {
                     it.first.kind == KParameter.Kind.INSTANCE ||
                     it.first.type == it.second.type
                 }
 
-        val resolveMethod = ScriptDependenciesResolver::resolve
+        val resolveFunctions = getResolveFunctions()
-        val resolverMethodAnnotations =
+
-                dependencyResolver::class.memberFunctions.find { function ->
+        dependencyResolver.unwrap()::class.memberFunctions
-                    function.name == resolveMethod.name &&
+                .filter { function -> resolveFunctions.any { sameSignature(function, it) } }
-                    sameSignature(function, resolveMethod)
+                .flatMap { it.annotations }
-                }?.annotations?.filterIsInstance<AcceptedAnnotations>()
+                .filterIsInstance<AcceptedAnnotations>()
-        resolverMethodAnnotations?.flatMap {
+                .flatMap { it.supportedAnnotationClasses.toList() }
-            it.supportedAnnotationClasses.toList()
+                .distinctBy { it.qualifiedName }
+    }
+
+    private fun getResolveFunctions(): List<KFunction<*>> {
+        // DependenciesResolver::resolve, ScriptDependenciesResolver::resolve, AsyncDependenciesResolver::resolveAsync
+        return AsyncDependenciesResolver::class.memberFunctions.filter { it.name == "resolve" || it.name == "resolveAsync" }.also {
+            assert(it.size == 3) {
+                AsyncDependenciesResolver::class.memberFunctions
+                        .joinToString(prefix = "${AsyncDependenciesResolver::class.qualifiedName} api changed, fix this code") { it.name }
+            }
         }
-        ?: emptyList()
     }
 
     override val name = template.simpleName!!
@@ -147,4 +156,12 @@ open class KotlinScriptDefinitionFromAnnotatedTemplate(
     companion object {
         internal val log = Logger.getInstance(KotlinScriptDefinitionFromAnnotatedTemplate::class.java)
     }
+}
+
+interface DependencyResolverWrapper<T : ScriptDependenciesResolver> {
+    val delegate: T
+}
+
+fun ScriptDependenciesResolver.unwrap(): ScriptDependenciesResolver {
+    return if (this is DependencyResolverWrapper<*>) delegate.unwrap() else this
 }

compiler/frontend.script/src/org/jetbrains/kotlin/script/legacyWrappers.kt

@@ -24,7 +24,11 @@ import kotlin.script.dependencies.Environment
 import kotlin.script.dependencies.ScriptDependencies
 import kotlin.script.dependencies.ScriptReport
 
-internal class LegacyPackageDependencyResolverWrapper(val legacyResolver: ScriptDependenciesResolver) : kotlin.script.dependencies.DependenciesResolver {
+// NOTE: not implementing resolver wrapper since it was not implemented initially and this is deprecated API
+@Deprecated("Compatibility with deprecated API")
+internal class LegacyPackageDependencyResolverWrapper(
+        val legacyResolver: ScriptDependenciesResolver
+) : kotlin.script.dependencies.DependenciesResolver {
     override fun resolve(
             scriptContents: kotlin.script.dependencies.ScriptContents,
             environment: Environment
@@ -62,15 +66,16 @@ internal class LegacyPackageDependencyResolverWrapper(val legacyResolver: Script
     private fun ScriptContents.Position.convertPosition(): ScriptReport.Position = ScriptReport.Position(line, col)
 }
 
-internal class ApiChangeDependencyResolverWrapper(val legacyResolver: kotlin.script.dependencies.ScriptDependenciesResolver)
+internal class ApiChangeDependencyResolverWrapper(
-    : kotlin.script.dependencies.DependenciesResolver {
+        override val delegate: kotlin.script.dependencies.ScriptDependenciesResolver
+) : kotlin.script.dependencies.DependenciesResolver, DependencyResolverWrapper<kotlin.script.dependencies.ScriptDependenciesResolver> {
 
     override fun resolve(
             scriptContents: kotlin.script.dependencies.ScriptContents,
             environment: Environment
     ): DependenciesResolver.ResolveResult {
         val reports = ArrayList<ScriptReport>()
-        val legacyDeps = legacyResolver.resolve(
+        val legacyDeps = delegate.resolve(
                 scriptContents,
                 environment,
                 { sev, msg, pos ->

compiler/testData/script/acceptedAnnotations.kts

@@ -0,0 +1 @@
+@file: [TestAnno1 TestAnno2]

compiler/tests/org/jetbrains/kotlin/scripts/ScriptTemplateTest.kt

@@ -13,7 +13,6 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
 package org.jetbrains.kotlin.scripts
 
 import com.intellij.openapi.util.Disposer
@@ -46,6 +45,7 @@ import java.lang.Exception
 import java.lang.reflect.InvocationTargetException
 import java.net.URL
 import java.net.URLClassLoader
+import java.util.concurrent.Future
 import kotlin.reflect.KClass
 import kotlin.script.dependencies.*
 import kotlin.script.dependencies.DependenciesResolver.ResolveResult
@@ -260,6 +260,24 @@ class ScriptTemplateTest {
         assertEqualsTrimmed(NUM_4_LINE + FIB_SCRIPT_OUTPUT_TAIL, out)
     }
 
+    @Test
+    fun testAcceptedAnnotationsSync() {
+        val aClass = compileScript("acceptedAnnotations.kts", ScriptWithAcceptedAnnotationsSyncResolver::class, null)
+        Assert.assertNotNull(aClass)
+    }
+
+    @Test
+    fun testAcceptedAnnotationsAsync() {
+        val aClass = compileScript("acceptedAnnotations.kts", ScriptWithAcceptedAnnotationsAsyncResolver::class, null)
+        Assert.assertNotNull(aClass)
+    }
+
+    @Test
+    fun testAcceptedAnnotationsLegacy() {
+        val aClass = compileScript("acceptedAnnotations.kts", ScriptWithAcceptedAnnotationsLegacyResolver::class, null)
+        Assert.assertNotNull(aClass)
+    }
+
     @Test
     fun testSmokeScriptException() {
         val aClass = compileScript("smoke_exception.kts", ScriptWithArrayParam::class)
@@ -349,14 +367,15 @@ open class TestKotlinScriptDummyDependenciesResolver : DependenciesResolver {
             imports = listOf("org.jetbrains.kotlin.scripts.DependsOn", "org.jetbrains.kotlin.scripts.DependsOnTwo")
         ).asSuccess()
     }
-
-    protected fun classpathFromClassloader(): List<File> =
-            (TestKotlinScriptDependenciesResolver::class.java.classLoader as? URLClassLoader)?.urLs
-                    ?.mapNotNull(URL::toFile)
-                    ?.filter { it.path.contains("out") && it.path.contains("test") }
-            ?: emptyList()
 }
 
+private fun classpathFromClassloader(): List<File> =
+        (TestKotlinScriptDependenciesResolver::class.java.classLoader as? URLClassLoader)?.urLs
+                ?.mapNotNull(URL::toFile)
+                ?.filter { it.path.contains("out") && it.path.contains("test") }
+        ?: emptyList()
+
+
 open class TestKotlinScriptDependenciesResolver : TestKotlinScriptDummyDependenciesResolver() {
 
     private val kotlinPaths by lazy { PathUtil.getKotlinPathsForCompiler() }
@@ -417,6 +436,63 @@ class TestAsyncResolver : TestKotlinScriptDependenciesResolver(), AsyncDependenc
             super<AsyncDependenciesResolver>.resolve(scriptContents, environment)
 }
 
+@Target(AnnotationTarget.FILE)
+annotation class TestAnno1
+@Target(AnnotationTarget.FILE)
+annotation class TestAnno2
+@Target(AnnotationTarget.FILE)
+annotation class TestAnno3
+
+private val annotationFqNames = listOf(TestAnno1::class, TestAnno2::class, TestAnno3::class).map { it.qualifiedName!! }
+
+interface AcceptedAnnotationsCheck {
+    fun checkHasAnno1Annotation(scriptContents: ScriptContents): ResolveResult.Success {
+        val actualAnnotations = scriptContents.annotations
+        Assert.assertTrue(
+                "Loaded annotation: $actualAnnotations",
+                actualAnnotations.single().annotationClass.qualifiedName == TestAnno1::class.qualifiedName
+        )
+
+        return ScriptDependencies(
+                classpath = classpathFromClassloader(),
+                imports = annotationFqNames
+        ).asSuccess()
+    }
+}
+
+class TestAcceptedAnnotationsSyncResolver: DependenciesResolver, AcceptedAnnotationsCheck {
+    @AcceptedAnnotations(TestAnno1::class, TestAnno3::class)
+    override fun resolve(scriptContents: ScriptContents, environment: Environment): ResolveResult {
+        return checkHasAnno1Annotation(scriptContents)
+    }
+}
+
+class TestAcceptedAnnotationsAsyncResolver: AsyncDependenciesResolver, AcceptedAnnotationsCheck {
+    @AcceptedAnnotations(TestAnno1::class, TestAnno3::class)
+    override suspend fun resolveAsync(scriptContents: ScriptContents, environment: Environment): ResolveResult {
+        return checkHasAnno1Annotation(scriptContents)
+    }
+}
+
+class TestAcceptedAnnotationsLegacyResolver: ScriptDependenciesResolver, AcceptedAnnotationsCheck {
+    @AcceptedAnnotations(TestAnno1::class, TestAnno3::class)
+    override fun resolve(
+            script: ScriptContents,
+            environment: Environment?,
+            report: (ScriptDependenciesResolver.ReportSeverity, String, ScriptContents.Position?) -> Unit,
+            previousDependencies: KotlinScriptExternalDependencies?
+    ): Future<KotlinScriptExternalDependencies?> {
+        checkHasAnno1Annotation(script)
+        return object : KotlinScriptExternalDependencies {
+            override val classpath: Iterable<File>
+                get() = classpathFromClassloader()
+
+            override val imports: Iterable<String>
+                get() = annotationFqNames
+        }.asFuture()
+    }
+}
+
 @ScriptTemplateDefinition(
         scriptFilePattern =".*\\.kts",
         resolver = TestKotlinScriptDummyDependenciesResolver::class)
@@ -474,6 +550,15 @@ abstract class ScriptReportingErrors(val num: Int)
 @ScriptTemplateDefinition(resolver = TestAsyncResolver::class)
 abstract class ScriptWithAsyncResolver(val num: Int)
 
+@ScriptTemplateDefinition(resolver = TestAcceptedAnnotationsSyncResolver::class)
+abstract class ScriptWithAcceptedAnnotationsSyncResolver
+
+@ScriptTemplateDefinition(resolver = TestAcceptedAnnotationsAsyncResolver::class)
+abstract class ScriptWithAcceptedAnnotationsAsyncResolver
+
+@ScriptTemplateDefinition(resolver = TestAcceptedAnnotationsLegacyResolver::class)
+abstract class ScriptWithAcceptedAnnotationsLegacyResolver
+
 @Target(AnnotationTarget.FILE)
 @Retention(AnnotationRetention.RUNTIME)
 annotation class DependsOn(val path: String)

idea/testData/script/definition/highlighting/acceptedAnnotations/script.kts

@@ -0,0 +1,3 @@
+@file: [Anno custom.scriptDefinition.Anno2 <error>NonExistentAnno</error>]
+
+val c = 3

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

@@ -0,0 +1,29 @@
+package custom.scriptDefinition
+
+import java.io.File
+import kotlin.script.dependencies.*
+import kotlin.script.templates.ScriptTemplateDefinition
+
+class TestDependenciesResolver : DependenciesResolver {
+    @kotlin.script.templates.AcceptedAnnotations(Anno::class)
+    override fun resolve(
+            scriptContents: ScriptContents,
+            environment: Environment
+    ): DependenciesResolver.ResolveResult {
+        val annoFQN = Anno::class.qualifiedName!!
+        assert(scriptContents.annotations.single().annotationClass.qualifiedName == annoFQN)
+        return ScriptDependencies(
+                classpath = listOf(environment["template-classes"] as File),
+                imports = listOf(annoFQN)
+        ).asSuccess()
+    }
+}
+
+@Target(AnnotationTarget.FILE)
+annotation class Anno
+
+@Target(AnnotationTarget.FILE)
+annotation class Anno2
+
+@ScriptTemplateDefinition(TestDependenciesResolver::class, scriptFilePattern = "script.kts")
+class Template

idea/tests/org/jetbrains/kotlin/idea/script/ScriptConfigurationHighlightingTestGenerated.java

@@ -32,6 +32,12 @@ import java.util.regex.Pattern;
 @TestDataPath("$PROJECT_ROOT")
 @RunWith(JUnit3RunnerWithInners.class)
 public class ScriptConfigurationHighlightingTestGenerated extends AbstractScriptConfigurationHighlightingTest {
+    @TestMetadata("acceptedAnnotations")
+    public void testAcceptedAnnotations() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/script/definition/highlighting/acceptedAnnotations/");
+        doTest(fileName);
+    }
+
     @TestMetadata("additionalImports")
     public void testAdditionalImports() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("idea/testData/script/definition/highlighting/additionalImports/");