Merge pull request #10268 from tamasvajk/kotlin-local-function-comments

Kotlin: fix doc comment extraction for local functions

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

@@ -1140,10 +1140,6 @@ open class KotlinUsesExtractor(
         // Note not using `parentsWithSelf` as that only works if `d` is an IrDeclarationParent
         d.parents.any { (it as? IrAnnotationContainer)?.hasAnnotation(jvmWildcardSuppressionAnnotaton) == true }
 
-    protected fun IrFunction.isLocalFunction(): Boolean {
-        return this.visibility == DescriptorVisibilities.LOCAL
-    }
-
     /**
      * Class to hold labels for generated classes around local functions, lambdas, function references, and property references.
      */
@@ -1185,6 +1181,14 @@ open class KotlinUsesExtractor(
         return res
     }
 
+    fun getExistingLocallyVisibleFunctionLabel(f: IrFunction): Label<DbMethod>? {
+        if (!f.isLocalFunction()){
+            return null
+        }
+
+        return tw.lm.locallyVisibleFunctionLabelMapping[f]?.function
+    }
+
     // These are classes with Java equivalents, but whose methods don't all exist on those Java equivalents--
     // for example, the numeric classes define arithmetic functions (Int.plus, Long.or and so on) that lower to
     // primitive arithmetic on the JVM, but which we extract as calls to reflect the source syntax more closely.

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

@@ -2,6 +2,7 @@ package com.github.codeql.comments
 
 import com.github.codeql.*
 import com.github.codeql.utils.IrVisitorLookup
+import com.github.codeql.utils.isLocalFunction
 import com.github.codeql.utils.versions.Psi2Ir
 import com.intellij.psi.PsiComment
 import com.intellij.psi.PsiElement
@@ -92,31 +93,12 @@ class CommentExtractor(private val fileExtractor: KotlinFileExtractor, private v
                 file.accept(IrVisitorLookup(psi2Ir, ownerPsi, file), owners)
 
                 for (ownerIr in owners) {
-                    val ownerLabel =
+                    val ownerLabel = getLabel(ownerIr)
-                        if (ownerIr == file)
+                    if (ownerLabel != null) {
-                            fileLabel
-                        else {
-                            if (ownerIr is IrValueParameter && ownerIr.index == -1) {
-                                // Don't attribute comments to the implicit `this` parameter of a function.
-                                continue
-                            }
-                            val label: String
-                            val existingLabel = if (ownerIr is IrVariable) {
-                                label = "variable ${ownerIr.name.asString()}"
-                                tw.getExistingVariableLabelFor(ownerIr)
-                            } else {
-                                label = getLabel(ownerIr) ?: continue
-                                tw.getExistingLabelFor<DbTop>(label)
-                            }
-                            if (existingLabel == null) {
-                                logger.warn("Couldn't get existing label for $label")
-                                continue
-                            }
-                            existingLabel
-                        }
                         tw.writeKtCommentOwners(commentLabel, ownerLabel)
                     }
                 }
+            }
 
             private fun getKDocOwner(comment: KDoc) : PsiElement? {
                 val owner = comment.owner
@@ -126,11 +108,47 @@ class CommentExtractor(private val fileExtractor: KotlinFileExtractor, private v
                 return owner
             }
 
-            private fun getLabel(element: IrElement) : String? {
+            private fun getLabel(element: IrElement): Label<out DbTop>? {
+                if (element == file)
+                    return fileLabel
+
+                if (element is IrValueParameter && element.index == -1) {
+                    // Don't attribute comments to the implicit `this` parameter of a function.
+                    return null
+                }
+
+                val label: String
+                val existingLabel = if (element is IrVariable) {
+                    // local variables are not named globally, so we need to get them from the variable label cache
+                    label = "variable ${element.name.asString()}"
+                    tw.getExistingVariableLabelFor(element)
+                } else if (element is IrFunction && element.isLocalFunction()) {
+                    // local functions are not named globally, so we need to get them from the local function label cache
+                    label = "local function ${element.name.asString()}"
+                    fileExtractor.getExistingLocallyVisibleFunctionLabel(element)
+                }
+                else {
+                    label = getLabelForNamedElement(element) ?: return null
+                    tw.getExistingLabelFor<DbTop>(label)
+                }
+                if (existingLabel == null) {
+                    logger.warn("Couldn't get existing label for $label")
+                    return null
+                }
+                return existingLabel
+            }
+
+            private fun getLabelForNamedElement(element: IrElement) : String? {
                 when (element) {
                     is IrClass -> return fileExtractor.getClassLabel(element, listOf()).classLabel
                     is IrTypeParameter -> return fileExtractor.getTypeParameterLabel(element)
-                    is IrFunction -> return fileExtractor.getFunctionLabel(element, null)
+                    is IrFunction -> {
+                        return if (element.isLocalFunction()) {
+                            null
+                        } else {
+                            fileExtractor.getFunctionLabel(element, null)
+                        }
+                    }
                     is IrValueParameter -> return fileExtractor.getValueParameterLabel(element, null)
                     is IrProperty -> return fileExtractor.getPropertyLabel(element)
                     is IrField -> return fileExtractor.getFieldLabel(element)
@@ -144,10 +162,10 @@ class CommentExtractor(private val fileExtractor: KotlinFileExtractor, private v
                             return null
                         }
                         // Assign the comment to the class. The content of the `init` blocks might be extracted in multiple constructors.
-                        return getLabel(parentClass)
+                        return getLabelForNamedElement(parentClass)
                     }
 
-                    // Fresh entities:
+                    // Fresh entities, not named elements:
                     is IrBody -> return null
                     is IrExpression -> return null
 

java/kotlin-extractor/src/main/kotlin/utils/Helpers.kt

@@ -0,0 +1,8 @@
+package com.github.codeql.utils
+
+import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+
+fun IrFunction.isLocalFunction(): Boolean {
+    return this.visibility == DescriptorVisibilities.LOCAL
+}

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

@@ -7,9 +7,13 @@ comments
 | comments.kt:28:5:30:6 | /*\n    A block comment\n    */ | /*\n    A block comment\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:6 | /**\n     * A variable.\n    */ | /**\n     * A variable.\n    */ |
+| comments.kt:42:5:44:7 | /**\n     * A variable.\n     */ | /**\n     * A variable.\n     */ |
 | comments.kt:48:1:50:3 | /**\n * A type alias comment\n */ | /**\n * A type alias comment\n */ |
 | comments.kt:54:5:56:7 | /**\n     * An init block comment\n     */ | /**\n     * An init block comment\n     */ |
+| comments.kt:61:5:63:7 | /**\n     * A prop comment\n     */ | /**\n     * A prop comment\n     */ |
+| comments.kt:65:9:67:11 | /**\n         * An accessor comment\n         */ | /**\n         * An accessor comment\n         */ |
+| comments.kt:71:9:73:11 | /**\n         * An anonymous function comment\n         */ | /**\n         * An anonymous function comment\n         */ |
+| comments.kt:79:9:81:11 | /**\n         * A local function comment\n         */ | /**\n         * A local function comment\n         */ |
 commentOwners
 | 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 */ | comments.kt:12:1:31:1 | 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 */ | comments.kt:12:1:31:1 | Group |
@@ -19,9 +23,19 @@ commentOwners
 | 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:42:5:44:6 | /**\n     * A variable.\n    */ | comments.kt:45:5:45:13 | int a |
+| comments.kt:42:5:44:7 | /**\n     * A variable.\n     */ | comments.kt:45:5:45:13 | int a |
 | 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:71:9:73:11 | /**\n         * An anonymous function comment\n         */ | comments.kt:70:5:76:10 | getL |
+| comments.kt:71:9:73:11 | /**\n         * An anonymous function comment\n         */ | comments.kt:70:5:76:10 | l |
+| comments.kt:71:9:73:11 | /**\n         * An anonymous function comment\n         */ | comments.kt:70:5:76:10 | l |
+| comments.kt:79:9:81:11 | /**\n         * A local function comment\n         */ | comments.kt:82:9:82:24 | localFn |
+commentNoOwners
+| comments.kt:1:1:1:25 | /** Kdoc with no owner */ |
+| comments.kt:24:9:24:25 | // A line comment |
+| comments.kt:28:5:30:6 | /*\n    A block comment\n    */ |
 commentSections
 | comments.kt:1:1:1:25 | /** Kdoc with no owner */ | Kdoc with no owner |
 | 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 |
@@ -31,14 +45,22 @@ commentSections
 | 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:6 | /**\n     * A variable.\n    */ | A variable. |
+| 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 |
 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 init block comment | An init block comment |
 | Creates an empty group. | Creates an empty group. |
 | Kdoc with no owner | Kdoc with no owner |

java/ql/test/kotlin/library-tests/comments/comments.kt

@@ -56,3 +56,29 @@ class InitBlock {
      */
     init { }
 }
+
+class X {
+    /**
+     * A prop comment
+     */
+    val prop: Int
+        /**
+         * An accessor comment
+         */
+        get() = 5
+
+    val l: Lazy<Int> = lazy(
+        /**
+         * An anonymous function comment
+         */
+        fun(): Int {
+            return 5
+        })
+
+    fun fn() {
+        /**
+         * A local function comment
+         */
+        fun localFn() {}
+    }
+}

java/ql/test/kotlin/library-tests/comments/comments.ql

@@ -4,6 +4,8 @@ query predicate comments(KtComment c, string s) { c.getText() = s }
 
 query predicate commentOwners(KtComment c, Top t) { c.getOwner() = t }
 
+query predicate commentNoOwners(KtComment c) { not exists(c.getOwner()) }
+
 query predicate commentSections(KtComment c, KtCommentSection s) { c.getSections() = s }
 
 query predicate commentSectionContents(KtCommentSection s, string c) { s.getContent() = c }