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 resolverMethodAnnotations =
-                dependencyResolver::class.memberFunctions.find { function ->
-                    function.name == resolveMethod.name &&
-                    sameSignature(function, resolveMethod)
-                }?.annotations?.filterIsInstance<AcceptedAnnotations>()
-        resolverMethodAnnotations?.flatMap {
-            it.supportedAnnotationClasses.toList()
+        val resolveFunctions = getResolveFunctions()
+
+        dependencyResolver.unwrap()::class.memberFunctions
+                .filter { function -> resolveFunctions.any { sameSignature(function, it) } }
+                .flatMap { it.annotations }
+                .filterIsInstance<AcceptedAnnotations>()
+                .flatMap { 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)
-    : kotlin.script.dependencies.DependenciesResolver {
+internal class ApiChangeDependencyResolverWrapper(
+        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 ->