Create visitor generator for FIR (KT-24062 in progress)

buildSrc/src/main/kotlin/copyright.kt

@@ -0,0 +1,65 @@
+/*
+ * Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package tasks
+
+import groovy.util.Node
+import groovy.util.XmlParser
+import org.gradle.api.DefaultTask
+import org.gradle.api.Project
+import org.gradle.api.tasks.Input
+import org.gradle.api.tasks.InputFile
+import org.gradle.api.tasks.OutputFile
+import org.gradle.api.tasks.TaskAction
+import java.io.File
+import java.util.*
+
+
+
+
+open class WriteCopyrightToFile : DefaultTask() {
+
+    @InputFile
+    var path = project.file("${project.rootDir}/.idea/copyright/apache.xml")
+
+    @OutputFile
+    var outputFile: File? = null
+
+    @Input
+    var commented: Boolean = true
+
+    @TaskAction
+    fun write() {
+        if (commented) {
+            outputFile!!.writeText(project.readCopyrightCommented())
+        } else {
+            outputFile!!.writeText(project.readCopyright())
+        }
+    }
+
+
+    fun Project.readCopyright(): String {
+        val file = rootDir.resolve(".idea/copyright/apache.xml")
+
+        assert(file.exists()) {
+            "File $file with copyright not found"
+        }
+
+
+        val xmlParser = XmlParser()
+        val node = xmlParser.parse(file)
+        assert(node.attribute("name") == "CopyrightManager") {
+            "Copyright format changed occasionally?"
+        }
+
+        val copyrightBlock = node.children().filterIsInstance<Node>().single()
+        val noticeNode = copyrightBlock.children().filterIsInstance<Node>().single { it.attribute("name") == "notice" }
+        return noticeNode.attribute("value").toString().replace("&#36;today.year", GregorianCalendar()[Calendar.YEAR].toString())
+    }
+
+    fun Project.readCopyrightCommented(): String {
+        return "/*\n" + readCopyright().prependIndent(" * ") + "\n */"
+    }
+}

compiler/fir/tree/build.gradle.kts

