Extract source files for external types

java/kotlin-extractor/src/main/java/com/semmle/extractor/java/PopulateFile.java

@@ -71,15 +71,22 @@ public class PopulateFile {
 		}
 	}
 
-  public Label populateFile(File absoluteFile) {
-    String databasePath = transformer.fileAsDatabaseString(absoluteFile);
-    // Ensure the rewritten path is used from now on.
-    File normalisedFile = new File(databasePath);
-    Label result = tw.getLabelFor("@\"" + escapeKey(databasePath) + ";sourcefile" + "\"");
-    KotlinExtractorDbSchemeKt.writeFiles(tw, result, databasePath);
-    populateParents(normalisedFile, result);
-    return result;
-  }
+    public Label populateFile(File absoluteFile) {
+        return getFileLabel(absoluteFile, true);
+    }
+
+    public Label getFileLabel(File absoluteFile, boolean populateTables) {
+        String databasePath = transformer.fileAsDatabaseString(absoluteFile);
+        Label result = tw.getLabelFor("@\"" + escapeKey(databasePath) + ";sourcefile" + "\"");
+        // Ensure the rewritten path is used from now on.
+
+        if(populateTables) {
+            KotlinExtractorDbSchemeKt.writeFiles(tw, result, databasePath);
+            populateParents(new File(databasePath), result);
+        }
+
+        return result;
+    }
 
 	private Label addFolderTuple(String databasePath) {
 		Label result = tw.getLabelFor("@\"" + escapeKey(databasePath) + ";folder" + "\"");
@@ -102,17 +109,24 @@ public class PopulateFile {
 		KotlinExtractorDbSchemeKt.writeContainerparent(tw, parentLabel, label);
 	}
 
-	public Label relativeFileId(File jarFile, String pathWithinJar) {
+    public Label relativeFileId(File jarFile, String pathWithinJar) {
+        return getFileInJarLabel(jarFile, pathWithinJar, true);
+    }
+
+	public Label getFileInJarLabel(File jarFile, String pathWithinJar, boolean populateTables) {
 		if (pathWithinJar.contains("\\"))
 			throw new CatastrophicError("Invalid jar path: '" + pathWithinJar + "' should not contain '\\'.");
 
-		Label jarFileId = this.populateFile(jarFile);
-		Label jarFileLocation = tw.getLocation(jarFileId,0,0,0,0);
-		KotlinExtractorDbSchemeKt.writeHasLocation(tw, jarFileId, jarFileLocation);
+        String databasePath = transformer.fileAsDatabaseString(jarFile);
+        if(!populateTables)
+            return tw.getLabelFor("@\"" + databasePath + "/" + pathWithinJar + ";jarFile\"");
+
+		Label jarFileId = this.populateFile(jarFile);
+        Label jarFileLocation = tw.getLocation(jarFileId, 0, 0, 0, 0);
+        KotlinExtractorDbSchemeKt.writeHasLocation(tw, jarFileId, jarFileLocation);
 
-		String databasePath = transformer.fileAsDatabaseString(jarFile);
 		StringBuilder fullName = new StringBuilder(databasePath);
-		String[] split = pathWithinJar.split("/");
+        String[] split = pathWithinJar.split("/");
 		Label current = jarFileId;
 		for (int i = 0; i < split.length; i++) {
 			String shortName = split[i];

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

@@ -342,6 +342,21 @@ open class KotlinUsesExtractor(
             ?.let { pluginContext.referenceClass(it.asSingleFqName()) }
             ?.owner
 
+    /**
+     * Gets a KotlinUsesExtractor like this one, except it attributes locations to the file that declares the given class.
+     */
+    fun withSourceFileOfClass(cls: IrClass, populateFileTables: Boolean): KotlinUsesExtractor {
+        val clsFile = cls.fileOrNull
+
+        val newTrapWriter =
+            if (isExternalDeclaration(cls) || clsFile == null)
+                tw.withTargetFile(getIrClassBinaryPath(cls), NullSourceOffsetResolver, populateFileTables)
+            else
+                tw.withTargetFile(clsFile.path, FileSourceOffsetResolver(clsFile.fileEntry))
+
+        return KotlinUsesExtractor(logger, newTrapWriter, dependencyCollector, externalClassExtractor, pluginContext)
+    }
+
     fun useClassInstance(c: IrClass, typeArgs: List<IrTypeArgument>): UseClassInstanceResult {
         // TODO: only substitute in class and function signatures
         //       because within function bodies we can get things like Unit.INSTANCE
@@ -356,7 +371,7 @@ open class KotlinUsesExtractor(
             // If this is a generic type instantiation then it has no
             // source entity, so we need to extract it here
             if (typeArgs.isNotEmpty()) {
-                extractClassInstance(extractClass, typeArgs)
+                this.withSourceFileOfClass(c, false).extractClassInstance(extractClass, typeArgs)
             }
 
             // Extract both the Kotlin and equivalent Java classes, so that we have database entries
@@ -1178,7 +1193,7 @@ open class KotlinFileExtractor(
     fun extractValueParameter(vp: IrValueParameter, parent: Label<out DbCallable>, idx: Int) {
         val id = useValueParameter(vp)
         val type = useType(vp.type)
-        val locId = tw.getLocation(vp.startOffset, vp.endOffset)
+        val locId = tw.getLocation(vp)
         tw.writeParams(id, type.javaResult.id, type.kotlinResult.id, idx, parent, id)
         tw.writeHasLocation(id, locId)
         tw.writeParamName(id, vp.name.asString())

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

@@ -77,6 +77,13 @@ open class TrapWriter (val lm: TrapLabelManager, val bw: BufferedWriter) {
     fun flush() {
         bw.flush()
     }
+
+    /**
+     * Gets a FileTrapWriter like this one (using the same label manager, writer etc), but with the given
+     * default file used in getLocation etc.
+     */
+    fun withTargetFile(filePath: String, sourceOffsetResolver: SourceOffsetResolver, populateFileTables: Boolean = true) =
+        FileTrapWriter(lm, bw, filePath, sourceOffsetResolver, populateFileTables)
 }
 
 abstract class SourceOffsetResolver {
@@ -112,16 +119,17 @@ open class FileTrapWriter (
     lm: TrapLabelManager,
     bw: BufferedWriter,
     val filePath: String,
-    val sourceOffsetResolver: SourceOffsetResolver
+    val sourceOffsetResolver: SourceOffsetResolver,
+    populateFileTables: Boolean = true
 ): TrapWriter (lm, bw) {
     val populateFile = PopulateFile(this)
     val splitFilePath = filePath.split("!/")
     @Suppress("UNCHECKED_CAST")
     val fileId =
         (if(splitFilePath.size == 1)
-            populateFile.populateFile(File(filePath))
+            populateFile.getFileLabel(File(filePath), populateFileTables)
          else
-            populateFile.relativeFileId(File(splitFilePath.get(0)), splitFilePath.get(1))
+            populateFile.getFileInJarLabel(File(splitFilePath.get(0)), splitFilePath.get(1), populateFileTables)
         ) as Label<DbFile>
 
     fun getLocation(e: IrElement): Label<DbLocation> {
@@ -142,8 +150,7 @@ open class FileTrapWriter (
         val endLine =     if(unknownLoc) 0 else sourceOffsetResolver.getLineNumber(endOffset) + 1
         val endColumn =   if(unknownLoc) 0 else sourceOffsetResolver.getColumnNumber(endOffset)
         val endColumn2 =  if(zeroWidthLoc) endColumn + 1 else endColumn
-        val locFileId: Label<DbFile> = if (unknownLoc) unknownFileId else fileId
-        return getLocation(locFileId, startLine, startColumn, endLine, endColumn2)
+        return getLocation(fileId, startLine, startColumn, endLine, endColumn2)
     }
     fun getLocationString(e: IrElement): String {
         if (e.startOffset == -1 && e.endOffset == -1) {