Fix Kotlin fake source mapping (#10880)

Fix Kotlin fake source mapping

When there is no mapping to original source code, Kotlin SMAP
use the fake.kt file inside the SoruceDebugExtenstion.
Need to support this to avoid having line not correctly mapped and
that are beyond the end of the file

fix unit test

fix spotless

Co-authored-by: jean-philippe.bempel <jean-philippe.bempel@datadoghq.com>

Author: jpbempel

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() {