@@ -10,13 +10,39 @@ plugins {
 
 jvmTarget = "1.6"
 
+val generatorClasspath by configurations.creating
+
 dependencies {
     compile(project(":core:descriptors"))
+    generatorClasspath(project("visitors-generator"))
     // Necessary only to store bound PsiElement inside FirElement
     compileOnly(intellijCoreDep()) { includeJars("intellij-core", "annotations") }
 }
 
 sourceSets {
-    "main" { projectDefault() }
+    "main" {
+        projectDefault()
+        java.srcDir("visitors")
+    }
     "test" {}
 }
+
+val generateVisitors by tasks.creating(JavaExec::class) {
+    val generationRoot = "$projectDir/src/org/jetbrains/kotlin/fir/"
+    val output = "$projectDir/visitors"
+
+    val allSourceFiles = fileTree(generationRoot) {
+        include("**/*.kt")
+    }
+
+    inputs.files(allSourceFiles)
+    outputs.files(output)
+
+    classpath = generatorClasspath
+    args(generationRoot, output)
+    main = "org.jetbrains.kotlin.fir.visitors.generator.VisitorsGeneratorKt"
+}
+
+val compileKotlin by tasks
+
+compileKotlin.dependsOn(generateVisitors)

compiler/fir/tree/src/org/jetbrains/kotlin/fir/VisitedSupertype.kt

@@ -0,0 +1,10 @@
+/*
+ * Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir
+
+@Retention(AnnotationRetention.SOURCE)
+@Target(AnnotationTarget.TYPE)
+annotation class VisitedSupertype

compiler/fir/tree/visitors-generator/build.gradle.kts

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+import tasks.WriteCopyrightToFile
+
+plugins {
+    kotlin("jvm")
+    id("jps-compatible")
+}
+
+jvmTarget = "1.6"
+
+dependencies {
+    compile(project(":compiler:psi"))
+
+    compile(intellijCoreDep()) { includeJars("intellij-core") }
+    compile(intellijDep()) {
+        includeJars("trove4j", "picocontainer-1.2", "annotations", rootProject = rootProject)
+        isTransitive = false
+    }
+    compile(intellijDep()) { includeJars("guava", rootProject = rootProject) }
+}
+
+val writeCopyright by task<WriteCopyrightToFile> {
+    outputFile = file("$buildDir/copyright/notice.txt")
+    commented = true
+}
+
+val processResources by tasks
+processResources.dependsOn(writeCopyright)
+
+sourceSets {
+    "main" {
+        projectDefault()
+        resources.srcDir("$buildDir/copyright")
+    }
+    "test" {}
+}
+

compiler/fir/tree/visitors-generator/src/DataCollector.kt

@@ -0,0 +1,80 @@
+/*
+ * Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.visitors.generator
+
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.findDescendantOfType
+
+class DataCollector {
+
+    private val references = mutableMapOf<String, List<String>>()
+    private val packagePerClass = mutableMapOf<String, FqName>()
+
+    fun readFile(file: KtFile) {
+        file.acceptChildren(object : KtVisitorVoid() {
+
+            override fun visitKtElement(element: KtElement) {
+                super.visitKtElement(element)
+                element.acceptChildren(this)
+            }
+
+            override fun visitClass(klass: KtClass) {
+                val className = klass.name
+                if (klass.isInterface() && className != null) {
+                    packagePerClass[className] = file.packageFqName
+
+                    val manual = klass.superTypeListEntries.find {
+                        it.typeReference?.findDescendantOfType<KtAnnotationEntry> { annotationEntry ->
+                            annotationEntry.shortName?.asString() == VISITED_SUPERTYPE_ANNOTATION_NAME
+                        } != null
+                    }
+                    if (manual != null) {
+                        references[className] = listOfNotNull(manual.typeAsUserType?.referencedName)
+                    } else {
+                        references[className] =
+                                klass.superTypeListEntries.mapNotNull {
+                                    it.typeAsUserType?.referencedName
+                                }
+                    }
+                }
+
+                super.visitClass(klass)
+            }
+        })
+    }
+
+
+    private fun Map<String, List<String>>.computeBackReferences(): Map<String, List<String>> {
+        val result = mutableMapOf<String, List<String>>()
+
+        this.forEach { (k, v) ->
+            v.forEach {
+                result.merge(it, listOf(k)) { a, b -> a + b }
+            }
+        }
+        return result
+    }
+
+
+    fun computeResult(): ReferencesData {
+        val back = references.computeBackReferences()
+
+        val keysToKeep =
+            generateSequence(listOf(FIR_ELEMENT_CLASS_NAME)) { firName ->
+                firName.flatMap { back[it].orEmpty() }.takeUnless { it.isEmpty() }
+            }.flatten().toSet()
+
+        val cleanBack = back.filterKeys { it in keysToKeep }
+        return ReferencesData(
+            cleanBack.computeBackReferences(),
+            cleanBack,
+            packagePerClass.filterKeys { it in keysToKeep }.values.distinct()
+        )
+    }
+
+    data class ReferencesData(val direct: Map<String, List<String>>, val back: Map<String, List<String>>, val usedPackages: List<FqName>)
+}

compiler/fir/tree/visitors-generator/src/PsiSetup.kt

@@ -0,0 +1,70 @@
+/*
+ * Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.visitors.generator
+
+import com.intellij.core.CoreApplicationEnvironment
+import com.intellij.core.CoreProjectEnvironment
+import com.intellij.lang.MetaLanguage
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.extensions.Extensions
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.Disposer
+import org.jetbrains.kotlin.idea.KotlinFileType
+import org.jetbrains.kotlin.parsing.KotlinParserDefinition
+import org.jetbrains.kotlin.script.KotlinScriptDefinition
+import org.jetbrains.kotlin.script.ScriptDefinitionProvider
+import org.jetbrains.kotlin.script.StandardScriptDefinition
+
+internal data class PsiSetup(
+    val applicationEnvironment: CoreApplicationEnvironment,
+    val projectEnvironment: CoreProjectEnvironment,
+    val project: Project,
+    val disposable: Disposable
+)
+
+internal fun setup(): PsiSetup {
+    val disposable = Disposer.newDisposable()
+
+    val applicationEnvironment = CoreApplicationEnvironment(disposable, false)
+
+    val projectEnvironment = CoreProjectEnvironment(disposable, applicationEnvironment)
+
+    CoreApplicationEnvironment.registerExtensionPoint(Extensions.getRootArea(), MetaLanguage.EP_NAME, MetaLanguage::class.java)
+
+
+    applicationEnvironment.registerApplicationService(ScriptDefinitionProvider::class.java, NoopScriptDefinitionProvider())
+    applicationEnvironment.registerFileType(KotlinFileType.INSTANCE, "kt")
+    applicationEnvironment.registerParserDefinition(KotlinParserDefinition())
+
+
+    val project = projectEnvironment.project
+    return PsiSetup(applicationEnvironment, projectEnvironment, project, disposable)
+}
+
+internal inline fun <T> withPsiSetup(l: PsiSetup.() -> T): T {
+    val setup = setup()
+    val t = setup.l()
+    Disposer.dispose(setup.disposable)
+    return t
+}
+
+private class NoopScriptDefinitionProvider : ScriptDefinitionProvider {
+    override fun getDefaultScriptDefinition(): KotlinScriptDefinition {
+        return StandardScriptDefinition
+    }
+
+    override fun getKnownFilenameExtensions(): Sequence<String> {
+        return emptySequence()
+    }
+
+    override fun isScript(fileName: String): Boolean {
+        return false
+    }
+
+    override fun findScriptDefinition(fileName: String): KotlinScriptDefinition? {
+        return null
+    }
+}

compiler/fir/tree/visitors-generator/src/VisitorsGenerator.kt

@@ -0,0 +1,288 @@
+/*
+ * Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+package org.jetbrains.kotlin.fir.visitors.generator
+
+import com.intellij.openapi.vfs.StandardFileSystems
+import com.intellij.openapi.vfs.VirtualFileManager
+import com.intellij.openapi.vfs.local.CoreLocalFileSystem
+import com.intellij.psi.PsiManager
+import com.intellij.psi.SingleRootFileViewProvider
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.utils.Printer
+import java.io.File
+
+
+const val FIR_ELEMENT_CLASS_NAME = "FirElement"
+const val VISITOR_PACKAGE = "org.jetbrains.kotlin.fir.visitors"
+const val SIMPLE_VISITOR_NAME = "FirVisitor"
+const val UNIT_VISITOR_NAME = "FirVisitorVoid"
+const val VISITED_SUPERTYPE_ANNOTATION_NAME = "VisitedSupertype"
+
+const val WARNING_GENERATED_FILE =
+    """/** This file generated by :compiler:fir:tree:generateVisitors. DO NOT MODIFY MANUALLY! */"""
+
+fun main(args: Array<String>) {
+
+    val rootPath = File(args[0])
+    val output = File(args[1])
+
+    val packageDirectory = output.resolve(VISITOR_PACKAGE.replace('.', '/'))
+    packageDirectory.mkdirs()
+
+    withPsiSetup {
+        val psiManager = PsiManager.getInstance(project)
+        val vfm = VirtualFileManager.getInstance()
+
+
+        val dataCollector = DataCollector()
+
+        rootPath.walkTopDown()
+            .filter {
+                it.extension == "kt"
+            }.map {
+                (vfm.getFileSystem(StandardFileSystems.FILE_PROTOCOL) as CoreLocalFileSystem).findFileByIoFile(it)
+            }.flatMap {
+                SingleRootFileViewProvider(psiManager, it!!).allFiles.asSequence()
+            }
+            .filterIsInstance<KtFile>()
+            .forEach(dataCollector::readFile)
+
+
+        val data = dataCollector.computeResult()
+
+
+        packageDirectory
+            .resolve("${SIMPLE_VISITOR_NAME}Generated.kt")
+            .writeText(
+                SimpleVisitorGenerator(data).generate()
+            )
+
+        packageDirectory
+            .resolve("${UNIT_VISITOR_NAME}Generated.kt")
+            .writeText(
+                UnitVisitorGenerator(data).generate()
+            )
+    }
+}
+
+
+val String.classNameWithoutFir get() = this.removePrefix("Fir")
+
+
+fun DataCollector.ReferencesData.walkHierarchyTopDown(from: String, l: (p: String, e: String) -> Unit) {
+    val referents = back[from] ?: return
+    for (referent in referents) {
+        l(from, referent)
+        walkHierarchyTopDown(referent, l)
+    }
+}
+
+abstract class AbstractVisitorGenerator(val referencesData: DataCollector.ReferencesData) {
+    fun Printer.generateFunction(
+        name: String,
+        parameters: Map<String, String>,
+        returnType: String,
+        override: Boolean = false,
+        final: Boolean = false,
+        body: (Printer.() -> Unit)?
+    ) {
+
+        if (body == null) {
+            print("abstract ")
+        } else {
+            printIndent()
+            if (!final) {
+                printWithNoIndent("open ")
+            }
+            if (override) {
+                if (final) {
+                    printWithNoIndent("final ")
+                }
+                printWithNoIndent("override ")
+            }
+        }
+        printWithNoIndent("fun ", name, "(")
+        parameters
+            .flatMap { (a, b) ->
+                listOf(a, ": ", b, ", ")
+            }.dropLast(1)
+            .forEach {
+                printWithNoIndent(it)
+            }
+        printWithNoIndent(")")
+        if (returnType != "Unit") {
+            printWithNoIndent(": ", returnType)
+        }
+        if (body != null) {
+            printlnWithNoIndent(" {")
+            indented {
+                body()
+            }
+            println("}")
+        } else {
+            printlnWithNoIndent()
+        }
+        println()
+    }
+
+    protected inline fun Printer.indented(l: () -> Unit) {
+        pushIndent()
+        l()
+        popIndent()
+    }
+
+    protected fun Printer.generateCall(name: String, args: List<String>) {
+        printWithNoIndent(name, "(")
+        separatedOneLine(args, ", ")
+        printWithNoIndent(")")
+    }
+
+    protected fun Printer.separatedOneLine(iterable: Iterable<Any>, separator: Any) {
+        var first = true
+        for (element in iterable) {
+            if (!first) {
+                printWithNoIndent(separator)
+            } else {
+                first = false
+            }
+            printWithNoIndent(element)
+        }
+    }
+
+    protected fun Printer.generateDefaultImports() {
+        referencesData.usedPackages.forEach {
+            println("import ", it.asString(), ".*")
+        }
+        println()
+        println()
+    }
+
+    val String.safeName
+        get() = when (this) {
+            "class" -> "klass"
+            else -> this
+        }
+
+    fun generate(): String {
+        val builder = StringBuilder()
+        val printer = Printer(builder, "    ")
+        printer.apply {
+            println(javaClass.getResource("/notice.txt").readText())
+
+            println("package $VISITOR_PACKAGE")
+            println()
+            generateDefaultImports()
+            println(WARNING_GENERATED_FILE)
+
+            printer.generateContent()
+        }
+        return builder.toString()
+    }
+
+    abstract fun Printer.generateContent()
+
+}
+
+
+class SimpleVisitorGenerator(referencesData: DataCollector.ReferencesData) : AbstractVisitorGenerator(referencesData) {
+    override fun Printer.generateContent() {
+        println("abstract class $SIMPLE_VISITOR_NAME<R, D> {")
+        indented {
+            generateFunction(
+                "visitElement",
+                parameters = mapOf(
+                    "element" to FIR_ELEMENT_CLASS_NAME,
+                    "data" to "D"
+                ),
+                returnType = "R",
+                body = null
+            )
+            referencesData.walkHierarchyTopDown(from = FIR_ELEMENT_CLASS_NAME) { parent, element ->
+                generateVisit(element, parent)
+            }
+        }
+        println("}")
+    }
+
+    private fun Printer.generateVisit(className: String, parent: String) {
+        val shortcutName = className.classNameWithoutFir
+        val parameterName = shortcutName.decapitalize().safeName
+        generateFunction(
+            name = "visit$shortcutName",
+            parameters = mapOf(
+                parameterName to className,
+                "data" to "D"
+            ),
+            returnType = "R"
+        ) {
+            print("return ")
+            generateCall("visit${parent.classNameWithoutFir}", listOf(parameterName, "data"))
+            println()
+        }
+    }
+}
+
+
+class UnitVisitorGenerator(referencesData: DataCollector.ReferencesData) : AbstractVisitorGenerator(referencesData) {
+    override fun Printer.generateContent() {
+        println("abstract class $UNIT_VISITOR_NAME : $SIMPLE_VISITOR_NAME<Unit, Nothing?>() {")
+        indented {
+            generateFunction(
+                "visitElement",
+                parameters = mapOf(
+                    "element" to FIR_ELEMENT_CLASS_NAME
+                ),
+                returnType = "Unit",
+                body = null
+            )
+
+            referencesData.walkHierarchyTopDown(FIR_ELEMENT_CLASS_NAME) { parent, klass ->
+                generateVisit(klass, parent)
+            }
+
+            referencesData.back.keys.forEach {
+                generateTrampolineVisit(it)
+            }
+
+        }
+        println("}")
+    }
+
+
+    private fun Printer.generateVisit(className: String, parent: String) {
+        val shortcutName = className.classNameWithoutFir
+        val parameterName = shortcutName.decapitalize().safeName
+        generateFunction(
+            name = "visit$shortcutName",
+            parameters = mapOf(
+                parameterName to className
+            ),
+            returnType = "Unit"
+        ) {
+            printIndent()
+            generateCall("visit${parent.classNameWithoutFir}", listOf(parameterName, "null"))
+            println()
+        }
+    }
+
+    private fun Printer.generateTrampolineVisit(className: String) {
+        val shortcutName = className.classNameWithoutFir
+        val parameterName = shortcutName.decapitalize().safeName
+        generateFunction(
+            name = "visit$shortcutName",
+            parameters = mapOf(
+                parameterName to className,
+                "data" to "Nothing?"
+            ),
+            returnType = "Unit",
+            override = true,
+            final = true
+        ) {
+            printIndent()
+            generateCall("visit$shortcutName", listOf(parameterName))
+            println()
+        }
+    }
+}

settings.gradle

@@ -29,6 +29,7 @@ include ":kotlin-build-common",
         ":compiler:serialization",
         ":compiler:psi",
         ":compiler:fir:tree",
+        ":compiler:fir:tree:visitors-generator",
         ":compiler:fir:psi2fir",
         ":compiler:frontend",
         ":compiler:frontend.java",