Fix Kotlin fake source mapping

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/symbol/SourceRemapper.java

@@ -13,11 +13,18 @@ static SourceRemapper getSourceRemapper(String sourceFile, SourceMap sourceMap)
     JvmLanguage jvmLanguage = JvmLanguage.of(sourceFile);
     switch (jvmLanguage) {
       case KOTLIN:
-        StratumExt stratum = sourceMap.getStratum("KotlinDebug");
-        if (stratum == null) {
-          throw new IllegalArgumentException("No stratum found for KotlinDebug");
+        StratumExt stratumMain = sourceMap.getStratum("Kotlin");
+        if (stratumMain == null) {
+          stratumMain = sourceMap.getStratum(sourceMap.getDefaultStratumName());
+          if (stratumMain == null) {
+            throw new IllegalArgumentException("No default stratum found");
+          }
         }
-        return new KotlinSourceRemapper(stratum);
+        StratumExt stratumDebug = sourceMap.getStratum("KotlinDebug");
+        if (stratumDebug == null) {
+          throw new IllegalArgumentException("No stratumDebug found for KotlinDebug");
+        }
+        return new KotlinSourceRemapper(stratumMain, stratumDebug);
       default:
         return NOOP_REMAPPER;
     }
@@ -33,19 +40,33 @@ public int remapSourceLine(int line) {
   }
 
   class KotlinSourceRemapper implements SourceRemapper {
-    private final StratumExt stratum;
+    private final StratumExt stratumMain;
+    private final StratumExt stratumDebug;
 
-    public KotlinSourceRemapper(StratumExt stratum) {
-      this.stratum = stratum;
+    public KotlinSourceRemapper(StratumExt stratumMain, StratumExt stratumDebug) {
+      this.stratumMain = stratumMain;
+      this.stratumDebug = stratumDebug;
     }
 
     @Override
     public int remapSourceLine(int line) {
-      Pair<String, Integer> pair = stratum.getInputLine(line);
-      if (pair == null || pair.getRight() == null) {
-        return line;
+      Pair<String, Integer> pairDebug = stratumDebug.getInputLine(line);
+      if (pairDebug == null || pairDebug.getRight() == null) {
+        Pair<String, Integer> pairMain = stratumMain.getInputLine(line);
+        if (pairMain == null || pairMain.getRight() == null) {
+          return line;
+        }
+        String fileId = pairMain.getLeft();
+        String sourceFileName = stratumMain.getSourceFileName(fileId);
+        if (sourceFileName == null) {
+          throw new IllegalArgumentException("Cannot find source filename for fileid=" + fileId);
+        }
+        if (sourceFileName.equals("fake.kt")) {
+          return -1; // no mapping possible
+        }
+        return pairMain.getRight();
       }
-      return pair.getRight();
+      return pairDebug.getRight();
     }
   }
 }

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/symbol/SymbolExtractor.java

@@ -506,7 +506,7 @@ private static MethodLineInfo extractMethodLineInfo(
       if (node.getType() == AbstractInsnNode.LINE) {
         LineNumberNode lineNumberNode = (LineNumberNode) node;
         int newLine = sourceRemapper.remapSourceLine(lineNumberNode.line);
-        if (dedupSet.add(newLine)) {
+        if (newLine > 0 && dedupSet.add(newLine)) {
           lineNo.add(newLine);
         }
         maxLine = Math.max(newLine, maxLine);

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/symbol/SourceRemapperTest.java

@@ -21,9 +21,11 @@ public void noopSourceRemapper() {
   @Test
   public void kotlinSourceRemapper() {
     SourceMap sourceMapMock = mock(SourceMap.class);
-    StratumExt stratumMock = mock(StratumExt.class);
-    when(sourceMapMock.getStratum(eq("KotlinDebug"))).thenReturn(stratumMock);
-    when(stratumMock.getInputLine(eq(42))).thenReturn(Pair.of("", 24));
+    StratumExt stratumMainMock = mock(StratumExt.class);
+    StratumExt stratumDebugMock = mock(StratumExt.class);
+    when(sourceMapMock.getStratum(eq("Kotlin"))).thenReturn(stratumMainMock);
+    when(sourceMapMock.getStratum(eq("KotlinDebug"))).thenReturn(stratumDebugMock);
+    when(stratumDebugMock.getInputLine(eq(42))).thenReturn(Pair.of("", 24));
     SourceRemapper sourceRemapper = SourceRemapper.getSourceRemapper("foo.kt", sourceMapMock);
     assertTrue(sourceRemapper instanceof SourceRemapper.KotlinSourceRemapper);
     assertEquals(24, sourceRemapper.remapSourceLine(42));

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/symbol/SymbolExtractionTransformerTest.java

@@ -1002,7 +1002,7 @@ public void symbolExtraction16() throws IOException, URISyntaxException {
     }
     assertEquals(2, symbolSinkMock.jarScopes.size());
     Scope classScope = symbolSinkMock.jarScopes.get(0).getScopes().get(0);
-    assertScope(classScope, ScopeType.CLASS, CLASS_NAME, 6, 17, SOURCE_FILE, 4, 1);
+    assertScope(classScope, ScopeType.CLASS, CLASS_NAME, 6, 23, SOURCE_FILE, 5, 1);
     assertLangSpecifics(
         classScope.getLanguageSpecifics(),
         asList("public", "final"),
@@ -1024,12 +1024,15 @@ public void symbolExtraction16() throws IOException, URISyntaxException {
     Scope f2MethodScope = classScope.getScopes().get(2);
     assertScope(f2MethodScope, ScopeType.METHOD, "f2", 10, 17, SOURCE_FILE, 3, 1);
     assertLineRanges(f2MethodScope, "10-10", "12-12", "14-14", "16-17");
+    Scope f3MethodScope = classScope.getScopes().get(3);
+    assertScope(f3MethodScope, ScopeType.METHOD, "f3", 21, 23, SOURCE_FILE, 1, 1);
+    assertLineRanges(f3MethodScope, "21-23");
     assertScope(
-        classScope.getScopes().get(3), ScopeType.METHOD, "<clinit>", 0, 0, SOURCE_FILE, 0, 0);
+        classScope.getScopes().get(4), ScopeType.METHOD, "<clinit>", 0, 0, SOURCE_FILE, 0, 0);
 
     Scope companionClassScope = symbolSinkMock.jarScopes.get(1).getScopes().get(0);
     assertScope(
-        companionClassScope, ScopeType.CLASS, CLASS_NAME + "$Companion", 22, 23, SOURCE_FILE, 3, 0);
+        companionClassScope, ScopeType.CLASS, CLASS_NAME + "$Companion", 28, 29, SOURCE_FILE, 3, 0);
     assertLangSpecifics(
         classScope.getLanguageSpecifics(),
         asList("public", "final"),
@@ -1047,8 +1050,8 @@ public void symbolExtraction16() throws IOException, URISyntaxException {
         0,
         0);
     Scope mainMethodScope = companionClassScope.getScopes().get(1);
-    assertScope(mainMethodScope, ScopeType.METHOD, "main", 22, 23, SOURCE_FILE, 1, 1);
-    assertLineRanges(mainMethodScope, "22-23");
+    assertScope(mainMethodScope, ScopeType.METHOD, "main", 28, 29, SOURCE_FILE, 1, 1);
+    assertLineRanges(mainMethodScope, "28-29");
   }
 
   @Test

dd-java-agent/agent-debugger/src/test/resources/CapturedSnapshot301.kt

@@ -15,6 +15,12 @@ class CapturedSnapshot301 {
     return value
   }
 
+  fun f3(value: Int): Int {
+    val list = listOf(value, 2, 3)
+    val max = list.maxOf { it -> it > 0 }
+    return value
+  }
+
   companion object {
     fun main(arg: String): Int {
       val c = CapturedSnapshot301()

dd-java-agent/agent-debugger/src/test/resources/com/datadog/debugger/symboltest/SymbolExtraction16.kt

@@ -17,6 +17,12 @@ class SymbolExtraction16 {
     return value
   }
 
+  fun f3(value: Int): Int {
+    val list = listOf(value, 2, 3)
+    val max = list.maxOf { it -> it > 0 }
+    return value
+  }
+
   companion object {
     fun main(arg: String): Int {
       val c = SymbolExtraction16()

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/stratum/StratumExt.java

@@ -4,12 +4,15 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Comparator;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 public class StratumExt extends AbstractStratum implements Stratum {
   private final List<FileInfo> fileInfo = new ArrayList<>();
+  private Map<String, FileInfo> fileInfoMap;
 
   private int[] lineStart = null;
 
@@ -57,11 +60,28 @@ public String getSourceFile(String fileId) {
     if (fileInfo.isEmpty()) {
       return null;
     }
-    return fileInfo.stream()
-        .filter(f -> f.getFileId().equals(fileId))
-        .findFirst()
-        .map(FileInfo::getInputFilePath)
-        .orElse(null);
+    populateFileInfoMap();
+    FileInfo fileInfo = fileInfoMap.get(fileId);
+    return fileInfo != null ? fileInfo.getInputFilePath() : null;
+  }
+
+  public String getSourceFileName(String fileId) {
+    if (fileInfo.isEmpty()) {
+      return null;
+    }
+    populateFileInfoMap();
+    FileInfo fileInfo = fileInfoMap.get(fileId);
+    return fileInfo != null ? fileInfo.getInputFileName() : null;
+  }
+
+  private void populateFileInfoMap() {
+    if (fileInfoMap != null) {
+      return;
+    }
+    fileInfoMap = new HashMap<>();
+    for (FileInfo fileInfo : fileInfo) {
+      fileInfoMap.put(fileInfo.getFileId(), fileInfo);
+    }
   }
 
   public List<FileInfo> getFileInfo() {