Fix hoisting local vars for Kotlin code (#7758)

javac begins the range of a local var in the LocalVariableTable just
after the store instruction at the init.
kotlinc is more free to instroduce nop/other instructions in-between.
need to scan backward from the start of the range to find the last
store of this slot to begin the rewrite of the slot needed to hoist
correctly local vars.

Author: jpbempel

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

@@ -467,7 +467,7 @@ private Collection<LocalVariableNode> initAndHoistLocalVars(InsnList insnList) {
     }
     Map<String, LocalVariableNode> localVarsByName = new HashMap<>();
     Map<Integer, LocalVariableNode> localVarsBySlot = new HashMap<>();
-    Map<String, List<LocalVariableNode>> hoistableVarByName = new HashMap<>();
+    Map<String, Set<LocalVariableNode>> hoistableVarByName = new HashMap<>();
     for (LocalVariableNode localVar : methodNode.localVariables) {
       int idx = localVar.index - localVarBaseOffset;
       if (idx < argOffset) {
@@ -479,13 +479,14 @@ private Collection<LocalVariableNode> initAndHoistLocalVars(InsnList insnList) {
     removeDuplicatesFromArgs(hoistableVarByName, localVarsBySlotArray);
     // hoist vars
     List<LocalVariableNode> results = new ArrayList<>();
-    for (Map.Entry<String, List<LocalVariableNode>> entry : hoistableVarByName.entrySet()) {
-      List<LocalVariableNode> hoistableVars = entry.getValue();
+    for (Map.Entry<String, Set<LocalVariableNode>> entry : hoistableVarByName.entrySet()) {
+      Set<LocalVariableNode> hoistableVars = entry.getValue();
       LocalVariableNode newVarNode;
       if (hoistableVars.size() > 1) {
         // merge variables
-        String name = hoistableVars.get(0).name;
-        String desc = hoistableVars.get(0).desc;
+        LocalVariableNode firstHoistableVar = hoistableVars.iterator().next();
+        String name = firstHoistableVar.name;
+        String desc = firstHoistableVar.desc;
         Type localVarType = getType(desc);
         int newSlot = newVar(localVarType); // new slot for the local variable
         newVarNode = new LocalVariableNode(name, desc, null, null, null, newSlot);
@@ -502,7 +503,7 @@ private Collection<LocalVariableNode> initAndHoistLocalVars(InsnList insnList) {
         methodNode.localVariables.add(newVarNode);
       } else {
         // hoist the single variable and rewrite all its local var instructions
-        newVarNode = hoistableVars.get(0);
+        newVarNode = hoistableVars.iterator().next();
         int oldIndex = newVarNode.index;
         newVarNode.index = newVar(getType(newVarNode.desc)); // new slot for the local variable
         rewriteLocalVarInsn(newVarNode, oldIndex, newVarNode.index);
@@ -514,7 +515,7 @@ private Collection<LocalVariableNode> initAndHoistLocalVars(InsnList insnList) {
   }
 
   private void removeDuplicatesFromArgs(
-      Map<String, List<LocalVariableNode>> hoistableVarByName,
+      Map<String, Set<LocalVariableNode>> hoistableVarByName,
       LocalVariableNode[] localVarsBySlotArray) {
     for (int idx = 0; idx < argOffset; idx++) {
       LocalVariableNode localVar = localVarsBySlotArray[idx];
@@ -573,13 +574,13 @@ private void checkHoistableLocalVar(
       LocalVariableNode localVar,
       Map<String, LocalVariableNode> localVarsByName,
       Map<Integer, LocalVariableNode> localVarsBySlot,
-      Map<String, List<LocalVariableNode>> hoistableVarByName) {
+      Map<String, Set<LocalVariableNode>> hoistableVarByName) {
     LocalVariableNode previousVarBySlot = localVarsBySlot.putIfAbsent(localVar.index, localVar);
     LocalVariableNode previousVarByName = localVarsByName.putIfAbsent(localVar.name, localVar);
     if (previousVarBySlot != null) {
       // there are multiple local variables with the same slot but different names
       // by hoisting in a new slot, we can avoid the conflict
-      hoistableVarByName.computeIfAbsent(localVar.name, k -> new ArrayList<>()).add(localVar);
+      hoistableVarByName.computeIfAbsent(localVar.name, k -> new HashSet<>()).add(localVar);
     }
     if (previousVarByName != null) {
       // there are multiple local variables with the same name
@@ -601,18 +602,11 @@ private void checkHoistableLocalVar(
       // Merge variables because compatible type
     }
     // by default, there is no conflict => hoistable
-    hoistableVarByName.computeIfAbsent(localVar.name, k -> new ArrayList<>()).add(localVar);
+    hoistableVarByName.computeIfAbsent(localVar.name, k -> new HashSet<>()).add(localVar);
   }
 
   private void rewriteLocalVarInsn(LocalVariableNode localVar, int oldSlot, int newSlot) {
-    // previous insn could be a store to index that need to be rewritten as well
-    AbstractInsnNode previous = localVar.start.getPrevious();
-    if (previous instanceof VarInsnNode) {
-      VarInsnNode varInsnNode = (VarInsnNode) previous;
-      if (varInsnNode.var == oldSlot) {
-        varInsnNode.var = newSlot;
-      }
-    }
+    rewritePreviousStoreInsn(localVar, oldSlot, newSlot);
     for (AbstractInsnNode insn = localVar.start;
         insn != null && insn != localVar.end;
         insn = insn.getNext()) {
@@ -631,6 +625,42 @@ private void rewriteLocalVarInsn(LocalVariableNode localVar, int oldSlot, int ne
     }
   }
 
+  // Previous insn(s) to local var range could be a store to index that need to be rewritten as well
+  // LocalVariableTable ranges starts after the init of the local var.
+  // ex:
+  //    9: iconst_0
+  //   10: astore_1
+  //   11: aload_1
+  // range for slot 1 starts at 11
+  // javac always starts the range right after the init of the local var, so we can just look for
+  // the previous instruction
+  // for kotlinc, many instructions can separate the init and the range start
+  // ex:
+  //    LocalVariableTable:
+  //    Start  Length  Slot  Name   Signature
+  //    70     121     4 $i$f$map   I
+  //
+  //    64: istore        4
+  //    66: aload_2
+  //    67: iconst_2
+  //    68: iconst_1
+  //    69: bastore
+  //    70: aload_3
+  private static void rewritePreviousStoreInsn(
+      LocalVariableNode localVar, int oldSlot, int newSlot) {
+    AbstractInsnNode previous = localVar.start.getPrevious();
+    while (previous != null
+        && (!(previous instanceof VarInsnNode) || ((VarInsnNode) previous).var != oldSlot)) {
+      previous = previous.getPrevious();
+    }
+    if (previous != null) {
+      VarInsnNode varInsnNode = (VarInsnNode) previous;
+      if (varInsnNode.var == oldSlot) {
+        varInsnNode.var = newSlot;
+      }
+    }
+  }
+
   private void createInProbeFinallyHandler(LabelNode inProbeStartLabel, LabelNode inProbeEndLabel) {
     LabelNode handlerLabel = new LabelNode();
     InsnList handler = new InsnList();

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

@@ -631,6 +631,27 @@ public void suspendKotlin() {
     }
   }
 
+  @Test
+  public void hoistVarKotlin() {
+    final String CLASS_NAME = "CapturedSnapshot303";
+    TestSnapshotListener listener =
+        installProbes(
+            CLASS_NAME, createProbeAtExit(PROBE_ID, CLASS_NAME + "$Companion", "main", null));
+    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);
+    } 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/CapturedSnapshot303.kt

@@ -0,0 +1,11 @@
+class CapturedSnapshot303 {
+  companion object {
+    fun main(arg: String): Int {
+      listOf(1, 2, 3, 4, 5)
+        .map { it * 2 }
+        .filter { it % 2 == 0 }
+        .forEach { println(it) }
+      return arg.length
+    }
+  }
+}