Change location of methods inside parameterized types

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

@@ -8,7 +8,6 @@ import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.util.isFileClass
 import org.jetbrains.kotlin.ir.util.packageFqName
-import org.jetbrains.kotlin.ir.util.parentAsClass
 import org.jetbrains.kotlin.ir.util.parentClassOrNull
 import java.io.File
 import java.util.ArrayList
@@ -63,10 +62,7 @@ class ExternalDeclExtractor(val logger: FileLogger, val invocationTrapFile: Stri
                             val trapFile = manager.file
                             val trapTmpFile = File.createTempFile("${trapFile.nameWithoutExtension}.", ".${trapFile.extension}.tmp", trapFile.parentFile)
 
-                            val containingClass = when(irDecl) {
-                                is IrClass -> irDecl
-                                else -> irDecl.parentClassOrNull
-                            }
+                            val containingClass = getContainingClassOrSelf(irDecl)
                             if (containingClass == null) {
                                 logger.warnElement("Unable to get containing class", irDecl)
                                 return

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

@@ -287,14 +287,7 @@ open class KotlinFileExtractor(
             extractClassModifiers(c, id)
             extractClassSupertypes(c, id, if (argsIncludingOuterClasses == null) ExtractSupertypesMode.Raw else ExtractSupertypesMode.Specialised(argsIncludingOuterClasses))
 
-            val locId = if (argsIncludingOuterClasses != null && argsIncludingOuterClasses.isNotEmpty()) {
-                val binaryPath = getIrClassBinaryPath(c)
-                val newTrapWriter = tw.makeFileTrapWriter(binaryPath, true)
-                newTrapWriter.getWholeFileLocation()
-            } else {
-                tw.getLocation(c)
-            }
-
+            val locId = getLocation(c, argsIncludingOuterClasses)
             tw.writeHasLocation(id, locId)
 
             // Extract the outer <-> inner class relationship, passing on any type arguments in excess to this class' parameters.
@@ -304,6 +297,21 @@ open class KotlinFileExtractor(
         }
     }
 
+    private fun getLocation(decl: IrDeclaration, typeArgs: List<IrTypeArgument>?): Label<DbLocation> {
+        return if (typeArgs != null && typeArgs.isNotEmpty()) {
+            val c = getContainingClassOrSelf(decl)
+            if (c == null) {
+                tw.getLocation(decl)
+            } else {
+                val binaryPath = getIrClassBinaryPath(c)
+                val newTrapWriter = tw.makeFileTrapWriter(binaryPath, true)
+                newTrapWriter.getWholeFileLocation()
+            }
+        } else {
+            tw.getLocation(decl)
+        }
+    }
+
     // `typeArgs` can be null to describe a raw generic type.
     // For non-generic types it will be zero-length list.
     fun extractMemberPrototypes(c: IrClass, argsIncludingOuterClasses: List<IrTypeArgument>?, id: Label<out DbClassorinterface>) {
@@ -630,8 +638,6 @@ open class KotlinFileExtractor(
 
                 getFunctionTypeParameters(f).mapIndexed { idx, tp -> extractTypeParameter(tp, idx) }
 
-                val locId = tw.getLocation(f)
-
                 val id =
                     if (idOverride != null)
                         idOverride
@@ -686,6 +692,7 @@ open class KotlinFileExtractor(
                     tw.writeMethodsKotlinType(methodId, returnType.kotlinResult.id)
                 }
 
+                val locId = getLocation(f, classTypeArgsIncludingOuterClasses)
                 tw.writeHasLocation(id, locId)
                 val body = f.body
                 if (body != null && extractBody) {

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

@@ -8,6 +8,7 @@ import org.jetbrains.kotlin.load.kotlin.KotlinJvmBinarySourceElement
 
 import com.intellij.openapi.vfs.VirtualFile
 import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.util.parentClassOrNull
 import org.jetbrains.kotlin.load.kotlin.JvmPackagePartSource
 
 // Adapted from Kotlin's interpreter/Utils.kt function 'internalName'
@@ -77,3 +78,10 @@ fun getIrClassBinaryPath(irClass: IrClass): String {
   // Otherwise, make up a fake location:
     ?: "/!unknown-binary-location/${getIrDeclBinaryName(irClass).replace(".", "/")}.class"
 }
+
+fun getContainingClassOrSelf(decl: IrDeclaration): IrClass? {
+    return when(decl) {
+        is IrClass -> decl
+        else -> decl.parentClassOrNull
+    }
+}