Kotlin: Populate numfiles

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

@@ -84,7 +84,7 @@ class ExternalDeclExtractor(val logger: FileLogger, val invocationTrapFile: Stri
                                     // file information if needed:
                                     val ftw = tw.makeFileTrapWriter(binaryPath, irDecl is IrClass)
 
-                                    val fileExtractor = KotlinFileExtractor(logger, ftw, binaryPath, manager, this, primitiveTypeMapping, pluginContext, KotlinFileExtractor.DeclarationStack(), globalExtensionState)
+                                    val fileExtractor = KotlinFileExtractor(logger, ftw, null, binaryPath, manager, this, primitiveTypeMapping, pluginContext, KotlinFileExtractor.DeclarationStack(), globalExtensionState)
 
                                     if (irDecl is IrClass) {
                                         // Populate a location and compilation-unit package for the file. This is similar to

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

@@ -322,7 +322,8 @@ private fun doFile(
                 // file information
                 val sftw = tw.makeSourceFileTrapWriter(srcFile, true)
                 val externalDeclExtractor = ExternalDeclExtractor(logger, invocationTrapFile, srcFilePath, primitiveTypeMapping, pluginContext, globalExtensionState, fileTrapWriter)
-                val fileExtractor = KotlinFileExtractor(logger, sftw, srcFilePath, null, externalDeclExtractor, primitiveTypeMapping, pluginContext, KotlinFileExtractor.DeclarationStack(), globalExtensionState)
+                val linesOfCode = LinesOfCode(logger, sftw, srcFile)
+                val fileExtractor = KotlinFileExtractor(logger, sftw, linesOfCode, srcFilePath, null, externalDeclExtractor, primitiveTypeMapping, pluginContext, KotlinFileExtractor.DeclarationStack(), globalExtensionState)
 
                 fileExtractor.extractFileContents(srcFile, sftw.fileId)
                 externalDeclExtractor.extractExternalClasses()

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

@@ -35,6 +35,7 @@ import kotlin.collections.ArrayList
 open class KotlinFileExtractor(
     override val logger: FileLogger,
     override val tw: FileTrapWriter,
+    val linesOfCode: LinesOfCode?,
     val filePath: String,
     dependencyCollector: OdasaOutput.TrapFileManager?,
     externalClassExtractor: ExternalDeclExtractor,
@@ -90,6 +91,8 @@ open class KotlinFileExtractor(
             if (!declarationStack.isEmpty()) {
                 logger.errorElement("Declaration stack is not empty after processing the file", file)
             }
+
+            linesOfCode?.linesOfCodeInFile(id)
         }
     }
 
@@ -459,6 +462,8 @@ open class KotlinFileExtractor(
                 extractClassModifiers(c, id)
                 extractClassSupertypes(c, id, inReceiverContext = true) // inReceiverContext = true is specified to force extraction of member prototypes of base types
 
+                linesOfCode?.linesOfCodeInDeclaration(c, id)
+
                 return id
             }
         }
@@ -1038,6 +1043,8 @@ open class KotlinFileExtractor(
                     addModifiers(id, "suspend")
                 }
 
+                linesOfCode?.linesOfCodeInDeclaration(f, id)
+
                 return id
             }
         }

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

@@ -138,13 +138,13 @@ open class KotlinUsesExtractor(
             val newTrapWriter = tw.makeFileTrapWriter(filePath, true)
             val newLoggerTrapWriter = logger.tw.makeFileTrapWriter(filePath, false)
             val newLogger = FileLogger(logger.loggerBase, newLoggerTrapWriter)
-            return KotlinFileExtractor(newLogger, newTrapWriter, filePath, dependencyCollector, externalClassExtractor, primitiveTypeMapping, pluginContext, newDeclarationStack, globalExtensionState)
+            return KotlinFileExtractor(newLogger, newTrapWriter, null, filePath, dependencyCollector, externalClassExtractor, primitiveTypeMapping, pluginContext, newDeclarationStack, globalExtensionState)
         }
 
         val newTrapWriter = tw.makeSourceFileTrapWriter(clsFile, true)
         val newLoggerTrapWriter = logger.tw.makeSourceFileTrapWriter(clsFile, false)
         val newLogger = FileLogger(logger.loggerBase, newLoggerTrapWriter)
-        return KotlinFileExtractor(newLogger, newTrapWriter, clsFile.path, dependencyCollector, externalClassExtractor, primitiveTypeMapping, pluginContext, newDeclarationStack, globalExtensionState)
+        return KotlinFileExtractor(newLogger, newTrapWriter, null, clsFile.path, dependencyCollector, externalClassExtractor, primitiveTypeMapping, pluginContext, newDeclarationStack, globalExtensionState)
     }
 
     // The Kotlin compiler internal representation of Outer<T>.Inner<S>.InnerInner<R> is InnerInner<R, S, T>. This function returns just `R`.

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

