[FIR generator] Write kDocs for accept & transform methods in FirElement

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/FirElement.kt

@@ -19,18 +19,53 @@ import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 interface FirElement {
     val source: KtSourceElement?
 
+    /**
+     * Runs the provided [visitor] on the FIR subtree with the root at this node.
+     *
+     * @param visitor The visitor to accept.
+     * @param data An arbitrary context to pass to each invocation of [visitor]'s methods.
+     * @return The value returned by the topmost `visit*` invocation.
+     */
     fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R =
         visitor.visitElement(this, data)
 
+    /**
+     * Runs the provided [transformer] on the FIR subtree with the root at this node.
+     *
+     * @param transformer The transformer to use.
+     * @param data An arbitrary context to pass to each invocation of [transformer]'s methods.
+     * @return The transformed node.
+     */
     @Suppress("UNCHECKED_CAST")
     fun <E : FirElement, D> transform(transformer: FirTransformer<D>, data: D): E =
         transformer.transformElement(this, data) as E
 
     fun accept(visitor: FirVisitorVoid) = accept(visitor, null)
 
+    /**
+     * Runs the provided [visitor] on subtrees with roots in this node's children.
+     *
+     * Basically, calls `accept(visitor, data)` on each child of this node.
+     *
+     * Does **not** run [visitor] on this node itself.
+     *
+     * @param visitor The visitor for children to accept.
+     * @param data An arbitrary context to pass to each invocation of [visitor]'s methods.
+     */
     fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D)
 
     fun acceptChildren(visitor: FirVisitorVoid) = acceptChildren(visitor, null)
 
+    /**
+     * Recursively transforms this node's children *in place* using [transformer].
+     *
+     * Basically, executes `this.child = this.child.transform(transformer, data)` for each child of this node.
+     *
+     * Does **not** run [transformer] on this node itself.
+     *
+     * @param transformer The transformer to use for transforming the children.
+     * @param data An arbitrary context to pass to each invocation of [transformer]'s methods.
+     * @return `this`
+     */
     fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirElement
 }

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

