Kotlin: extract KDoc sections under the K2/FIR frontend

The comment extractor has two implementations selected at runtime: the
PSI-based `CommentExtractorPSI` (used when the IR is built from PSI, i.e.
the K1 frontend) and `CommentExtractorLighterAST` (used under the K2/FIR
frontend, where no PSI is available and comments are recovered from the
FIR lighter AST). The PSI extractor writes a row per KDoc section
(`comment.getAllSections()`: the default section plus `@property`,
`@constructor`, ... tag sections), but the lighter-AST extractor did not,
because the KDOC node in the FIR lighter AST is a leaf: its section/tag
structure is never expanded there. As a result, `ktCommentSections`,
`ktCommentSectionNames` and `ktCommentSectionSubjectNames` were entirely
absent under K2, producing a large divergence from the K1 output.

KDoc section structure can only be recovered by parsing the KDoc text
into real PSI, which needs a `Project`. The compiler passes one to the
component registrar's `registerProjectComponents`, so we capture it there
into `KDocProjectHolder` (a process-global; a weak reference guarded by
`isDisposed` avoids keeping a disposed project alive). The lighter-AST
extractor then re-parses each KDoc's text with `KtPsiFactory` and reads
`getAllSections()`, writing exactly the same rows as the PSI extractor.
Because both paths delegate to the same compiler KDoc parser, the section
content, names and subject names are reproduced identically.

Trade-offs:

- A process-global project holder is used rather than threading the
  project through the extension/extractor constructors, which would touch
  a lot of unrelated wiring. Each CodeQL extraction is a single compiler
  invocation, so one project per process is a safe assumption; the weak
  reference and `isDisposed` check bound the lifetime risk.

- The re-parse forces the KDoc text into a doc-comment position via a
  throwaway trailing declaration. A KDoc is only recognised as a doc
  comment when it precedes a declaration; the appended declaration is
  inert and does not affect section parsing.

- Section output is byte-identical between K1 (2.3.20) and K2 (2.4.0)
  today because both use the compiler's own KDoc parser. This is a
  version-coupled compatibility shim rather than a guaranteed invariant;
  the shared test suite will catch any future drift.

Only the K2 expectations gain rows; the K1 output is unchanged. The
`test-kotlin2` comment section relations now match `test-kotlin1`
byte-for-byte, except for the line-1 KDoc whose source text still differs
between the two suites (addressed separately).

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

java/kotlin-extractor/src/main/kotlin/comments/KDocProjectHolder.kt

@@ -0,0 +1,24 @@
+package com.github.codeql.comments
+
+import com.intellij.openapi.project.Project
+import java.lang.ref.WeakReference
+
+// Holds the compiler's `Project`, captured in the component registrar's
+// `registerProjectComponents`. The K2/FIR comment extractor
+// (`CommentExtractorLighterAST`) needs it to build a `KtPsiFactory` and parse
+// KDoc section structure, which is not present in the FIR lighter AST (the KDOC
+// node is a leaf there). Under K1 the PSI-based extractor is used instead and
+// this holder is not consulted.
+//
+// A weak reference is used so this process-global does not keep a disposed
+// project alive; callers must tolerate a null/disposed project (sections are
+// then omitted, matching the pre-existing behaviour).
+object KDocProjectHolder {
+    private var ref: WeakReference<Project>? = null
+
+    var project: Project?
+        get() = ref?.get()?.takeUnless { it.isDisposed }
+        set(value) {
+            ref = value?.let { WeakReference(it) }
+        }
+}

java/kotlin-extractor/src/main/kotlin/utils/versions/v_1_8_0/Kotlin2ComponentRegistrar.kt

@@ -21,6 +21,7 @@ abstract class Kotlin2ComponentRegistrar : ComponentRegistrar {
         configuration: CompilerConfiguration
     ) {
         this.project = project
+        com.github.codeql.comments.KDocProjectHolder.project = project
         doRegisterExtensions(configuration)
     }
 

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

@@ -37,6 +37,40 @@ class CommentExtractorLighterAST(
         return true
     }
 
