Implement Kotlin Script support in UAST

#KT-18353 Fixed

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUClass.kt

@@ -18,17 +18,17 @@ package org.jetbrains.uast.kotlin
 
 import com.intellij.psi.*
 import org.jetbrains.kotlin.asJava.classes.KtLightClass
+import org.jetbrains.kotlin.asJava.classes.KtLightClassForScript
 import org.jetbrains.kotlin.asJava.elements.KtLightMethod
 import org.jetbrains.kotlin.asJava.toLightMethods
 import org.jetbrains.kotlin.load.java.JvmAbi
-import org.jetbrains.kotlin.psi.KtEnumEntry
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.KtObjectDeclaration
 import org.jetbrains.uast.*
 import org.jetbrains.uast.java.AbstractJavaUClass
 import org.jetbrains.uast.kotlin.declarations.KotlinUMethod
 import org.jetbrains.uast.kotlin.declarations.UastLightIdentifier
 
-class KotlinUClass private constructor(
+open class KotlinUClass private constructor(
         psi: KtLightClass,
         override val uastParent: UElement?
 ) : AbstractJavaUClass(), PsiClass by psi {
@@ -37,9 +37,7 @@ class KotlinUClass private constructor(
 
     override val psi = unwrap<UClass, PsiClass>(psi)
 
-    override fun getOriginalElement(): PsiElement? {
+    override fun getOriginalElement(): PsiElement? = super.getOriginalElement()
-        return super.getOriginalElement()
-    }
 
     override fun getNameIdentifier() = UastLightIdentifier(psi, ktClass)
 
@@ -67,34 +65,12 @@ class KotlinUClass private constructor(
 
     override fun getMethods(): Array<UMethod> {
         val primaryConstructor = ktClass?.primaryConstructor?.toLightMethods()?.firstOrNull()
-        val initBlocks = ktClass?.getAnonymousInitializers() ?: emptyList()
 
         fun createUMethod(psiMethod: PsiMethod): UMethod {
-            return if (psiMethod is KtLightMethod && psiMethod.isConstructor && initBlocks.isNotEmpty()
+            return if (psiMethod is KtLightMethod &&
-                    && (primaryConstructor == null || psiMethod == primaryConstructor)) {
+                       psiMethod.isConstructor &&
-                object : KotlinUMethod(psiMethod, this@KotlinUClass) {
+                       (primaryConstructor == null || psiMethod == primaryConstructor)) {
-                    override val uastBody by lz {
+                KotlinPrimaryConstructorUMethod(ktClass, psiMethod, this)
-                        val initializers = ktClass?.getAnonymousInitializers() ?: return@lz UastEmptyExpression
-                        val containingMethod = this
-
-                        object : UBlockExpression {
-                            override val psi: PsiElement?
-                                 get() = null
-
-                            override val uastParent: UElement?
-                                get() = containingMethod
-
-                            override val annotations: List<UAnnotation>
-                                get() = emptyList()
-
-                            override val expressions by lz {
-                                initializers.map {
-                                    getLanguagePlugin().convertOpt<UExpression>(it.body, this) ?: UastEmptyExpression
-                                }
-                            }
-                        }
-                    }
-                }
             } else {
                 getLanguagePlugin().convert(psiMethod, this)
             }
@@ -112,11 +88,22 @@ class KotlinUClass private constructor(
     private fun PsiClass.isEnumEntryLightClass() = (this as? KtLightClass)?.kotlinOrigin is KtEnumEntry
 
     companion object {
-        fun create(psi: KtLightClass, containingElement: UElement?): UClass {
+        fun create(psi: KtLightClass, containingElement: UElement?): UClass = when (psi) {
-            return if (psi is PsiAnonymousClass)
+            is PsiAnonymousClass -> KotlinUAnonymousClass(psi, containingElement)
-                KotlinUAnonymousClass(psi, containingElement)
+            is KtLightClassForScript -> KotlinScriptUClass(psi, containingElement)
-            else
+            else -> KotlinUClass(psi, containingElement)
-                KotlinUClass(psi, containingElement)
+        }
+    }
+
+    class KotlinPrimaryConstructorUMethod(
+            private val ktClass: KtClassOrObject?,
+            override val psi: KtLightMethod,
+            override val uastParent: UElement?
+    ): KotlinUMethod(psi, uastParent) {
+        override val uastBody: UExpression? by lz {
+            ktClass?.getAnonymousInitializers()
+                    ?.takeIf { it.isNotEmpty() }
+                    ?.let { KotlinUBlockExpression.create(it, this) }
         }
     }
 }
@@ -128,9 +115,7 @@ class KotlinUAnonymousClass(
 
     override val psi: PsiAnonymousClass = unwrap<UAnonymousClass, PsiAnonymousClass>(psi)
 
-    override fun getOriginalElement(): PsiElement? {
+    override fun getOriginalElement(): PsiElement? = super<AbstractJavaUClass>.getOriginalElement()
-        return super<AbstractJavaUClass>.getOriginalElement()
-    }
 
     override fun getSuperClass(): UClass? = super<AbstractJavaUClass>.getSuperClass()
     override fun getFields(): Array<UField> = super<AbstractJavaUClass>.getFields()
@@ -143,4 +128,44 @@ class KotlinUAnonymousClass(
             val ktClassOrObject = (psi.originalElement as? KtLightClass)?.kotlinOrigin as? KtObjectDeclaration ?: return null 
             return UIdentifier(ktClassOrObject.getObjectKeyword(), this)
         }
+}
+
+class KotlinScriptUClass(
+        psi: KtLightClassForScript,
+        override val uastParent: UElement?
+) : AbstractJavaUClass(), PsiClass by psi {
+
+    override val psi = unwrap<UClass, KtLightClassForScript>(psi)
+
+    override fun getSuperClass(): UClass? = super.getSuperClass()
+
+    override fun getFields(): Array<UField> = super.getFields()
+
+    override fun getInitializers(): Array<UClassInitializer> = super.getInitializers()
+
+    override fun getInnerClasses(): Array<UClass> = super.getInnerClasses()
+
+    override fun getMethods(): Array<UMethod> = psi.methods.map(this::createUMethod).toTypedArray()
+
+    private fun createUMethod(method: PsiMethod): UMethod {
+        return if (method.isConstructor) {
+            KotlinScriptConstructorUMethod(psi.script, method as KtLightMethod, this)
+        }
+        else {
+            getLanguagePlugin().convert(method, this)
+        }
+    }
+
+    override fun getOriginalElement(): PsiElement? = psi.originalElement
+
+    class KotlinScriptConstructorUMethod(
+            script: KtScript,
+            override val psi: KtLightMethod,
+            override val uastParent: UElement?
+    ) : KotlinUMethod(psi, uastParent) {
+        override val uastBody: UExpression? by lz {
+            val initializers = script.declarations.filterIsInstance<KtScriptInitializer>()
+            KotlinUBlockExpression.create(initializers, this)
+        }
+    }
 }

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUFile.kt

@@ -46,14 +46,11 @@ class KotlinUFile(override val psi: KtFile, override val languagePlugin: UastLan
     override val imports by lz { psi.importDirectives.map { KotlinUImportStatement(it, this) } }
 
     override val classes by lz {
-        fun PsiClass.toUClass() = languagePlugin.convert<UClass>(this, this@KotlinUFile)
+        val facadeOrScriptClass = if (psi.isScript()) psi.script?.toLightClass() else psi.findFacadeClass()
+        val classes = psi.declarations.mapNotNull { (it as? KtClassOrObject)?.toLightClass()?.toUClass() }
 
-        val classes = psi.findFacadeClass()?.let { mutableListOf(it.toUClass()) } ?: mutableListOf()
+        (facadeOrScriptClass?.let { listOf(it.toUClass()) } ?: emptyList()) + classes
-
-        for (declaration in psi.declarations) {
-            (declaration as? KtClassOrObject)?.toLightClass()?.let { classes += it.toUClass() }
-        }
-
-        return@lz classes
     }
+
+    private fun PsiClass.toUClass() = languagePlugin.convert<UClass>(this, this@KotlinUFile)
 }

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinUBlockExpression.kt

@@ -16,13 +16,33 @@
 
 package org.jetbrains.uast.kotlin
 
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.psi.KtAnonymousInitializer
 import org.jetbrains.kotlin.psi.KtBlockExpression
-import org.jetbrains.uast.UBlockExpression
+import org.jetbrains.uast.*
-import org.jetbrains.uast.UElement
 
 class KotlinUBlockExpression(
         override val psi: KtBlockExpression,
         override val uastParent: UElement?
 ) : KotlinAbstractUExpression(), UBlockExpression, KotlinUElementWithType {
+
     override val expressions by lz { psi.statements.map { KotlinConverter.convertOrEmpty(it, this) } }
+
+    private class KotlinLazyUBlockExpression(
+            override val uastParent: UElement?,
+            expressionProducer: (expressionParent: UElement?) -> List<UExpression>
+    ) : UBlockExpression {
+        override val psi: PsiElement? = null
+        override val annotations: List<UAnnotation> = emptyList()
+        override val expressions by lz { expressionProducer(this) }
+    }
+
+    companion object {
+        fun create(initializers: List<KtAnonymousInitializer>, uastParent: UElement): UBlockExpression {
+            val languagePlugin = uastParent.getLanguagePlugin()
+            return KotlinLazyUBlockExpression(uastParent) { expressionParent ->
+                initializers.map { languagePlugin.convertOpt<UExpression>(it.body, expressionParent) ?: UastEmptyExpression }
+            }
+        }
+    }
 }

plugins/uast-kotlin/testData/InnerClasses.kt

@@ -0,0 +1,8 @@
+class Foo {
+    class Bar(val a: Int, val b: Int) {
+        fun getAPlusB() = a + b
+        class Baz {
+            fun doNothing() = Unit
+        }
+    }
+}

plugins/uast-kotlin/testData/InnerClasses.log.txt

@@ -0,0 +1,23 @@
+UFile (package = )
+    UClass (name = Foo)
+        UAnnotationMethod (name = Foo)
+        UClass (name = Bar)
+            UField (name = a)
+                UAnnotation (fqName = null)
+            UField (name = b)
+                UAnnotation (fqName = null)
+            UAnnotationMethod (name = getAPlusB)
+                UBinaryExpression (operator = +)
+                    USimpleNameReferenceExpression (identifier = a)
+                    USimpleNameReferenceExpression (identifier = b)
+            UAnnotationMethod (name = getA)
+            UAnnotationMethod (name = getB)
+            UAnnotationMethod (name = Bar)
+                UParameter (name = a)
+                    UAnnotation (fqName = null)
+                UParameter (name = b)
+                    UAnnotation (fqName = null)
+            UClass (name = Baz)
+                UAnnotationMethod (name = doNothing)
+                    USimpleNameReferenceExpression (identifier = Unit)
+                UAnnotationMethod (name = Baz)

plugins/uast-kotlin/testData/InnerClasses.render.txt

@@ -0,0 +1,15 @@
+public final class Foo {
+    public fun Foo() = UastEmptyExpression
+    public static final class Bar {
+        private final var a: int
+        private final var b: int
+        public final fun getAPlusB() : int = a + b
+        public final fun getA() : int = UastEmptyExpression
+        public final fun getB() : int = UastEmptyExpression
+        public fun Bar(a: int, b: int) = UastEmptyExpression
+        public static final class Baz {
+            public final fun doNothing() : void = Unit
+            public fun Baz() = UastEmptyExpression
+        }
+    }
+}

plugins/uast-kotlin/testData/SimpleScript.kts

@@ -0,0 +1,20 @@
+println("Hello World!")
+
+fun getBarOrNull(flag: Boolean): Bar? {
+    return if (flag) Bar(42) else null
+}
+
+class Bar(val a: Int) {
+    val b: Int = 0
+
+    fun getAPlusB() = a + b
+
+    class Baz {
+        fun doSomething() {
+
+        }
+    }
+}
+
+getBarOrNull(true)
+println("Goodbye World!")

plugins/uast-kotlin/testData/SimpleScript.log.txt

@@ -0,0 +1,49 @@
+UFile (package = )
+    UClass (name = SimpleScript)
+        UAnnotationMethod (name = getBarOrNull)
+            UParameter (name = flag)
+                UAnnotation (fqName = null)
+            UBlockExpression
+                UReturnExpression
+                    UIfExpression
+                        USimpleNameReferenceExpression (identifier = flag)
+                        UCallExpression (kind = UastCallKind(name='constructor_call'), argCount = 1))
+                            UIdentifier (Identifier (Bar))
+                            USimpleNameReferenceExpression (identifier = <init>)
+                            ULiteralExpression (value = 42)
+                        ULiteralExpression (value = null)
+        UAnnotationMethod (name = SimpleScript)
+            UParameter (name = p)
+                UAnnotation (fqName = null)
+            UBlockExpression
+                UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1))
+                    UIdentifier (Identifier (println))
+                    USimpleNameReferenceExpression (identifier = println)
+                    ULiteralExpression (value = "Hello World!")
+                UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1))
+                    UIdentifier (Identifier (getBarOrNull))
+                    USimpleNameReferenceExpression (identifier = getBarOrNull)
+                    ULiteralExpression (value = true)
+                UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1))
+                    UIdentifier (Identifier (println))
+                    USimpleNameReferenceExpression (identifier = println)
+                    ULiteralExpression (value = "Goodbye World!")
+        UClass (name = Bar)
+            UField (name = b)
+                UAnnotation (fqName = null)
+                ULiteralExpression (value = 0)
+            UField (name = a)
+                UAnnotation (fqName = null)
+            UAnnotationMethod (name = getB)
+            UAnnotationMethod (name = getAPlusB)
+                UBinaryExpression (operator = +)
+                    USimpleNameReferenceExpression (identifier = a)
+                    USimpleNameReferenceExpression (identifier = b)
+            UAnnotationMethod (name = getA)
+            UAnnotationMethod (name = Bar)
+                UParameter (name = a)
+                    UAnnotation (fqName = null)
+            UClass (name = Baz)
+                UAnnotationMethod (name = doSomething)
+                    UBlockExpression
+                UAnnotationMethod (name = Baz)

plugins/uast-kotlin/testData/SimpleScript.render.txt

@@ -0,0 +1,23 @@
+public class SimpleScript : kotlin.script.templates.standard.ScriptTemplateWithArgs {
+    public final fun getBarOrNull(flag: boolean) : SimpleScript.Bar {
+        return if (flag) <init>(42) else null
+    }
+    public fun SimpleScript(p: java.lang.String[]) {
+        println("Hello World!")
+        getBarOrNull(true)
+        println("Goodbye World!")
+    }
+    public static final class Bar {
+        private final var b: int = 0
+        private final var a: int
+        public final fun getB() : int = UastEmptyExpression
+        public final fun getAPlusB() : int = a + b
+        public final fun getA() : int = UastEmptyExpression
+        public fun Bar(a: int) = UastEmptyExpression
+        public static final class Baz {
+            public final fun doSomething() : void {
+            }
+            public fun Baz() = UastEmptyExpression
+        }
+    }
+}

plugins/uast-kotlin/tests/AbstractKotlinUastTest.kt

@@ -19,6 +19,8 @@ import org.jetbrains.kotlin.config.CommonConfigurationKeys
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.addKotlinSourceRoot
 import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisHandlerExtension
+import org.jetbrains.kotlin.script.KotlinScriptDefinitionProvider
+import org.jetbrains.kotlin.script.StandardScriptDefinition
 import org.jetbrains.kotlin.test.ConfigurationKind
 import org.jetbrains.kotlin.test.KotlinTestUtils
 import org.jetbrains.kotlin.test.TestJdkKind
@@ -43,13 +45,14 @@ abstract class AbstractKotlinUastTest : AbstractUastTest() {
     private var kotlinCoreEnvironment: KotlinCoreEnvironment? = null
 
     override fun getVirtualFile(testName: String): VirtualFile {
-        val projectDir = TEST_KOTLIN_MODEL_DIR
+        val testFile = TEST_KOTLIN_MODEL_DIR.listFiles { pathname -> pathname.nameWithoutExtension == testName }.first()
-        val testFile = File(TEST_KOTLIN_MODEL_DIR, testName.substringBefore('/') + ".kt")
 
         super.initializeEnvironment(testFile)
 
         initializeKotlinEnvironment()
 
+        KotlinScriptDefinitionProvider.getInstance(project)?.addScriptDefinition(StandardScriptDefinition)
+
         val trace = CliLightClassGenerationSupport.NoScopeRecordCliBindingTrace()
 
         val environment = kotlinCoreEnvironment!!
@@ -60,7 +63,7 @@ abstract class AbstractKotlinUastTest : AbstractUastTest() {
         val vfs = VirtualFileManager.getInstance().getFileSystem(URLUtil.FILE_PROTOCOL)
 
         val ideaProject = project
-        ideaProject.baseDir = vfs.findFileByPath(projectDir.canonicalPath)
+        ideaProject.baseDir = vfs.findFileByPath(TEST_KOTLIN_MODEL_DIR.canonicalPath)
 
         return vfs.findFileByPath(testFile.canonicalPath)!!
     }

plugins/uast-kotlin/tests/SimpleKotlinRenderLogTest.kt

@@ -30,4 +30,8 @@ class SimpleKotlinRenderLogTest : AbstractKotlinRenderLogTest() {
     @Test fun testPropertyDelegate() = doTest("PropertyDelegate")
 
     @Test fun testPropertyWithAnnotation() = doTest("PropertyWithAnnotation")
+
+    @Test fun testInnerClasses() = doTest("InnerClasses")
+
+    @Test fun testSimpleScript() = doTest("SimpleScript")
 }