Fix mixed local vars for suspend funs in Kotlin (#7748)

Recent Kotlin compiler generates multiple range for same local vars
and mess up our algorithm for adjustment.
Fix this by made a list of unique local vars

Author: jpbempel

dd-java-agent/agent-debugger/build.gradle

@@ -52,7 +52,8 @@ dependencies {
   testImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-web', version: '2.3.5.RELEASE'
   testImplementation group: 'org.freemarker', name: 'freemarker', version: '2.3.30'
   testImplementation group: 'org.jooq', name: 'joor-java-8', version: '0.9.13'
-  testImplementation group: 'org.jetbrains.kotlin', name: 'kotlin-compiler-embeddable', version: "1.5.21"
+  testImplementation group: 'org.jetbrains.kotlin', name: 'kotlin-compiler-embeddable', version: "1.9.25"
+  testImplementation group: 'org.jetbrains.kotlinx', name: 'kotlinx-coroutines-core', version: "1.0.0"
   testImplementation project(':dd-trace-core')
   testImplementation project(':dd-java-agent:agent-builder')
   testImplementation project(':remote-config:remote-config-core')

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/ASMHelper.java

@@ -244,21 +244,35 @@ public static void adjustLocalVarsBasedOnArgs(
     // https://docs.oracle.com/javase/specs/jvms/se8/html/jvms-2.html#jvms-2.6.1
     // so we reassigned local var in arg slots if they are empty
     if (argTypes.length < localVars.length) {
+      List<LocalVariableNode> uniqueSortedLocalVars = dedupLocalVars(sortedLocalVars);
       int slot = isStatic ? 0 : 1;
       int localVarTableIdx = slot;
       for (org.objectweb.asm.Type t : argTypes) {
         if (slot >= localVars.length) {
           break;
         }
-        if (localVars[slot] == null && localVarTableIdx < sortedLocalVars.size()) {
-          localVars[slot] = sortedLocalVars.get(localVarTableIdx);
+        if (localVars[slot] == null && localVarTableIdx < uniqueSortedLocalVars.size()) {
+          localVars[slot] = uniqueSortedLocalVars.get(localVarTableIdx);
         }
         slot += t.getSize();
         localVarTableIdx++;
       }
     }
   }
 
+  private static List<LocalVariableNode> dedupLocalVars(List<LocalVariableNode> sortedLocalVars) {
+    List<LocalVariableNode> uniqueSortedLocalVars = new ArrayList<>();
+    int maxIndex = sortedLocalVars.get(sortedLocalVars.size() - 1).index;
+    boolean[] usedIndexes = new boolean[maxIndex + 1];
+    for (LocalVariableNode localVariableNode : sortedLocalVars) {
+      if (!usedIndexes[localVariableNode.index]) {
+        uniqueSortedLocalVars.add(localVariableNode);
+        usedIndexes[localVariableNode.index] = true;
+      }
+    }
+    return uniqueSortedLocalVars;
+  }
+
   public static void newInstance(InsnList insnList, org.objectweb.asm.Type type) {
     insnList.add(new TypeInsnNode(Opcodes.NEW, type.getInternalName()));
   }

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturedSnapshotTest.java

@@ -590,8 +590,9 @@ public void sourceFileProbeKotlin() {
     URL resource = CapturedSnapshotTest.class.getResource("/" + CLASS_NAME + ".kt");
     assertNotNull(resource);
     List<File> filesToDelete = new ArrayList<>();
-    Class<?> testClass = KotlinHelper.compileAndLoad(CLASS_NAME, resource.getFile(), filesToDelete);
     try {
+      Class<?> testClass =
+          KotlinHelper.compileAndLoad(CLASS_NAME, resource.getFile(), filesToDelete);
       Object companion = Reflect.onClass(testClass).get("Companion");
       int result = Reflect.on(companion).call("main", "").get();
       assertEquals(48, result);
@@ -607,6 +608,29 @@ public void sourceFileProbeKotlin() {
     }
   }
 
+  @Test
+  public void suspendKotlin() {
+    final String CLASS_NAME = "CapturedSnapshot302";
+    TestSnapshotListener listener =
+        installProbes(CLASS_NAME, createSourceFileProbe(PROBE_ID, CLASS_NAME + ".kt", 9));
+    URL resource = CapturedSnapshotTest.class.getResource("/" + CLASS_NAME + ".kt");
+    assertNotNull(resource);
+    List<File> filesToDelete = new ArrayList<>();
+    try {
+      Class<?> testClass =
+          KotlinHelper.compileAndLoad(CLASS_NAME, resource.getFile(), filesToDelete);
+      Object companion = Reflect.onClass(testClass).get("Companion");
+      int result = Reflect.on(companion).call("main", "").get();
+      assertEquals(0, result);
+      Snapshot snapshot = assertOneSnapshot(listener);
+      assertCaptureFields(snapshot.getCaptures().getLines().get(9), "intField", "int", "42");
+      assertCaptureFields(
+          snapshot.getCaptures().getLines().get(9), "strField", String.class.getTypeName(), "foo");
+    } finally {
+      filesToDelete.forEach(File::delete);
+    }
+  }
+
   @Test
   public void fieldExtractor() throws IOException, URISyntaxException {
     final String CLASS_NAME = "CapturedSnapshot04";

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

@@ -0,0 +1,49 @@
+import kotlinx.coroutines.*
+
+class CapturedSnapshot302 {
+  var intField = 42
+  var strField = "foo"
+  val list: MutableList<Int> = mutableListOf()
+
+  suspend fun process(arg: String): Int {
+    println(intField)
+    try {
+      list.add(1)
+      val content = download(arg)
+      list.add(2)
+      subprocess() {
+        println(arg)
+        println(list)
+        delay(1L)
+        intField = 43
+        strField = "bar"
+      }
+      println(strField)
+      println(content)
+    } catch (e: Exception) {
+      println(list)
+      println(e)
+    }
+    return arg.length
+  }
+
+  suspend fun download(arg: String): String {
+    delay(10L)
+    return arg
+  }
+
+  suspend fun subprocess(f: suspend () -> Unit): Unit {
+    println(intField)
+    f()
+    println(strField)
+  }
+
+
+
+  companion object {
+    fun main(arg: String): Int = runBlocking {
+      val c = CapturedSnapshot302()
+      c.process(arg)
+    }
+  }
+}