@@ -0,0 +1,131 @@
+package com.github.codeql
+
+import com.github.codeql.utils.versions.Psi2Ir
+import com.intellij.psi.PsiComment
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiWhiteSpace
+import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.kdoc.psi.api.KDocElement
+import org.jetbrains.kotlin.psi.KtCodeFragment
+import org.jetbrains.kotlin.psi.KtVisitor
+
+class LinesOfCode(
+    val logger: FileLogger,
+    val tw: FileTrapWriter,
+    val file: IrFile
+) {
+    val psi2Ir = Psi2Ir(logger)
+
+    fun linesOfCodeInFile(id: Label<DbFile>) {
+        val ktFile = psi2Ir.getKtFile(file)
+        if (ktFile == null) {
+            logger.warnElement("Cannot find PSI for file", file)
+            println("No KtFile")
+            return
+        }
+        linesOfCodeInPsi(id, ktFile, file)
+    }
+
+    fun linesOfCodeInDeclaration(d: IrDeclaration, id: Label<out DbSourceline>) {
+        val p = psi2Ir.findPsiElement(d, file)
+        if (p == null) {
+            logger.warnElement("Cannot find PSI for declaration: " + d.javaClass, d)
+            println("No p")
+            return
+        }
+        linesOfCodeInPsi(id, p, d)
+    }
+
+    private fun linesOfCodeInPsi(id: Label<out DbSourceline>, root: PsiElement, e: IrElement) {
+        val document = root.getContainingFile().getViewProvider().getDocument()
+        if (document == null) {
+            logger.errorElement("Cannot find document for PSI", e)
+            tw.writeNumlines(id, 0, 0, 0)
+            return
+        }
+
+        val rootRange = root.getTextRange()
+        val rootFirstLine = document.getLineNumber(rootRange.getStartOffset())
+        val rootLastLine = document.getLineNumber(rootRange.getEndOffset())
+        if (rootLastLine < rootFirstLine) {
+            logger.errorElement("PSI ends before it starts", e)
+            tw.writeNumlines(id, 0, 0, 0)
+            return
+        }
+        val numLines = 1 + rootLastLine - rootFirstLine
+        val lineContents = Array(numLines) { LineContent() }
+
+        val visitor =
+            object : KtVisitor<Unit, Unit>() {
+                override fun visitElement(element: PsiElement) {
+                    val isComment = element is PsiComment
+                    // Comments may include nodes that aren't PsiComments,
+                    // so we don't want to visit them or we'll think they
+                    // are code.
+                    if (!isComment) {
+                        element.acceptChildren(this)
+                    }
+
+                    if (element is PsiWhiteSpace) {
+                        return
+                    }
+                    // Leaf nodes are assumed to be tokens, and
+                    // therefore we count any lines that they are on.
+                    // For comments, we actually need to look at the
+                    // outermost node, as the leaves of KDocs don't
+                    // necessarily cover all lines.
+                    if (isComment || element.getChildren().size == 0) {
+                        val range = element.getTextRange()
+                        val startOffset = range.getStartOffset()
+                        val endOffset = range.getEndOffset()
+                        // The PSI doesn't seem to have anything like
+                        // the IR's UNDEFINED_OFFSET and SYNTHETIC_OFFSET,
+                        // but < 0 still seem to represent bad/unknown
+                        // locations.
+                        if (startOffset < 0 || endOffset < 0) {
+                            logger.errorElement("PSI has negative offset", e)
+                            return
+                        }
+                        if (startOffset > endOffset) {
+                            return
+                        }
+                        // We might get e.g. an import list for a file
+                        // with no imports, which claims to have start
+                        // and end offsets of 0. Anything of 0 width
+                        // we therefore just skip.
+                        if (startOffset == endOffset) {
+                            return
+                        }
+                        val firstLine = document.getLineNumber(startOffset)
+                        val lastLine = document.getLineNumber(endOffset)
+                        if (firstLine < rootFirstLine) {
+                            logger.errorElement("PSI element starts before root", e)
+                            return
+                        } else if (lastLine > rootLastLine) {
+                            logger.errorElement("PSI element ends after root", e)
+                            return
+                        }
+                        for (line in firstLine..lastLine) {
+                            val lineContent = lineContents[line - rootFirstLine]
+                            if (isComment) {
+                                lineContent.containsComment = true
+                            } else {
+                                lineContent.containsCode = true
+                            }
+                        }
+                    }
+                }
+            }
+        root.accept(visitor)
+        val total = lineContents.size
+        val code = lineContents.count { it.containsCode }
+        val comment = lineContents.count { it.containsComment }
+        tw.writeNumlines(id, total, code, comment)
+    }
+
+    private class LineContent {
+        var containsComment = false
+        var containsCode = false
+    }
+}