@@ -23,16 +23,14 @@ private class ElementPrinter(printer: SmartPrinter) : AbstractElementPrinter<Ele
     override fun SmartPrinter.printAdditionalMethods(element: Element) {
         val kind = element.kind ?: error("Expected non-null element kind")
         with(element) {
-            // TODO: Add a kDoc for `accept`
+            printAcceptMethod(element, firVisitorType, hasImplementation = true, treeName = "FIR")
-            printAcceptMethod(element, firVisitorType, hasImplementation = true, kDoc = null)
 
-            // TODO: Add a kDoc for `transform`
             printTransformMethod(
                 element = element,
                 transformerClass = firTransformerType,
                 implementation = "transformer.transform${element.name}(this, data)",
                 returnType = TypeVariable("E", listOf(AbstractFirTreeBuilder.baseFirElement)),
-                kDoc = null,
+                treeName = "FIR",
             )
 
             fun Field.replaceDeclaration(override: Boolean, overridenType: TypeRefWithNullability? = null, forceNullable: Boolean = false) {
@@ -72,14 +70,20 @@ private class ElementPrinter(printer: SmartPrinter) : AbstractElementPrinter<Ele
                 println()
                 println("fun accept(visitor: ", firVisitorVoidType.render(), ") = accept(visitor, null)")
 
-                // TODO: Add a kDoc for `acceptChildren`
+                printAcceptChildrenMethod(
-                printAcceptChildrenMethod(element, firVisitorType, visitorResultType = TypeVariable("R"), kDoc = null)
+                    element = element,
+                    visitorClass = firVisitorType,
+                    visitorResultType = TypeVariable("R"),
+                )
                 println()
                 println()
                 println("fun acceptChildren(visitor: ", firVisitorVoidType.render(), ") = acceptChildren(visitor, null)")
 
-                // TODO: Add a kDoc for `transformChildren`
+                printTransformChildrenMethod(
-                printTransformChildrenMethod(element, firTransformerType, returnType = AbstractFirTreeBuilder.baseFirElement, kDoc = null)
+                    element = element,
+                    transformerClass = firTransformerType,
+                    returnType = AbstractFirTreeBuilder.baseFirElement,
+                )
                 println()
             }
         }

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

@@ -152,7 +152,7 @@ fun SmartPrinter.printImplementation(implementation: Implementation) {
             fun Field.acceptString(): String = "${name}${call()}accept(visitor, data)"
 
             if (hasAcceptChildrenMethod) {
-                printAcceptChildrenMethod(this, firVisitorType, TypeVariable("R"), override = true, kDoc = null)
+                printAcceptChildrenMethod(this, firVisitorType, TypeVariable("R"), override = true)
                 print(" {")
 
                 val walkableFields = walkableChildren
@@ -222,7 +222,6 @@ fun SmartPrinter.printImplementation(implementation: Implementation) {
                     this,
                     modality = Modality.ABSTRACT.takeIf { isAbstract },
                     override = true,
-                    kDoc = null,
                 )
                 if (!isInterface && !isAbstract) {
                     println(" {")

compiler/ir/ir.tree/tree-generator/src/org/jetbrains/kotlin/ir/generator/print/Elements.kt

@@ -36,36 +36,20 @@ private class ElementPrinter(printer: SmartPrinter) : AbstractElementPrinter<Ele
     override fun SmartPrinter.printAdditionalMethods(element: Element) {
         element.generationCallback?.invoke(this@ImportCollector, this)
 
-        if (element.hasAcceptMethod) {
+        printAcceptMethod(
-            printAcceptMethod(
+            element = element,
-                element = element,
+            visitorClass = elementVisitorType,
-                visitorClass = elementVisitorType,
+            hasImplementation = !element.isRootElement,
-                hasImplementation = !element.isRootElement,
+            treeName = "IR",
-                kDoc = """
+        )
-                Runs the provided [visitor] on the IR subtree with the root at this node.
-                
-                @param visitor The visitor to accept.
-                @param data An arbitrary context to pass to each invocation of [visitor]'s methods.
-                @return The value returned by the topmost `visit*` invocation.
-                """.trimIndent().takeIf { element.isRootElement }
-            )
-        }
 
-        if (element.hasTransformMethod) {
+        printTransformMethod(
-            printTransformMethod(
+            element = element,
-                element = element,
+            transformerClass = elementTransformerType,
-                transformerClass = elementTransformerType,
+            implementation = "accept(transformer, data)".takeIf { !element.isRootElement },
-                implementation = "accept(transformer, data)".takeIf { !element.isRootElement },
+            returnType = element,
-                returnType = element,
+            treeName = "IR",
-                kDoc = """
+        )
-                    Runs the provided [transformer] on the IR subtree with the root at this node.
-        
-                    @param transformer The transformer to use.
-                    @param data An arbitrary context to pass to each invocation of [transformer]'s methods.
-                    @return The transformed node.
-                    """.trimIndent().takeIf { element.isRootElement }
-            )
-        }
 
         if (element.hasAcceptChildrenMethod) {
             printAcceptChildrenMethod(
@@ -73,16 +57,6 @@ private class ElementPrinter(printer: SmartPrinter) : AbstractElementPrinter<Ele
                 visitorClass = elementVisitorType,
                 visitorResultType = StandardTypes.unit,
                 override = !element.isRootElement,
-                kDoc = """
-                    Runs the provided [visitor] on subtrees with roots in this node's children.
-                    
-                    Basically, calls `accept(visitor, data)` on each child of this node.
-                    
-                    Does **not** run [visitor] on this node itself.
-                    
-                    @param visitor The visitor for children to accept.
-                    @param data An arbitrary context to pass to each invocation of [visitor]'s methods.
-                    """.trimIndent().takeIf { element.isRootElement }
             )
 
             if (!element.isRootElement) {
@@ -116,16 +90,6 @@ private class ElementPrinter(printer: SmartPrinter) : AbstractElementPrinter<Ele
                 transformerClass = elementTransformerType,
                 returnType = StandardTypes.unit,
                 override = !element.isRootElement,
-                kDoc = """
-                    Recursively transforms this node's children *in place* using [transformer].
-                    
-                    Basically, executes `this.child = this.child.transform(transformer, data)` for each child of this node.
-                    
-                    Does **not** run [transformer] on this node itself.
-                    
-                    @param transformer The transformer to use for transforming the children.
-                    @param data An arbitrary context to pass to each invocation of [transformer]'s methods.
-                    """.trimIndent().takeIf { element.isRootElement }
             )
             if (!element.isRootElement) {
                 println(" {")

generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/printer/printUtils.kt

@@ -156,18 +156,26 @@ fun SmartPrinter.printAcceptMethod(
     element: AbstractElement<*, *>,
     visitorClass: ClassRef<PositionTypeParameterRef>,
     hasImplementation: Boolean,
-    kDoc: String?,
+    treeName: String,
 ) {
     if (!element.hasAcceptMethod) return
     println()
-    printKDoc(kDoc)
     val resultTP = TypeVariable("R")
+    val visitorParameter = FunctionParameter("visitor", visitorClass.withArgs(resultTP, dataTP))
+    if (element.isRootElement) {
+        printKDoc(
+            """
+            Runs the provided [${visitorParameter.name}] on the $treeName subtree with the root at this node.
+            
+            @param ${visitorParameter.name} The visitor to accept.
+            @param ${dataParameter.name} An arbitrary context to pass to each invocation of [${visitorParameter.name}]'s methods.
+            @return The value returned by the topmost `visit*` invocation.
+            """.trimIndent()
+        )
+    }
     printFunctionDeclaration(
         name = "accept",
-        parameters = listOf(
+        parameters = listOf(visitorParameter, dataParameter),
-            FunctionParameter("visitor", visitorClass.withArgs(resultTP, dataTP)),
-            dataParameter,
-        ),
         returnType = resultTP,
         typeParameters = listOf(resultTP, dataTP),
         override = !element.isRootElement,
@@ -175,7 +183,7 @@ fun SmartPrinter.printAcceptMethod(
     if (hasImplementation) {
         println(" =")
         withIndent {
-            print("visitor.", element.visitFunctionName, "(this, ", dataParameter.name, ")")
+            print(visitorParameter.name, ".", element.visitFunctionName, "(this, ", dataParameter.name, ")")
         }
     }
     println()
@@ -187,20 +195,28 @@ fun SmartPrinter.printTransformMethod(
     transformerClass: ClassRef<PositionTypeParameterRef>,
     implementation: String?,
     returnType: TypeRefWithNullability,
-    kDoc: String?,
+    treeName: String,
 ) {
     if (!element.hasTransformMethod) return
     println()
-    printKDoc(kDoc)
+    val transformerParameter = FunctionParameter("transformer", transformerClass.withArgs(dataTP))
+    if (element.isRootElement) {
+        printKDoc(
+            """
+            Runs the provided [${transformerParameter.name}] on the $treeName subtree with the root at this node.
+            
+            @param ${transformerParameter.name} The transformer to use.
+            @param ${dataParameter.name} An arbitrary context to pass to each invocation of [${transformerParameter.name}]'s methods.
+            @return The transformed node.
+            """.trimIndent()
+        )
+    }
     if (returnType is TypeParameterRef && implementation != null) {
         println("@Suppress(\"UNCHECKED_CAST\")")
     }
     printFunctionDeclaration(
         name = "transform",
-        parameters = listOf(
+        parameters = listOf(transformerParameter, dataParameter),
-            FunctionParameter("transformer", transformerClass.withArgs(dataTP)),
-            dataParameter,
-        ),
         returnType = returnType,
         typeParameters = listOfNotNull(returnType as? TypeVariable, dataTP),
         override = !element.isRootElement,
@@ -221,17 +237,27 @@ fun SmartPrinter.printAcceptChildrenMethod(
     visitorResultType: TypeRef,
     modality: Modality? = null,
     override: Boolean = false,
-    kDoc: String?,
 ) {
     if (!element.hasAcceptChildrenMethod) return
     println()
-    printKDoc(kDoc)
+    val visitorParameter = FunctionParameter("visitor", visitorClass.withArgs(visitorResultType, dataTP))
+    if (!override) {
+        printKDoc(
+            """
+            Runs the provided [${visitorParameter.name}] on subtrees with roots in this node's children.
+            
+            Basically, calls `accept(${visitorParameter.name}, ${dataParameter.name})` on each child of this node.
+            
+            Does **not** run [${visitorParameter.name}] on this node itself.
+            
+            @param ${visitorParameter.name} The visitor for children to accept.
+            @param ${dataParameter.name} An arbitrary context to pass to each invocation of [${visitorParameter.name}]'s methods.
+            """.trimIndent()
+        )
+    }
     printFunctionDeclaration(
         name = "acceptChildren",
-        parameters = listOf(
+        parameters = listOf(visitorParameter, dataParameter),
-            FunctionParameter("visitor", visitorClass.withArgs(visitorResultType, dataTP)),
-            dataParameter,
-        ),
         returnType = StandardTypes.unit,
         typeParameters = listOfNotNull(visitorResultType as? TypeVariable, dataTP),
         modality = modality,
@@ -246,17 +272,27 @@ fun SmartPrinter.printTransformChildrenMethod(
     returnType: TypeRef,
     modality: Modality? = null,
     override: Boolean = false,
-    kDoc: String?,
 ) {
     if (!element.hasTransformChildrenMethod) return
     println()
-    printKDoc(kDoc)
+    val transformerParameter = FunctionParameter("transformer", transformerClass.withArgs(dataTP))
+    if (!override) {
+        printKDoc(
+            """
+            Recursively transforms this node's children *in place* using [${transformerParameter.name}].
+            
+            Basically, executes `this.child = this.child.transform(${transformerParameter.name}, ${dataParameter.name})` for each child of this node.
+            
+            Does **not** run [${transformerParameter.name}] on this node itself.
+            
+            @param ${transformerParameter.name} The transformer to use for transforming the children.
+            @param ${dataParameter.name} An arbitrary context to pass to each invocation of [${transformerParameter.name}]'s methods.
+            """.trimIndent() + (if (returnType == StandardTypes.unit) "" else "\n@return `this`")
+        )
+    }
     printFunctionDeclaration(
         name = "transformChildren",
-        parameters = listOf(
+        parameters = listOf(transformerParameter, dataParameter),
-            FunctionParameter("transformer", transformerClass.withArgs(dataTP)),
-            dataParameter,
-        ),
         returnType = returnType,
         typeParameters = listOf(dataTP),
         modality = modality,