Unify Kotlin test suites

java/kotlin-extractor/src/main/kotlin/TrapWriter.kt

@@ -303,6 +303,20 @@ open class FileTrapWriter(
     /** The ID for the file that we are extracting from. */
     val fileId = mkFileId(filePath, populateFileTables)
 
+    /**
+     * When non-null, any location computed with a start/end offset pair exactly equal to
+     * [scopedOffsetRemap]`.first` is emitted as if it had the offset pair
+     * [scopedOffsetRemap]`.second`.
+     *
+     * This is used to normalise the source range of the synthesised body expressions of a
+     * delegated-property accessor: the K1 frontend anchors them at the enclosing
+     * `KtPropertyDelegate` (i.e. including the `by` keyword), whereas K2 anchors them at the
+     * delegate expression itself. The remap is set only while extracting such an accessor
+     * body and matches the full `by`-inclusive range exactly, so it never affects unrelated
+     * locations.
+     */
+    var scopedOffsetRemap: Pair<Pair<Int, Int>, Pair<Int, Int>>? = null
+
     override fun getDiagnosticTrapWriter(): DiagnosticTrapWriter {
         return dtw
     }
@@ -347,6 +361,15 @@ open class FileTrapWriter(
     fun getLocation(e: IrElement): Label<DbLocation> {
         return getLocation(getStartOffset(e), getEndOffset(e))
     }
+    /**
+     * Gets a label for the location starting at `e`'s (adjusted) start offset and ending at the
+     * given `endOffset`. This preserves the IrCall start-offset adjustment that [getLocation]
+     * applies, while allowing the caller to override the end offset (for example to widen an
+     * assignment's span to include its right-hand value).
+     */
+    fun getLocation(e: IrElement, endOffset: Int): Label<DbLocation> {
+        return getLocation(getStartOffset(e), endOffset)
+    }
     /**
      * Gets a label for the location corresponding to `startOffset` and `endOffset` within this
      * file.
@@ -398,6 +421,12 @@ class SourceFileTrapWriter(
     private val fileEntry = irFile.fileEntry
 
     override fun getLocation(startOffset: Int, endOffset: Int): Label<DbLocation> {
+        scopedOffsetRemap?.let { (from, to) ->
+            if (startOffset == from.first && endOffset == from.second) {
+                return getLocation(to.first, to.second)
+            }
+        }
+
         if (startOffset == UNDEFINED_OFFSET || endOffset == UNDEFINED_OFFSET) {
             if (startOffset != endOffset) {
                 loggerBase.warn(

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

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.kdoc.psi.api.KDoc
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.KtVisitor
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 
@@ -95,8 +96,21 @@ class CommentExtractorPSI(
                     }
                 }
 
-                // Only storing the owner of doc comments:
-                val ownerPsi = getKDocOwner(comment) ?: return
+                val ownerPsi = getKDocOwner(comment)
+                if (ownerPsi == null) {
+                    // A file-level KDoc (one written before the file's package
+                    // directive) has no declaration owner. The K2/FIR extractor
+                    // attributes such a comment to the compilation unit (the KDOC
+                    // is a direct child of the FILE node there), so we do the same
+                    // here to keep the two frontends in agreement.
+                    if (comment.parent is KtFile) {
+                        val fileOwnerLabel = getLabel(file)
+                        if (fileOwnerLabel != null) {
+                            tw.writeKtCommentOwners(commentLabel, fileOwnerLabel)
+                        }
+                    }
+                    return
+                }
 
                 val owners = mutableListOf<IrElement>()
                 file.accept(IrVisitorLookup(psi2Ir, ownerPsi, file), owners)
@@ -111,7 +125,7 @@ class CommentExtractorPSI(
 
             private fun getKDocOwner(comment: KDoc): PsiElement? {
                 val owner = comment.owner
-                if (owner == null) {
+                if (owner == null && comment.parent !is KtFile) {
                     logger.warn(
                         "Couldn't get owner of KDoc. The comment is extracted without an owner."
                     )

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

@@ -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,11 +35,52 @@ class CommentExtractorLighterAST(
         }
 
         val owners = findKDocOwners(file)
+        addMetadataLessKDocOwners(treeStructure, file, owners)
         extractComments(treeStructure.root, treeStructure, owners)
         return true
     }
 
-    private fun findKDocOwners(file: IrFile): Map<Int, List<IrElement>> {
+    // 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: Exception) {
+            // Never swallow IntelliJ's ProcessCanceledException: it is a control-flow
+            // exception that must propagate for cancellation to work. We match it by
+            // name rather than by type because the class is not resolvable on every
+            // supported compiler version (its embeddable classpath varies), and a
+            // compile-time reference breaks the cross-version build.
+            if (e.javaClass.name == "com.intellij.openapi.progress.ProcessCanceledException") {
+                throw e
+            }
+            logger.warn("Couldn't parse KDoc sections: ${e}")
+            null
+        }
+    }
+
+    private fun findKDocOwners(file: IrFile): MutableMap<Int, MutableList<IrElement>> {
         fun LighterASTNode.isKDocComment() = this.tokenType == KDocTokens.KDOC
 
         val kDocOwners = mutableMapOf<Int, MutableList<IrElement>>()
@@ -76,10 +119,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)) {
@@ -90,7 +198,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 -> {
@@ -117,8 +228,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/integration-tests/kotlin/all-platforms/enhanced-nullability/test.expected

@@ -2,7 +2,7 @@ exprs
 | Test.java:5:19:5:25 | Integer | Integer |
 | Test.java:5:38:5:44 | Integer | Integer |
 | Test.java:5:58:5:58 | p | Integer |
-| user.kt:2:7:2:7 | x | int |
+| user.kt:2:3:2:16 | x | int |
 | user.kt:2:11:2:11 | t | Test |
 | user.kt:2:11:2:16 | f(...) | Integer |
 | user.kt:2:13:2:16 | <implicit not null> | int |

java/ql/test-kotlin1/library-tests/annotation_classes/PrintAst.expected

@@ -44,13 +44,7 @@ def.kt:
 #   50|         1: [LocalVariableDeclStmt] var ...;
 #   50|           1: [LocalVariableDeclExpr] x
 #   50|             0: [IntegerLiteral] 10
-#   53|     3: [FieldDeclaration] int p;
-#-----|       -2: (Annotations)
-#   56|         1: [Annotation] Annot0k
-#    0|           1: [IntegerLiteral] 0
-#   53|       -1: [TypeAccess] int
-#   57|       0: [IntegerLiteral] 5
-#   57|     4: [Method] getP
+#   57|     3: [Method] getP
 #-----|       1: (Annotations)
 #   54|         1: [Annotation] Annot0k
 #    0|           1: [IntegerLiteral] 0
@@ -59,6 +53,12 @@ def.kt:
 #   57|         0: [ReturnStmt] return ...
 #   57|           0: [VarAccess] DefKt.p
 #   57|             -1: [TypeAccess] DefKt
+#   57|     4: [FieldDeclaration] int p;
+#-----|       -2: (Annotations)
+#   56|         1: [Annotation] Annot0k
+#    0|           1: [IntegerLiteral] 0
+#   57|       -1: [TypeAccess] int
+#   57|       0: [IntegerLiteral] 5
 #   57|     5: [Method] setP
 #-----|       1: (Annotations)
 #   55|         1: [Annotation] Annot0k
@@ -72,7 +72,7 @@ def.kt:
 #   57|           0: [AssignExpr] ...=...
 #   57|             0: [VarAccess] DefKt.p
 #   57|               -1: [TypeAccess] DefKt
-#   57|             1: [VarAccess] <set-?>
+#   53|             1: [VarAccess] <set-?>
 #   59|     6: [ExtensionMethod] myExtension
 #   59|       3: [TypeAccess] Unit
 #-----|       4: (Parameters)
@@ -175,8 +175,8 @@ def.kt:
 #    0|     3: [Method] valueOf
 #    0|       3: [TypeAccess] Y
 #-----|       4: (Parameters)
-#    0|         0: [Parameter] value
-#    0|           0: [TypeAccess] String
+#   34|         0: [Parameter] value
+#   34|           0: [TypeAccess] String
 #    0|     4: [Method] values
 #    0|       3: [TypeAccess] Y[]
 #    0|         0: [TypeAccess] Y
@@ -223,7 +223,7 @@ def.kt:
 #-----|       1: (Annotations)
 #   41|         1: [Annotation] Annot0k
 #    0|           1: [IntegerLiteral] 0
-#   41|       5: [BlockStmt] { ... }
+#   42|       5: [BlockStmt] { ... }
 #   42|         0: [SuperConstructorInvocationStmt] super(...)
 #   42|         1: [BlockStmt] { ... }
 use.java:

java/ql/test-kotlin1/library-tests/annotation_classes/classes.expected

@@ -32,9 +32,9 @@ annotations
 | def.kt:41:5:41:12 | Annot0k | def.kt:42:5:42:19 | Z | def.kt:5:1:21:60 | Annot0k |
 | def.kt:45:1:45:8 | Annot0k | def.kt:46:1:51:1 | fn | def.kt:5:1:21:60 | Annot0k |
 | def.kt:46:21:46:28 | Annot0k | def.kt:46:21:46:39 | a | def.kt:5:1:21:60 | Annot0k |
-| def.kt:54:1:54:12 | Annot0k | def.kt:57:1:57:23 | getP | def.kt:5:1:21:60 | Annot0k |
-| def.kt:55:1:55:12 | Annot0k | def.kt:57:1:57:23 | setP | def.kt:5:1:21:60 | Annot0k |
-| def.kt:56:1:56:14 | Annot0k | def.kt:53:1:57:23 | p | def.kt:5:1:21:60 | Annot0k |
+| def.kt:54:1:54:12 | Annot0k | def.kt:57:1:57:19 | getP | def.kt:5:1:21:60 | Annot0k |
+| def.kt:55:1:55:12 | Annot0k | def.kt:57:1:57:19 | setP | def.kt:5:1:21:60 | Annot0k |
+| def.kt:56:1:56:14 | Annot0k | def.kt:57:1:57:23 | p | def.kt:5:1:21:60 | Annot0k |
 | def.kt:59:5:59:21 | Annot0k | def.kt:59:5:59:28 | <this> | def.kt:5:1:21:60 | Annot0k |
 | use.java:10:5:10:21 | Annot0j | use.java:14:18:14:18 | Z | Annot0j.java:1:19:1:25 | Annot0j |
 | use.java:11:5:11:90 | Annot1j | use.java:14:18:14:18 | Z | Annot1j.java:1:19:1:25 | Annot1j |