+    // KDoc section structure (the default section plus `@property`, `@constructor`,
+    // ... tag sections) is not represented in the FIR lighter AST: the KDOC node is
+    // a leaf there. The PSI-based extractor writes these sections, so to produce the
+    // same output under K2 we re-parse the KDoc text into PSI using the compiler's
+    // `Project` and read its sections. Returns null if no project was captured (in
+    // which case sections are omitted, matching the previous behaviour).
+    private val ktPsiFactory: org.jetbrains.kotlin.psi.KtPsiFactory? by lazy {
+        KDocProjectHolder.project?.let {
+            org.jetbrains.kotlin.psi.KtPsiFactory(it, markGenerated = false)
+        }
+    }
+
+    private fun parseKDocSections(
+        commentText: String
+    ): List<org.jetbrains.kotlin.kdoc.psi.impl.KDocSection>? {
+        val factory = ktPsiFactory ?: return null
+        return try {
+            // A KDoc is only recognised as a doc comment when it precedes a
+            // declaration, so we append a throwaway declaration before parsing.
+            val ktFile = factory.createFile("$commentText\nval __codeql_kdoc__ = 0")
+            val kdoc =
+                com.intellij.psi.util.PsiTreeUtil.findChildOfType(
+                    ktFile,
+                    org.jetbrains.kotlin.kdoc.psi.api.KDoc::class.java
+                )
+            kdoc?.getAllSections()
+        } catch (e: com.intellij.openapi.progress.ProcessCanceledException) {
+            throw e
+        } catch (e: Exception) {
+            logger.warn("Couldn't parse KDoc sections: ${e}")
+            null
+        }
+    }
+
     private fun findKDocOwners(file: IrFile): Map<Int, List<IrElement>> {
         fun LighterASTNode.isKDocComment() = this.tokenType == KDocTokens.KDOC
 
@@ -117,8 +151,20 @@ class CommentExtractorLighterAST(
             return
         }
 
-        // TODO: The PSI comment extractor extracts comment.getAllSections()
-        // here, so we should too
+        // Mirror the PSI extractor: write a row per KDoc section (default section
+        // and each `@tag` section) with its content, name, and subject name.
+        parseKDocSections(comment.toString())?.forEach { sec ->
+            val commentSectionLabel = tw.getFreshIdLabel<DbKtcommentsection>()
+            tw.writeKtCommentSections(commentSectionLabel, commentLabel, sec.getContent())
+            val name = sec.name
+            if (name != null) {
+                tw.writeKtCommentSectionNames(commentSectionLabel, name)
+            }
+            val subjectName = sec.getSubjectName()
+            if (subjectName != null) {
+                tw.writeKtCommentSectionSubjectNames(commentSectionLabel, subjectName)
+            }
+        }
 
         for (owner in owners.getOrDefault(comment.startOffset, listOf())) {
             val ownerLabel = getLabel(owner)

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

@@ -22,6 +22,7 @@ abstract class Kotlin2ComponentRegistrar : ComponentRegistrar {
         configuration: CompilerConfiguration
     ) {
         this.project = project
+        com.github.codeql.comments.KDocProjectHolder.project = project
         doRegisterExtensions(configuration)
     }
 

java/kotlin-extractor/src/main/kotlin/utils/versions/v_2_4_0/Kotlin2ComponentRegistrar.kt

@@ -24,6 +24,7 @@ abstract class Kotlin2ComponentRegistrar :
     ) {
         // Registration is done via ExtensionStorage in Kotlin 2.4+.
         // This legacy entry point remains for compatibility with service discovery.
+        com.github.codeql.comments.KDocProjectHolder.project = project
     }
 
     private var extensionStorage: CompilerPluginRegistrar.ExtensionStorage? = null

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

@@ -37,6 +37,41 @@ commentNoOwners
 | 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 |
+| comments.kt:4:1:11:3 | /**\n * A group of *members*.\n *\n * This class has no useful logic; it's just a documentation example.\n *\n * @property name the name of this group.\n * @constructor Creates an empty group.\n */ | A group of *members*.\n\nThis class has no useful logic; it's just a documentation example.\n\n |
+| comments.kt:4:1:11:3 | /**\n * A group of *members*.\n *\n * This class has no useful logic; it's just a documentation example.\n *\n * @property name the name of this group.\n * @constructor Creates an empty group.\n */ | Creates an empty group. |
+| comments.kt:4:1:11:3 | /**\n * A group of *members*.\n *\n * This class has no useful logic; it's just a documentation example.\n *\n * @property name the name of this group.\n * @constructor Creates an empty group.\n */ | the name of this group. |
+| comments.kt:14:5:16:7 | /**\n     * Members of this group.\n     */ | Members of this group. |
+| comments.kt:19:5:22:7 | /**\n     * Adds a [member] to this group.\n     * @return the new size of the group.\n     */ | Adds a [member] to this group.\n |
+| comments.kt:35:5:35:34 | /** Medium is in the middle */ | Medium is in the middle |
+| comments.kt:37:5:37:23 | /** This is high */ | This is high |
+| comments.kt:42:5:44:7 | /**\n     * A variable.\n     */ | A variable. |
+| comments.kt:48:1:50:3 | /**\n * A type alias comment\n */ | A type alias comment |
+| comments.kt:54:5:56:7 | /**\n     * An init block comment\n     */ | An init block comment |
+| comments.kt:61:5:63:7 | /**\n     * A prop comment\n     */ | A prop comment |
+| comments.kt:65:9:67:11 | /**\n         * An accessor comment\n         */ | An accessor comment |
+| comments.kt:71:9:73:11 | /**\n         * An anonymous function comment\n         */ | An anonymous function comment |
+| comments.kt:79:9:81:11 | /**\n         * A local function comment\n         */ | A local function comment |
+| comments.kt:88:10:90:11 | /**\n         * An anonymous object comment\n         */ | An anonymous object comment |
 commentSectionContents
+| A group of *members*.\n\nThis class has no useful logic; it's just a documentation example.\n\n | A group of *members*.\n\nThis class has no useful logic; it's just a documentation example.\n\n |
+| A local function comment | A local function comment |
+| A prop comment | A prop comment |
+| A type alias comment | A type alias comment |
+| A variable. | A variable. |
+| Adds a [member] to this group.\n | Adds a [member] to this group.\n |
+| An accessor comment | An accessor comment |
+| An anonymous function comment | An anonymous function comment |
+| An anonymous object comment | An anonymous object comment |
+| An init block comment | An init block comment |
+| Creates an empty group. | Creates an empty group. |
+| Kdoc owned by CompilationUnit | Kdoc owned by CompilationUnit |
+| Medium is in the middle | Medium is in the middle |
+| Members of this group. | Members of this group. |
+| This is high | This is high |
+| the name of this group. | the name of this group. |
 commentSectionNames
+| Creates an empty group. | constructor |
+| the name of this group. | property |
 commentSectionSubjectNames
+| the name of this group. | name |