Kotlin: attribute enum-entry and init-block KDoc owners under K2/FIR

Under the K2/FIR frontend the lighter-AST comment extractor finds a KDoc's
owner by walking the IR and, for each element that carries FIR metadata,
checking whether that metadata's source node has a KDOC child. Enum
entries (`IrEnumEntry`) and anonymous initializers
(`IrAnonymousInitializer`) carry no FIR metadata, so the search could
never reach them: their KDoc comments were recorded with no owner, whereas
the PSI frontend attributes them (an enum entry to itself, an init block
to its enclosing class).

These IR declarations do retain valid source offsets, however. This adds a
supplementary pass that walks the file's lighter AST, finds the
`ENUM_ENTRY` and `CLASS_INITIALIZER` nodes that carry a KDOC child, and
matches each to the metadata-less IR declaration whose source offset lies
within the node's range. Matching is fail-closed: a node is only
attributed when exactly one candidate declaration falls inside it, so
ambiguous or unexpected shapes leave the comment ownerless rather than
guessing.

Offset containment is used rather than per-kind heuristics (matching enum
entries by name, initializers by order) because it is uniform across both
node kinds and does not depend on names being unique or declaration order
being stable.

Only the K2 expectations change: the `test-kotlin2` `comments` enum-entry
and init-block KDoc owners now match `test-kotlin1` exactly. The one
remaining owner difference in that test is a KDoc on an anonymous function
passed as a call argument, which the PSI frontend attributes to the
enclosing property via its owner walk-up and the FIR frontend does not;
that discrepancy is left as-is (see the updated TODO).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: andersfugmann

java/kotlin-extractor/src/main/kotlin/utils/versions/v_1_9_0-Beta/CommentExtractorLighterAST.kt

@@ -13,10 +13,12 @@ import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 import org.jetbrains.kotlin.ir.visitors.acceptVoid
 import org.jetbrains.kotlin.kdoc.lexer.KDocTokens
 import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.util.getChildren
 
-// TODO: This doesn't give owners to as many comments as the PSI extractor does.
-// See the library-tests/comments tests for details.
+// TODO: This still doesn't give owners to quite as many comments as the PSI
+// extractor does (e.g. a KDoc on an anonymous function passed as a call
+// argument). See the library-tests/comments tests for details.
 
 class CommentExtractorLighterAST(
     fileExtractor: KotlinFileExtractor,
@@ -33,6 +35,7 @@ class CommentExtractorLighterAST(
         }
 
         val owners = findKDocOwners(file)
+        addMetadataLessKDocOwners(treeStructure, file, owners)
         extractComments(treeStructure.root, treeStructure, owners)
         return true
     }
@@ -71,7 +74,7 @@ class CommentExtractorLighterAST(
         }
     }
 
