[FIR] Support sealed classes and interfaces in FIR tree generator

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/context/AbstractFirTreeBuilder.kt

@@ -35,6 +35,17 @@ abstract class AbstractFirTreeBuilder {
             elements += it
         }
 
+    protected fun sealedElement(
+        name: String,
+        kind: Element.Kind,
+        vararg dependencies: Element,
+        init: Element.() -> Unit = {}
+    ): Element {
+        return element(name, kind, *dependencies, init = init).apply {
+            isSealed = true
+        }
+    }
+
     val configurations: MutableMap<Element, () -> Unit> = mutableMapOf()
 
     fun applyConfigurations() {

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/model/Element.kt

@@ -35,10 +35,22 @@ interface AbstractElement : FieldContainer, KindOwner {
     val overridenFields: Map<Field, Map<Importable, Boolean>>
     val useNullableForReplace: Set<Field>
 
+    val isSealed: Boolean
+        get() = false
+
     override val allParents: List<KindOwner> get() = parents
 }
 
 class Element(val name: String, kind: Kind) : AbstractElement {
+    companion object {
+        private val allowedKinds = setOf(
+            Implementation.Kind.Interface,
+            Implementation.Kind.SealedInterface,
+            Implementation.Kind.AbstractClass,
+            Implementation.Kind.SealedClass
+        )
+    }
+
     override val fields = mutableSetOf<Field>()
     override val type: String = "Fir$name"
     override val packageName: String = BASE_PACKAGE + kind.packageName.let { if (it.isBlank()) it else "." + it }
@@ -50,13 +62,15 @@ class Element(val name: String, kind: Kind) : AbstractElement {
     override val parentsArguments = mutableMapOf<AbstractElement, MutableMap<Importable, Importable>>()
     override var kind: Implementation.Kind? = null
         set(value) {
-            if (value != Implementation.Kind.Interface && value != Implementation.Kind.AbstractClass) {
+            if (value !in allowedKinds) {
                 throw IllegalArgumentException(value.toString())
             }
             field = value
         }
     var _needTransformOtherChildren: Boolean = false
 
+    override var isSealed: Boolean = false
+
     override var baseTransformerType: Element? = null
     override val transformerType: Element get() = baseTransformerType ?: this
 

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/model/Implementation.kt

@@ -79,6 +79,8 @@ class Implementation(val element: Element, val name: String?) : FieldContainer,
         FinalClass("class", true),
         OpenClass("open class", true),
         AbstractClass("abstract class", false),
+        SealedClass("sealed class", false),
+        SealedInterface("sealed interface", false),
         Object("object", false)
     }
 }

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/printer/element.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.tree.generator.printer
 
 import org.jetbrains.kotlin.fir.tree.generator.context.AbstractFirTreeBuilder
 import org.jetbrains.kotlin.fir.tree.generator.model.*
+import org.jetbrains.kotlin.fir.tree.generator.model.Implementation.Kind
 import org.jetbrains.kotlin.fir.tree.generator.pureAbstractElementType
 import org.jetbrains.kotlin.fir.tree.generator.util.get
 import org.jetbrains.kotlin.util.SmartPrinter
@@ -35,7 +36,7 @@ fun Element.generateCode(generationPath: File): GeneratedFile {
 
 fun SmartPrinter.printElement(element: Element) {
     with(element) {
-        val isInterface = kind == Implementation.Kind.Interface
+        val isInterface = kind == Kind.Interface || kind == Kind.SealedInterface
 
         fun abstract() {
             if (!isInterface) {
@@ -53,7 +54,7 @@ fun SmartPrinter.printElement(element: Element) {
         if (typeArguments.isNotEmpty()) {
             print(typeArguments.joinToString(", ", "<", ">") { it.toString() })
         }
-        val needPureAbstractElement = !isInterface && !allParents.any { it.kind == Implementation.Kind.AbstractClass }
+        val needPureAbstractElement = !isInterface && !allParents.any { it.kind == Kind.AbstractClass || it.kind == Kind.SealedClass }
 
         if (parents.isNotEmpty() || needPureAbstractElement) {
             print(" : ")

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/printer/implementation.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.tree.generator.printer
 
 import org.jetbrains.kotlin.fir.tree.generator.model.*
+import org.jetbrains.kotlin.fir.tree.generator.model.Implementation.Kind
 import org.jetbrains.kotlin.fir.tree.generator.pureAbstractElementType
 import org.jetbrains.kotlin.util.SmartPrinter
 import org.jetbrains.kotlin.util.withIndent
@@ -57,8 +58,8 @@ fun SmartPrinter.printImplementation(implementation: Implementation) {
         print("${kind!!.title} $type")
         print(element.typeParameters)
 
-        val isInterface = kind == Implementation.Kind.Interface
+        val isInterface = kind == Kind.Interface || kind == Kind.SealedInterface
-        val isAbstract = kind == Implementation.Kind.AbstractClass
+        val isAbstract = kind == Kind.AbstractClass || kind == Kind.SealedClass
 
         fun abstract() {
             if (isAbstract) {
@@ -80,7 +81,7 @@ fun SmartPrinter.printImplementation(implementation: Implementation) {
         }
 
         print(" : ")
-        if (!isInterface && !allParents.any { it.kind == Implementation.Kind.AbstractClass }) {
+        if (!isInterface && !allParents.any { it.kind == Kind.AbstractClass || it.kind == Kind.SealedClass }) {
             print("${pureAbstractElementType.type}(), ")
         }
         print(allParents.joinToString { "${it.typeWithArguments}${it.kind.braces()}" })

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/printer/utils.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.tree.generator.printer
 import org.jetbrains.kotlin.fir.tree.generator.context.AbstractFirTreeBuilder
 import org.jetbrains.kotlin.fir.tree.generator.firImplementationDetailType
 import org.jetbrains.kotlin.fir.tree.generator.model.*
+import org.jetbrains.kotlin.fir.tree.generator.model.Implementation.Kind
 import org.jetbrains.kotlin.fir.tree.generator.pureAbstractElementType
 import java.io.File
 
@@ -85,7 +86,7 @@ private fun List<String>.filterRedundantImports(
 
 
 val KindOwner.needPureAbstractElement: Boolean
-    get() = (kind != Implementation.Kind.Interface) && !allParents.any { it.kind == Implementation.Kind.AbstractClass }
+    get() = (kind != Kind.Interface && kind != Kind.SealedInterface) && !allParents.any { it.kind == Kind.AbstractClass || it.kind == Kind.SealedClass }
 
 
 val Field.isVal: Boolean get() = this is FieldList || (this is FieldWithDefault && origin is FieldList) || !isMutable
@@ -126,9 +127,9 @@ fun Element.multipleUpperBoundsList(): String {
     } ?: " "
 }
 
-fun Implementation.Kind?.braces(): String = when (this) {
+fun Kind?.braces(): String = when (this) {
-    Implementation.Kind.Interface -> ""
+    Kind.Interface, Kind.SealedInterface -> ""
-    Implementation.Kind.OpenClass, Implementation.Kind.AbstractClass -> "()"
+    Kind.OpenClass, Kind.AbstractClass, Kind.SealedClass -> "()"
     else -> throw IllegalStateException(this.toString())
 }
 

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/util/InterfaceAbstractClassSolver.kt

@@ -18,6 +18,7 @@ fun configureInterfacesAndAbstractClasses(builder: AbstractFirTreeBuilder) {
     val solution = solve2sat(elements, elementMapping)
     processRequirementsFromConfig(solution, elementMapping)
     updateKinds(solution, elementMapping)
+    updateSealedKinds(elements)
 }
 
 private class ElementMapping(elements: Collection<KindOwner>) {
@@ -73,6 +74,20 @@ private fun updateKinds(solution: List<Boolean>, elementMapping: ElementMapping)
     }
 }
 
+private fun updateSealedKinds(elements: Collection<KindOwner>) {
+    for (element in elements) {
+        if (element is Element) {
+            if (element.isSealed) {
+                element.kind = when (element.kind) {
+                    Implementation.Kind.AbstractClass -> Implementation.Kind.SealedClass
+                    Implementation.Kind.Interface -> Implementation.Kind.SealedInterface
+                    else -> error("element $element with kind ${element.kind} can not be sealed")
+                }
+            }
+        }
+    }
+}
+
 private fun processRequirementsFromConfig(solution: MutableList<Boolean>, elementMapping: ElementMapping) {
     fun processParents(element: KindOwner) {
         val origin = element.origin