Implement annotation proxy generation for script dependencies resolvers

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

@@ -92,6 +92,13 @@ fun <TF> getFilePath(file: TF): String = when (file) {
     else -> throw IllegalArgumentException("Unsupported file type $file")
 }
 
+fun <TF> getFile(file: TF): File? = when (file) {
+    is PsiFile -> file.originalFile.run { File(virtualFile?.path) }
+    is VirtualFile -> File(file.path)
+    is File -> file
+    else -> throw IllegalArgumentException("Unsupported file type $file")
+}
+
 object StandardScriptDefinition : KotlinScriptDefinition {
     private val ARGS_NAME = Name.identifier("args")
 

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

@@ -19,63 +19,51 @@ package org.jetbrains.kotlin.script
 import com.intellij.openapi.util.text.StringUtil
 import org.jetbrains.kotlin.psi.*
 
-object SimpleAnnotationAst {
+internal class KtAnnotationWrapper(val psi: KtAnnotationEntry) {
+    val name: String
+        get() = (psi.typeReference?.typeElement as? KtUserType)?.referencedName.orAnonymous()
 
-    sealed class Node<out T>(val name: String) {
+    val valueArguments by lazy {
-        class empty(name: String) : Node<Unit>(name)
+        psi.valueArguments.map {
-        class str(name: String, val value: String) : Node<String>(name)
+            Pair(it.getArgumentName()?.toString(), convert(it.getArgumentExpression()!!))
-        class int(name: String, val value: Int) : Node<Int>(name)
+        }
-        class list<out T>(name: String, val value: List<Node<T>>) : Node<List<Node<T>>>(name)
-        class ann(name: String, val value: List<Node<Any>>) : Node<List<Node<Any>>>(name)
     }
 
-    class KtAnnotationWrapper(val psi: KtAnnotationEntry) {
+    internal fun String?.orAnonymous(kind: String = ""): String {
-        val name: String
+        return this ?: "<anonymous" + (if (kind.isNotBlank()) " $kind" else "") + ">"
-            get() = (psi.typeReference?.typeElement as? KtUserType)?.referencedName.orAnonymous()
+    }
 
-        val valueArguments by lazy {
+    internal fun convert(expression: KtExpression): Any? = when (expression) {
-            psi.valueArguments.map {
+        is KtStringTemplateExpression -> {
-                val name = it.getArgumentName()?.asName?.identifier.orAnonymous()
+            if (expression.entries.isEmpty())
-                convert(it.getArgumentExpression()!!)
+                ""
-            }
+            else if (expression.entries.size == 1)
+                convert(expression.entries[0])
+            else
+                ""
+            // TODO: parse expressions, etc. e.g.:
+            //      convertStringTemplateExpression(expression, parent, expression.entries.size - 1)
         }
+        else -> null
 
-        internal fun String?.orAnonymous(kind: String = ""): String {
+    }
-            return this ?: "<anonymous" + (if (kind.isNotBlank()) " $kind" else "") + ">"
+
+    internal fun convert(entry: KtStringTemplateEntry): Any? = when (entry) {
+        is KtStringTemplateEntryWithExpression -> convertOrEmpty(entry.expression)
+        is KtEscapeStringTemplateEntry -> entry.unescapedValue
+        else -> {
+            StringUtil.unescapeStringCharacters(entry.text)
         }
+    }
 
-        internal fun convert(expression: KtExpression): Node<Any> = when (expression) {
+    internal fun convertOrEmpty(expression: KtExpression?): Any? {
-            is KtStringTemplateExpression -> {
+        return if (expression != null) convert(expression) else null
-                if (expression.entries.isEmpty())
-                    SimpleAnnotationAst.Node.str(name, "")
-                else if (expression.entries.size == 1)
-                    convert(expression.entries[0])
-                else
-                    SimpleAnnotationAst.Node.str(name, "")
-                    // TODO: parse expressions, etc. e.g.:
-                    //      convertStringTemplateExpression(expression, parent, expression.entries.size - 1)
-            }
-            else -> Node.empty(name)
-
-        }
-
-        internal fun convert(entry: KtStringTemplateEntry): Node<Any> = when (entry) {
-            is KtStringTemplateEntryWithExpression -> convertOrEmpty(entry.expression)
-            is KtEscapeStringTemplateEntry -> Node.str(name, entry.unescapedValue)
-            else -> {
-                Node.str(name, StringUtil.unescapeStringCharacters(entry.text))
-            }
-        }
-
-        internal fun convertOrEmpty(expression: KtExpression?): Node<Any> {
-            return if (expression != null) convert(expression) else Node.empty(name)
-        }
     }
 }
 
-fun parseAnnotation(ann: KtAnnotationEntry): SimpleAnnotationAst.Node.list<Any> {
+fun parseAnnotation(ann: KtAnnotationEntry): Pair<String, Iterable<Any?>> {
-    val wann = SimpleAnnotationAst.KtAnnotationWrapper(ann)
+    val wann = KtAnnotationWrapper(ann)
     val vals = wann.valueArguments
-    return SimpleAnnotationAst.Node.list(wann.name, vals)
+    return Pair(wann.name, vals)
 }
 

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

@@ -29,19 +29,25 @@ import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtScript
 import org.jetbrains.kotlin.types.KotlinType
 import java.io.File
+import java.lang.reflect.InvocationHandler
+import java.lang.reflect.Method
+import java.lang.reflect.Proxy
+import java.util.*
 import kotlin.reflect.KClass
+import kotlin.reflect.memberProperties
 
 @Target(AnnotationTarget.CLASS)
 @Retention(AnnotationRetention.RUNTIME)
 annotation class ScriptFilePattern(val pattern: String)
 
 interface ScriptDependenciesResolver {
-    fun resolve(projectRoot: File?, scriptFile: File?, annotations: Iterable<KtAnnotationEntry>, context: Any?): KotlinScriptExternalDependencies? = null
+    fun resolve(scriptFile: File?, annotations: Iterable<Annotation>, context: Any?): KotlinScriptExternalDependencies? = null
 }
 
 @Target(AnnotationTarget.CLASS)
 @Retention(AnnotationRetention.RUNTIME)
-annotation class ScriptDependencyResolver(val resolver: KClass<out ScriptDependenciesResolver>)
+annotation class ScriptDependenciesResolverClass(val resolver: KClass<out ScriptDependenciesResolver>,
+                                                 vararg val supportedAnnotationClasses: KClass<out Any>)
 
 data class KotlinScriptDefinitionFromTemplate(val template: KClass<out Any>, val context: Any?) : KotlinScriptDefinition {
     override val name = template.simpleName!!
@@ -62,12 +68,51 @@ data class KotlinScriptDefinitionFromTemplate(val template: KClass<out Any>, val
     override fun getScriptName(script: KtScript): Name = ScriptNameUtil.fileNameWithExtensionStripped(script, KotlinParserDefinition.STD_SCRIPT_EXT)
 
     private val dependenciesResolvers by lazy {
-        template.annotations.mapNotNull { it as? ScriptDependencyResolver }.map { it.resolver.constructors.first().call() }
+        template.annotations.mapNotNull { it as? ScriptDependenciesResolverClass }.map { Pair(it, it.resolver.constructors.first().call()) }
+    }
+
+    private val supportedAnnotationClasses: List<KClass<out Any>> by lazy {
+        template.annotations.mapNotNull { it as? ScriptDependenciesResolverClass }
+                .flatMap {
+                    // it.supportedAnnotationClasses.asIterable() // TODO: find out why it fails with "java.lang.Class cannot be cast to kotlin.reflect.KClass"
+                    it.supportedAnnotationClasses.asIterable2()
+                }
+                .distinctBy { it.qualifiedName }
+    }
+
+    private fun Array<out KClass<out Any>>.asIterable2(): ArrayList<KClass<out Any>> {
+        val res = arrayListOf<KClass<out Any>>()
+        for (x in this) {
+            res.add(x)
+        }
+        return res
     }
 
     override fun <TF> getDependenciesFor(file: TF, project: Project): KotlinScriptExternalDependencies? {
         val fileAnnotations = getAnnotationEntries(file, project)
-        val fileDeps = dependenciesResolvers.mapNotNull { it.resolve(project.basePath?.let { File(it) }, File(getFilePath(file)), fileAnnotations, context) }
+                .map { KtAnnotationWrapper(it) }
+                .mapNotNull { wrappedAnn ->
+                    // TODO: consider advanced matching using semantic similar to actual resolving
+                    supportedAnnotationClasses.find { wrappedAnn.name == it.simpleName || wrappedAnn.name == it.qualifiedName }
+                        ?.let { Pair(it, wrappedAnn) }
+                }
+        val annotations = fileAnnotations.map {
+            try {
+                val handler = AnnProxyInvocationHandler(it.first, it.second.valueArguments)
+                val proxy = Proxy.newProxyInstance((template as Any).javaClass.classLoader, arrayOf(it.first.java), handler) as Annotation
+                Pair(it.first, proxy)
+            }
+            catch (ex: Exception) {
+                Pair(it.first, InvalidScriptResolverAnnotation(it.second.name, it.second.valueArguments, ex))
+            }
+        }
+        val fileDeps = dependenciesResolvers.mapNotNull { resolver ->
+            val supportedAnnotations = annotations.mapNotNull { ann ->
+                val annFQN = ann.first.qualifiedName
+                if (resolver.first.supportedAnnotationClasses.asIterable2().any { it.qualifiedName == annFQN }) ann.second else null
+            }
+            resolver.second.resolve(getFile(file), supportedAnnotations, context)
+        }
         return KotlinScriptExternalDependenciesUnion(fileDeps)
     }
 
@@ -87,10 +132,23 @@ data class KotlinScriptDefinitionFromTemplate(val template: KClass<out Any>, val
             else throw IllegalArgumentException("Unable to extract kotlin annotations from ${file.name} (${file.fileType})")
 
     private fun getAnnotationEntriesFromVirtualFile(file: VirtualFile, project: Project): Iterable<KtAnnotationEntry> {
-        val psifile: PsiFile = PsiManager.getInstance(project).findFile(file)
+        val psiFile: PsiFile = PsiManager.getInstance(project).findFile(file)
                                ?: throw java.lang.IllegalArgumentException("Unable to load PSI from ${file.canonicalPath}")
-        return getAnnotationEntriesFromPsiFile(psifile)
+        return getAnnotationEntriesFromPsiFile(psiFile)
     }
 }
 
+@Suppress("unused")
+class InvalidScriptResolverAnnotation(val name: String, val params: Iterable<Any?>, val error: Exception? = null) : Annotation
 
+class AnnProxyInvocationHandler<out K: KClass<out Any>>(val targetAnnClass: K, val annParams: List<Pair<String?, Any?>>) : InvocationHandler {
+    override fun invoke(proxy: Any?, method: Method?, params: Array<out Any>?): Any? {
+        if (method == null) return null
+        targetAnnClass.memberProperties.forEachIndexed { i, prop ->
+            if (prop.name == method.name) {
+                return if (i >= annParams.size) null else annParams[i].second
+            }
+        }
+        return null
+    }
+}

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

@@ -120,14 +120,12 @@ class TestKotlinScriptDependenciesResolver : ScriptDependenciesResolver {
 
     private val kotlinPaths by lazy { PathUtil.getKotlinPathsForCompiler() }
 
-    override fun resolve(projectRoot: File?, scriptFile: File?, annotations: Iterable<KtAnnotationEntry>, context: Any?): KotlinScriptExternalDependencies? {
+    override fun resolve(scriptFile: File?, annotations: Iterable<Annotation>, context: Any?): KotlinScriptExternalDependencies? {
-        val anns = annotations.map { parseAnnotation(it) }.filter { it.name == depends::class.simpleName }
+        val cp = annotations.flatMap {
-        val cp = anns.flatMap {
+            when (it) {
-            it.value.mapNotNull {
+                is depends -> listOf(if (it.path == "@{runtime}") kotlinPaths.runtimePath else File(it.path))
-                when (it) {
+                is InvalidScriptResolverAnnotation -> throw Exception("Invalid annotation ${it.name}", it.error)
-                    is SimpleAnnotationAst.Node.str -> if (it.value == "@{runtime}") kotlinPaths.runtimePath else File(it.value)
+                else -> throw Exception("Unknown annotation ${it.javaClass}")
-                    else -> null
-                }
             }
         }
         return object : KotlinScriptExternalDependencies {
@@ -143,15 +141,15 @@ class TestKotlinScriptDependenciesResolver : ScriptDependenciesResolver {
 }
 
 @ScriptFilePattern(".*\\.kts")
-@ScriptDependencyResolver(TestKotlinScriptDependenciesResolver::class)
+@ScriptDependenciesResolverClass(TestKotlinScriptDependenciesResolver::class, depends::class)
 abstract class ScriptWithIntParam(num: Int)
 
 @ScriptFilePattern(".*\\.kts")
-@ScriptDependencyResolver(TestKotlinScriptDependenciesResolver::class)
+@ScriptDependenciesResolverClass(TestKotlinScriptDependenciesResolver::class, depends::class)
 abstract class ScriptWithClassParam(param: TestParamClass)
 
 @ScriptFilePattern(".*\\.kts")
-@ScriptDependencyResolver(TestKotlinScriptDependenciesResolver::class)
+@ScriptDependenciesResolverClass(TestKotlinScriptDependenciesResolver::class, depends::class)
 abstract class ScriptWithBaseClass(num: Int, passthrough: Int) : TestDSLClassWithParam(passthrough)
 
 @Target(AnnotationTarget.FILE)