Fix local var hoisting (#7624)

jpbempel · web-flow · commit bb5d791e62d2 · 2024-09-17T13:56:39.000+02:00
Rewrite the algorithm for hoisting local vars:
Go through all local vars defined in the method and determine by name
all local var that are hoistable.
local vars are hoistable if there is no conflict by slot or by name
with another.
If conflict by slot we are allocating a new slot for one of the local.
If conflict by name and type are not compatible we cancel hoisting for
this name. If type compatible, we can add in hoistable list.
Then we are considering all hoistable variables by name and if there
are multiple occurrence for a name, we are merging those variables
in a new slot and rewrite var instructions to use this new slot.
if only one variable we just assign a new slot and rewrite the
instructions in the range.
diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerTransformer.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerTransformer.java
@@ -569,7 +569,9 @@ private List<ToInstrumentInfo> filterAndSortDefinitions(
       // Log and span decoration probe shared the same instrumentor: CaptureContextInstrumentor
       // and therefore need to be instrumented once
       // note: exception probes are log probes and are handled the same way
-      if (definition instanceof LogProbe || definition instanceof SpanDecorationProbe) {
+      if (definition instanceof LogProbe
+          || definition instanceof SpanDecorationProbe
+          || definition instanceof DebuggerProbe) {
         if (definition instanceof ExceptionProbe) {
           if (addedExceptionProbe) {
             continue;
diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/CapturedContextInstrumentor.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/CapturedContextInstrumentor.java
@@ -54,6 +54,7 @@
 import org.objectweb.asm.tree.ClassNode;
 import org.objectweb.asm.tree.FieldInsnNode;
 import org.objectweb.asm.tree.FieldNode;
+import org.objectweb.asm.tree.IincInsnNode;
 import org.objectweb.asm.tree.InsnList;
 import org.objectweb.asm.tree.InsnNode;
 import org.objectweb.asm.tree.JumpInsnNode;
@@ -459,75 +460,161 @@ private void instrumentMethodEnter() {
   }
 
   // Initialize and hoist local variables to the top of the method
-  // if there is name conflict, do nothing for the conflicting local variable
+  // if there is name/slot conflict, do nothing for the conflicting local variable
   private Collection<LocalVariableNode> initAndHoistLocalVars(InsnList insnList) {
     if (methodNode.localVariables == null || methodNode.localVariables.isEmpty()) {
       return Collections.emptyList();
     }
     Map<String, LocalVariableNode> localVarsByName = new HashMap<>();
-    Set<String> duplicateNames = new HashSet<>();
+    Map<Integer, LocalVariableNode> localVarsBySlot = new HashMap<>();
+    Map<String, List<LocalVariableNode>> hoistableVarByName = new HashMap<>();
     for (LocalVariableNode localVar : methodNode.localVariables) {
       int idx = localVar.index - localVarBaseOffset;
       if (idx < argOffset) {
         // this is an argument
         continue;
       }
-      if (!checkDuplicateLocalVar(localVar, localVarsByName)) {
-        duplicateNames.add(localVar.name);
+      checkHoistableLocalVar(localVar, localVarsByName, localVarsBySlot, hoistableVarByName);
+    }
+    // hoist vars
+    List<LocalVariableNode> results = new ArrayList<>();
+    for (Map.Entry<String, List<LocalVariableNode>> entry : hoistableVarByName.entrySet()) {
+      List<LocalVariableNode> hoistableVars = entry.getValue();
+      LocalVariableNode newVarNode;
+      if (hoistableVars.size() > 1) {
+        // merge variables
+        String name = hoistableVars.get(0).name;
+        String desc = hoistableVars.get(0).desc;
+        Type localVarType = getType(desc);
+        int newSlot = newVar(localVarType); // new slot for the local variable
+        newVarNode = new LocalVariableNode(name, desc, null, null, null, newSlot);
+        Set<LabelNode> endLabels = new HashSet<>();
+        for (LocalVariableNode localVar : hoistableVars) {
+          // rewrite each usage of the old var to the new slot
+          rewriteLocalVarInsn(localVar, localVar.index, newSlot);
+          endLabels.add(localVar.end);
+        }
+        hoistVar(insnList, newVarNode);
+        newVarNode.end = findLatestLabel(methodNode.instructions, endLabels);
+        // remove previous local variables
+        methodNode.localVariables.removeIf(hoistableVars::contains);
+        methodNode.localVariables.add(newVarNode);
+      } else {
+        // hoist the single variable and rewrite all its local var instructions
+        newVarNode = hoistableVars.get(0);
+        int oldIndex = newVarNode.index;
+        newVarNode.index = newVar(getType(newVarNode.desc)); // new slot for the local variable
+        rewriteLocalVarInsn(newVarNode, oldIndex, newVarNode.index);
+        hoistVar(insnList, newVarNode);
+      }
+      results.add(newVarNode);
+    }
+    return results;
+  }
+
+  private LabelNode findLatestLabel(InsnList instructions, Set<LabelNode> endLabels) {
+    for (AbstractInsnNode insn = instructions.getLast(); insn != null; insn = insn.getPrevious()) {
+      if (insn instanceof LabelNode && endLabels.contains(insn)) {
+        return (LabelNode) insn;
       }
     }
-    for (String name : duplicateNames) {
-      localVarsByName.remove(name);
+    return null;
+  }
+
+  private void hoistVar(InsnList insnList, LocalVariableNode varNode) {
+    LabelNode labelNode = new LabelNode(); // new start label for the local variable
+    insnList.add(labelNode);
+    varNode.start = labelNode; // update the start label of the local variable
+    Type localVarType = getType(varNode.desc);
+    addStore0Insn(insnList, varNode, localVarType);
+  }
+
+  private static void addStore0Insn(
+      InsnList insnList, LocalVariableNode localVar, Type localVarType) {
+    switch (localVarType.getSort()) {
+      case Type.BOOLEAN:
+      case Type.CHAR:
+      case Type.BYTE:
+      case Type.SHORT:
+      case Type.INT:
+        insnList.add(new InsnNode(Opcodes.ICONST_0));
+        break;
+      case Type.LONG:
+        insnList.add(new InsnNode(Opcodes.LCONST_0));
+        break;
+      case Type.FLOAT:
+        insnList.add(new InsnNode(Opcodes.FCONST_0));
+        break;
+      case Type.DOUBLE:
+        insnList.add(new InsnNode(Opcodes.DCONST_0));
+        break;
+      default:
+        insnList.add(new InsnNode(Opcodes.ACONST_NULL));
+        break;
+    }
+    insnList.add(new VarInsnNode(localVarType.getOpcode(Opcodes.ISTORE), localVar.index));
+  }
+
+  private void checkHoistableLocalVar(
+      LocalVariableNode localVar,
+      Map<String, LocalVariableNode> localVarsByName,
+      Map<Integer, LocalVariableNode> localVarsBySlot,
+      Map<String, List<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);
     }
-    for (LocalVariableNode localVar : localVarsByName.values()) {
-      Type localVarType = getType(localVar.desc);
-      switch (localVarType.getSort()) {
-        case Type.BOOLEAN:
-        case Type.CHAR:
-        case Type.BYTE:
-        case Type.SHORT:
-        case Type.INT:
-          insnList.add(new InsnNode(Opcodes.ICONST_0));
-          break;
-        case Type.LONG:
-          insnList.add(new InsnNode(Opcodes.LCONST_0));
-          break;
-        case Type.FLOAT:
-          insnList.add(new InsnNode(Opcodes.FCONST_0));
-          break;
-        case Type.DOUBLE:
-          insnList.add(new InsnNode(Opcodes.DCONST_0));
-          break;
-        default:
-          insnList.add(new InsnNode(Opcodes.ACONST_NULL));
-          break;
+    if (previousVarByName != null) {
+      // there are multiple local variables with the same name
+      // checking type to see if they are compatible
+      Type previousType = getType(previousVarByName.desc);
+      Type currentType = getType(localVar.desc);
+      if (!ASMHelper.isStoreCompatibleType(previousType, currentType)) {
+        reportWarning(
+            "Local variable "
+                + localVar.name
+                + " has multiple definitions with incompatible types: "
+                + previousVarByName.desc
+                + " and "
+                + localVar.desc);
+        // remove name from hoistable
+        hoistableVarByName.remove(localVar.name);
+        return;
       }
-      insnList.add(new VarInsnNode(localVarType.getOpcode(Opcodes.ISTORE), localVar.index));
+      // Merge variables because compatible type
     }
-    return localVarsByName.values();
+    // by default, there is no conflict => hoistable
+    hoistableVarByName.computeIfAbsent(localVar.name, k -> new ArrayList<>()).add(localVar);
   }
 
-  private boolean checkDuplicateLocalVar(
-      LocalVariableNode localVar, Map<String, LocalVariableNode> localVarsByName) {
-    LocalVariableNode previousVar = localVarsByName.putIfAbsent(localVar.name, localVar);
-    if (previousVar == null) {
-      return true;
+  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;
+      }
     }
-    // there are multiple local variables with the same name
-    // checking type to see if they  are compatible
-    Type previousType = getType(previousVar.desc);
-    Type currentType = getType(localVar.desc);
-    if (!ASMHelper.isStoreCompatibleType(previousType, currentType)) {
-      reportWarning(
-          "Local variable "
-              + localVar.name
-              + " has multiple definitions with incompatible types: "
-              + previousVar.desc
-              + " and "
-              + localVar.desc);
-      return false;
+    for (AbstractInsnNode insn = localVar.start;
+        insn != null && insn != localVar.end;
+        insn = insn.getNext()) {
+      if (insn instanceof VarInsnNode) {
+        VarInsnNode varInsnNode = (VarInsnNode) insn;
+        if (varInsnNode.var == oldSlot) {
+          varInsnNode.var = newSlot;
+        }
+      }
+      if (insn instanceof IincInsnNode) {
+        IincInsnNode iincInsnNode = (IincInsnNode) insn;
+        if (iincInsnNode.var == oldSlot) {
+          iincInsnNode.var = newSlot;
+        }
+      }
     }
-    return true;
   }
 
   private void createInProbeFinallyHandler(LabelNode inProbeStartLabel, LabelNode inProbeEndLabel) {
diff --git a/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturedSnapshotTest.java b/dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturedSnapshotTest.java
@@ -11,6 +11,7 @@
 import static datadog.trace.bootstrap.debugger.util.Redaction.REDACTED_VALUE;
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -1779,7 +1780,7 @@ public void uncaughtExceptionCaptureLocalVars() throws IOException, URISyntaxExc
         "java.lang.RuntimeException",
         "oops",
         "com.datadog.debugger.CapturedSnapshot31.uncaughtException",
-        24);
+        30);
   }
 
   @Test
@@ -1810,6 +1811,18 @@ public void methodProbeLocalVarsDeepScopes() throws IOException, URISyntaxExcept
     int result = Reflect.onClass(testClass).call("main", "deepScopes").get();
     assertEquals(4, result);
     Snapshot snapshot = assertOneSnapshot(listener);
+    // localVarL4 can not be captured/hoisted because same name, same slot
+    //    LocalVariableTable:
+    //    Start  Length  Slot  Name   Signature
+    //       59       3     6 localVarL4   I
+    //       65       3     6 localVarL4   I
+    //       71       3     6 localVarL4   I
+    //       29      45     5 localVarL3   I
+    //       20      54     4 localVarL2   I
+    //       13      64     3 localVarL1   I
+    //        0      79     0  this   Lcom/datadog/debugger/CapturedSnapshot31;
+    //        0      79     1   arg   Ljava/lang/String;
+    //        2      77     2 localVarL0   I
     assertEquals(6, snapshot.getCaptures().getReturn().getLocals().size());
     assertCaptureLocals(snapshot.getCaptures().getReturn(), "localVarL0", "int", "0");
     assertCaptureLocals(snapshot.getCaptures().getReturn(), "localVarL1", "int", "1");
@@ -1852,6 +1865,40 @@ public void methodProbeExceptionLocalVars() throws IOException, URISyntaxExcepti
         expectedFields);
   }
 
+  @Test
+  @DisabledIf(
+      value = "datadog.trace.api.Platform#isJ9",
+      disabledReason = "we cannot get local variable debug info")
+  public void overlappingLocalVarSlot() throws IOException, URISyntaxException {
+    final String CLASS_NAME = "com.datadog.debugger.CapturedSnapshot31";
+    LogProbe probe = createProbeAtExit(PROBE_ID, CLASS_NAME, "overlappingSlots", "(String)");
+    TestSnapshotListener listener = installProbes(CLASS_NAME, probe);
+    Class<?> testClass = compileAndLoadClass(CLASS_NAME);
+    int result = Reflect.onClass(testClass).call("main", "overlappingSlots").get();
+    assertEquals(5, result);
+    Snapshot snapshot = assertOneSnapshot(listener);
+    // i local var cannot be hoisted because of overlapping slots and name
+    assertFalse(snapshot.getCaptures().getReturn().getLocals().containsKey("i"));
+    assertTrue(snapshot.getCaptures().getReturn().getLocals().containsKey("subStr"));
+  }
+
+  @Test
+  @DisabledIf(
+      value = "datadog.trace.api.Platform#isJ9",
+      disabledReason = "we cannot get local variable debug info")
+  public void duplicateLocalDifferentScope() throws IOException, URISyntaxException {
+    final String CLASS_NAME = "com.datadog.debugger.CapturedSnapshot31";
+    LogProbe probe =
+        createProbeAtExit(PROBE_ID, CLASS_NAME, "duplicateLocalDifferentScope", "(String)");
+    TestSnapshotListener listener = installProbes(CLASS_NAME, probe);
+    Class<?> testClass = compileAndLoadClass(CLASS_NAME);
+    int result = Reflect.onClass(testClass).call("main", "duplicateLocalDifferentScope").get();
+    assertEquals(28, result);
+    Snapshot snapshot = assertOneSnapshot(listener);
+    assertCaptureLocals(
+        snapshot.getCaptures().getReturn(), "ch", Character.TYPE.getTypeName(), "e");
+  }
+
   @Test
   public void enumConstructorArgs() throws IOException, URISyntaxException {
     final String CLASS_NAME = "com.datadog.debugger.CapturedSnapshot23";
@@ -2436,7 +2483,7 @@ private TestSnapshotListener installProbes(
     Config config = mock(Config.class);
     when(config.isDebuggerEnabled()).thenReturn(true);
     when(config.isDebuggerClassFileDumpEnabled()).thenReturn(true);
-    when(config.isDebuggerVerifyByteCode()).thenReturn(true);
+    when(config.isDebuggerVerifyByteCode()).thenReturn(false);
     when(config.getFinalDebuggerSnapshotUrl())
         .thenReturn("http://localhost:8126/debugger/v1/input");
     when(config.getFinalDebuggerSymDBUrl()).thenReturn("http://localhost:8126/symdb/v1/input");
@@ -2449,7 +2496,12 @@ private TestSnapshotListener installProbes(
     instr.addTransformer(currentTransformer);
     DebuggerAgentHelper.injectSink(listener);
     DebuggerContext.initProbeResolver(
-        (encodedId) -> resolver(encodedId, logProbes, configuration.getSpanDecorationProbes()));
+        (encodedId) ->
+            resolver(
+                encodedId,
+                logProbes,
+                configuration.getSpanDecorationProbes(),
+                configuration.getDebuggerProbes()));
     DebuggerContext.initClassFilter(new DenyListHelper(null));
     DebuggerContext.initValueSerializer(new JsonSnapshotSerializer());
     if (logProbes != null) {
@@ -2471,7 +2523,8 @@ private TestSnapshotListener installProbes(
   private ProbeImplementation resolver(
       String encodedId,
       Collection<LogProbe> logProbes,
-      Collection<SpanDecorationProbe> spanDecorationProbes) {
+      Collection<SpanDecorationProbe> spanDecorationProbes,
+      Collection<DebuggerProbe> debuggerProbes) {
     for (LogProbe probe : logProbes) {
       if (probe.getProbeId().getEncodedId().equals(encodedId)) {
         return probe;
@@ -2482,6 +2535,11 @@ private ProbeImplementation resolver(
         return probe;
       }
     }
+    for (DebuggerProbe probe : debuggerProbes) {
+      if (probe.getProbeId().getEncodedId().equals(encodedId)) {
+        return probe;
+      }
+    }
     return null;
   }
 
diff --git a/dd-java-agent/agent-debugger/src/test/resources/com/datadog/debugger/CapturedSnapshot31.java b/dd-java-agent/agent-debugger/src/test/resources/com/datadog/debugger/CapturedSnapshot31.java