-    private fun findKDocOwners(file: IrFile): Map<Int, List<IrElement>> {
+    private fun findKDocOwners(file: IrFile): MutableMap<Int, MutableList<IrElement>> {
         fun LighterASTNode.isKDocComment() = this.tokenType == KDocTokens.KDOC
 
         val kDocOwners = mutableMapOf<Int, MutableList<IrElement>>()
@@ -110,10 +113,75 @@ class CommentExtractorLighterAST(
         return kDocOwners
     }
 
+    // Enum entries and anonymous initializers carry no FIR metadata under the
+    // K2/FIR frontend, so the metadata-based search in `findKDocOwners` cannot
+    // reach them and their KDoc comments would be left without an owner (the PSI
+    // frontend does attribute them). They do, however, retain valid IR source
+    // offsets. We recover their ownership by finding the ENUM_ENTRY /
+    // CLASS_INITIALIZER lighter-AST nodes that carry a KDOC child and matching
+    // each to the metadata-less IR declaration whose source offset falls within
+    // the node's range. Matching is fail-closed: a node is only attributed when
+    // exactly one candidate declaration lies within it.
+    private fun addMetadataLessKDocOwners(
+        treeStructure: FlyweightCapableTreeStructure<LighterASTNode>,
+        file: IrFile,
+        owners: MutableMap<Int, MutableList<IrElement>>
+    ) {
+        val candidates = mutableListOf<IrDeclaration>()
+        file.acceptVoid(
+            object : IrVisitorVoid() {
+                override fun visitElement(element: IrElement) {
+                    if (element is IrEnumEntry || element is IrAnonymousInitializer) {
+                        val decl = element as IrDeclaration
+                        if (
+                            decl.startOffset != UNDEFINED_OFFSET &&
+                                decl.startOffset != SYNTHETIC_OFFSET
+                        ) {
+                            candidates.add(decl)
+                        }
+                    }
+                    element.acceptChildrenVoid(this)
+                }
+            }
+        )
+        if (candidates.isEmpty()) {
+            return
+        }
+
+        fun visit(node: LighterASTNode) {
+            if (
+                node.tokenType == KtNodeTypes.ENUM_ENTRY ||
+                    node.tokenType == KtNodeTypes.CLASS_INITIALIZER
+            ) {
+                node.getChildren(treeStructure)
+                    .firstOrNull { it.tokenType == KDocTokens.KDOC }
+                    ?.let { kDoc ->
+                        val startOffset = kDoc.startOffset
+                        if (
+                            startOffset != UNDEFINED_OFFSET && startOffset != SYNTHETIC_OFFSET
+                        ) {
+                            val matches =
+                                candidates.filter {
+                                    it.startOffset >= node.startOffset &&
+                                        it.startOffset < node.endOffset
+                                }
+                            if (matches.size == 1) {
+                                owners
+                                    .getOrPut(startOffset, { mutableListOf<IrElement>() })
+                                    .add(matches[0])
+                            }
+                        }
+                    }
+            }
+            node.getChildren(treeStructure).forEach { visit(it) }
+        }
+        visit(treeStructure.root)
+    }
+
     private fun extractComments(
         node: LighterASTNode,
         treeStructure: FlyweightCapableTreeStructure<LighterASTNode>,
-        owners: Map<Int, List<IrElement>>
+        owners: MutableMap<Int, MutableList<IrElement>>
     ) {
         node.getChildren(treeStructure).forEach {
             if (KtTokens.COMMENTS.contains(it.tokenType)) {
@@ -124,7 +192,10 @@ class CommentExtractorLighterAST(
         }
     }
 
-    private fun extractComment(comment: LighterASTNode, owners: Map<Int, List<IrElement>>) {
+    private fun extractComment(
+        comment: LighterASTNode,
+        owners: MutableMap<Int, MutableList<IrElement>>
+    ) {
         val type: CommentType =
             when (comment.tokenType) {
                 KtTokens.EOL_COMMENT -> {

java/ql/test-kotlin2/library-tests/comments/comments.expected

@@ -22,7 +22,10 @@ commentOwners
 | comments.kt:14:5:16:7 | /**\n     * Members of this group.\n     */ | comments.kt:17:13:17:46 | members |
 | comments.kt:14:5:16:7 | /**\n     * Members of this group.\n     */ | comments.kt:17:13:17:46 | members |
 | comments.kt:19:5:22:7 | /**\n     * Adds a [member] to this group.\n     * @return the new size of the group.\n     */ | comments.kt:23:5:26:5 | add |
+| comments.kt:35:5:35:34 | /** Medium is in the middle */ | comments.kt:36:5:36:14 | Medium |
+| comments.kt:37:5:37:23 | /** This is high */ | comments.kt:38:5:38:11 | High |
 | comments.kt:48:1:50:3 | /**\n * A type alias comment\n */ | comments.kt:51:1:51:24 | MyType |
+| comments.kt:54:5:56:7 | /**\n     * An init block comment\n     */ | comments.kt:53:1:58:1 | InitBlock |
 | comments.kt:61:5:63:7 | /**\n     * A prop comment\n     */ | comments.kt:64:5:68:17 | prop |
 | comments.kt:65:9:67:11 | /**\n         * An accessor comment\n         */ | comments.kt:68:9:68:17 | getProp |
 | comments.kt:79:9:81:11 | /**\n         * A local function comment\n         */ | comments.kt:82:9:82:24 | localFn |
@@ -31,10 +34,7 @@ commentOwners
 commentNoOwners
 | comments.kt:24:9:24:25 | // A line comment |
 | comments.kt:28:5:30:6 | /*\n    A block comment\n    */ |
-| comments.kt:35:5:35:34 | /** Medium is in the middle */ |
-| comments.kt:37:5:37:23 | /** This is high */ |
 | comments.kt:42:5:44:7 | /**\n     * A variable.\n     */ |
-| comments.kt:54:5:56:7 | /**\n     * An init block comment\n     */ |
 | comments.kt:71:9:73:11 | /**\n         * An anonymous function comment\n         */ |
 commentSections
 | comments.kt:1:1:1:36 | /** Kdoc owned by CompilationUnit */ | Kdoc owned by CompilationUnit |