From e29d9ddacb42a98efcdaa85aa04c866c5d742153 Mon Sep 17 00:00:00 2001
From: Chris Smowton <chris@smowton.net>
Date: Wed, 27 Nov 2024 16:49:37 +0000
Subject: [PATCH] Restore location and name reporting for symbols

---
 .../src/main/kotlin/KotlinFileExtractor.kt    | 19 +++++++------
 .../src/main/kotlin/entities/Class.kt         |  2 +-
 .../src/main/kotlin/entities/Function.kt      |  2 +-
 .../src/main/kotlin/utils/Logger.kt           | 28 ++++++++++++++++++-
 4 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/java/kotlin-extractor2/src/main/kotlin/KotlinFileExtractor.kt b/java/kotlin-extractor2/src/main/kotlin/KotlinFileExtractor.kt
index 7d949abc971..bc6346a3b9c 100644
--- a/java/kotlin-extractor2/src/main/kotlin/KotlinFileExtractor.kt
+++ b/java/kotlin-extractor2/src/main/kotlin/KotlinFileExtractor.kt
@@ -117,14 +117,15 @@ open class KotlinFileExtractor(
         val metaAnnotationSupport = MetaAnnotationSupport(logger, pluginContext, this)
     */
 
-    inline fun <T> with(kind: String, element: KtElement?, f: () -> T): T {
-        val name =
-            when (element) {
-                is KtFile -> element.virtualFilePath
-                is KtNamed -> element.getNameAsName()?.asString() ?: "<missing name>"
-                else -> "<no name>"
-            }
-        val loc = element?.let { tw.getLocationString(it) } ?: "<unknown location>"
+    inline fun <T> with(kind: String, element: PsiElement, f: () -> T) = with(kind, PsiElementWrapper(element), f)
+    inline fun <T> with(kind: String, element: KaSymbol, f: () -> T) =
+        with(kind,
+            element.psiSafe<PsiElement>()?.let { PsiElementWrapper(it) } ?: SymbolWrapper(element),
+            f)
+
+    inline fun <T> with(kind: String, element: PsiElementOrSymbol, f: () -> T): T {
+        val name = element.getName()
+        val loc = element.getLocationString(tw)
         val context = logger.loggerState.extractorContextStack
         context.push(ExtractorContext(kind, element, name, loc))
         try {
@@ -299,7 +300,7 @@ open class KotlinFileExtractor(
                 extractAnnotations: Boolean
         */
     ) {
-        with("declaration", declaration.psiSafe()) {
+        with("declaration", declaration) {
 /*
 OLD: KE1
             if (!shouldExtractDecl(declaration, extractPrivateMembers)) return
diff --git a/java/kotlin-extractor2/src/main/kotlin/entities/Class.kt b/java/kotlin-extractor2/src/main/kotlin/entities/Class.kt
index bb00156677d..9c98f4cba1b 100644
--- a/java/kotlin-extractor2/src/main/kotlin/entities/Class.kt
+++ b/java/kotlin-extractor2/src/main/kotlin/entities/Class.kt
@@ -20,7 +20,7 @@ fun KotlinFileExtractor.extractClassSource(
     extractFunctionBodies: Boolean
      */
 ): Label<out DbClassorinterface> {
-    with("class source", c.psiSafe()) {
+    with("class source", c) {
         // OLD: KE1: DeclarationStackAdjuster(c).use {
             val id = useClassSource(c)
             val pkg = c.classId?.packageFqName?.asString() ?: ""
diff --git a/java/kotlin-extractor2/src/main/kotlin/entities/Function.kt b/java/kotlin-extractor2/src/main/kotlin/entities/Function.kt
index 81885251179..08112355d06 100644
--- a/java/kotlin-extractor2/src/main/kotlin/entities/Function.kt
+++ b/java/kotlin-extractor2/src/main/kotlin/entities/Function.kt
@@ -409,7 +409,7 @@ private fun KotlinFileExtractor.forceExtractFunction(
             overriddenAttributes: OverriddenFunctionAttributes? = null
     */
 ): Label<out DbCallable> {
-    with("function", f.psiSafe()) {
+    with("function", f) {
 /*
 OLD: KE1
             DeclarationStackAdjuster(f, overriddenAttributes).use {
diff --git a/java/kotlin-extractor2/src/main/kotlin/utils/Logger.kt b/java/kotlin-extractor2/src/main/kotlin/utils/Logger.kt
index b5656da5a88..a3c0e0af9d6 100644
--- a/java/kotlin-extractor2/src/main/kotlin/utils/Logger.kt
+++ b/java/kotlin-extractor2/src/main/kotlin/utils/Logger.kt
@@ -1,6 +1,10 @@
 package com.github.codeql
 
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.analysis.api.symbols.KaSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.markers.KaNamedSymbol
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtNamed
 import java.io.BufferedWriter
 import java.io.File
 import java.io.FileWriter
@@ -99,9 +103,31 @@ class LogMessage(private val kind: String, private val message: String) {
     }
 }
 
+sealed class PsiElementOrSymbol {
+    abstract fun getLocationString(ftw: FileTrapWriter): String
+    abstract fun getName(): String
+}
+
+data class PsiElementWrapper(val e: PsiElement) : PsiElementOrSymbol() {
+    override fun getLocationString(ftw: FileTrapWriter) = ftw.getLocationString(e)
+    override fun getName() = when (e) {
+        is KtFile -> e.virtualFilePath
+        is KtNamed -> e.nameAsName?.asString() ?: "<missing name>"
+        else -> "<no name>"
+    }
+}
+
+data class SymbolWrapper(val e: KaSymbol) : PsiElementOrSymbol() {
+    override fun getLocationString(ftw: FileTrapWriter) = "file://${ftw.filePath}"
+    override fun getName() = when (e) {
+        is KaNamedSymbol -> e.name.asString()
+        else -> "<no name>"
+    }
+}
+
 data class ExtractorContext(
     val kind: String,
-    val element: PsiElement?,
+    val element: PsiElementOrSymbol,
     val name: String,
     val loc: String
 )