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,29 +93,10 @@ class CommentExtractor(private val fileExtractor: KotlinFileExtractor, private v
                 file.accept(IrVisitorLookup(psi2Ir, ownerPsi, file), owners)
 
                 for (ownerIr in owners) {
-                    val ownerLabel =
-                        if (ownerIr == file)
-                            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)
+                    val ownerLabel = getLabel(ownerIr)
+                    if (ownerLabel != null) {
+                        tw.writeKtCommentOwners(commentLabel, ownerLabel)
+                    }
                 }
             }
 
@@ -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
+}