improve local variable inlining

EpicPlayerA10 · EpicPlayerA10 · commit 1fa52fb8cf0e · 2025-06-02T16:03:06.000+02:00
diff --git a/deobfuscator-api/src/main/java/org/objectweb/asm/tree/AbstractInsnNode.java b/deobfuscator-api/src/main/java/org/objectweb/asm/tree/AbstractInsnNode.java
@@ -33,6 +33,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.function.Consumer;
+import java.util.function.Function;
 import java.util.function.Predicate;
 
 import org.jetbrains.annotations.Nullable;
@@ -633,6 +634,53 @@ public AbstractInsnNode walkNextUntil(
     return current;
   }
 
+  /**
+   * Follows jumps until a condition is met. Follows only GOTO
+   *
+   * @param until Condition to stop walking
+   * @return The instruction that matches the condition or {@code null} if no such instruction is found
+   */
+  @Nullable
+  public AbstractInsnNode followJumpsUntil(
+      Function<AbstractInsnNode, PredicateResult> until,
+      Predicate<AbstractInsnNode> filter,
+      Consumer<AbstractInsnNode> consumer
+  ) {
+    AbstractInsnNode current = this;
+    PredicateResult result = null;
+    while (current != null && (result = until.apply(current)) == PredicateResult.CONTINUE) {
+      if (current instanceof JumpInsnNode jumpInsn) {
+        if (current.getOpcode() == GOTO) {
+          current = jumpInsn.label;
+          continue;
+        } else {
+          // If it's not a GOTO, we stop here
+          return null;
+        }
+      }
+
+      if (filter.test(current)) {
+        consumer.accept(current);
+      }
+
+      current = current.getNext();
+    }
+
+    if (result == PredicateResult.SUCCESS) {
+      return current;
+    } else if (result == PredicateResult.FAILED) {
+      return null; // Failed to find a matching instruction
+    } else {
+      return null;
+    }
+  }
+
+  public enum PredicateResult {
+    FAILED,
+    CONTINUE,
+    SUCCESS
+  }
+
   public String namedOpcode() {
     return NamedOpcodes.map(this.getOpcode());
   }
diff --git a/deobfuscator-api/src/main/java/uwu/narumi/deobfuscator/api/helper/AsmHelper.java b/deobfuscator-api/src/main/java/uwu/narumi/deobfuscator/api/helper/AsmHelper.java
@@ -7,8 +7,11 @@
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.Type;
 import org.objectweb.asm.tree.*;
+import org.objectweb.asm.tree.analysis.Frame;
+import org.objectweb.asm.tree.analysis.OriginalSourceValue;
 import org.objectweb.asm.tree.analysis.Value;
 import uwu.narumi.deobfuscator.api.asm.ClassWrapper;
+import uwu.narumi.deobfuscator.api.asm.InsnContext;
 import uwu.narumi.deobfuscator.api.asm.MethodContext;
 import uwu.narumi.deobfuscator.api.asm.MethodRef;
 import uwu.narumi.deobfuscator.api.context.Context;
@@ -271,4 +274,40 @@ public static List<AbstractInsnNode> getInstructionsBetween(
 
     return instructions;
   }
+
+  /**
+   * Counts variable usages in the method.
+   *
+   * @param methodContext Method context to analyze
+   * @return A map of var store instructions -> usage count (load, iinc) in the method.
+   */
+  public static Map<VarInsnNode, Integer> getVarUsages(MethodContext methodContext) {
+    Map<VarInsnNode, Integer> varUsages = new HashMap<>();
+
+    for (AbstractInsnNode insn : methodContext.methodNode().instructions.toArray()) {
+      if ((insn instanceof VarInsnNode && !insn.isVarStore()) || insn instanceof IincInsnNode) {
+        InsnContext insnContext = methodContext.at(insn);
+
+        Frame<OriginalSourceValue> frame = insnContext.frame();
+        if (frame == null) continue;
+
+        int varIndex;
+        if (insn instanceof VarInsnNode varInsnNode) {
+          varIndex = varInsnNode.var;
+        } else {
+          varIndex = ((IincInsnNode) insn).var;
+        }
+
+        OriginalSourceValue localVariableSourceValue = frame.getLocal(varIndex);
+        for (AbstractInsnNode sourceInsn : localVariableSourceValue.insns) {
+          // Save var stores in use
+          if (sourceInsn.isVarStore()) {
+            varUsages.merge((VarInsnNode) sourceInsn, 1, Integer::sum);
+          }
+        }
+      }
+    }
+
+    return varUsages;
+  }
 }
diff --git a/deobfuscator-transformers/src/main/java/uwu/narumi/deobfuscator/core/other/impl/clean/peephole/PopUnUsedLocalVariablesTransformer.java b/deobfuscator-transformers/src/main/java/uwu/narumi/deobfuscator/core/other/impl/clean/peephole/PopUnUsedLocalVariablesTransformer.java
@@ -1,15 +1,12 @@
 package uwu.narumi.deobfuscator.core.other.impl.clean.peephole;
 
 import org.objectweb.asm.tree.AbstractInsnNode;
-import org.objectweb.asm.tree.IincInsnNode;
 import org.objectweb.asm.tree.VarInsnNode;
-import org.objectweb.asm.tree.analysis.Frame;
-import org.objectweb.asm.tree.analysis.OriginalSourceValue;
 import uwu.narumi.deobfuscator.api.asm.InsnContext;
 import uwu.narumi.deobfuscator.api.asm.MethodContext;
+import uwu.narumi.deobfuscator.api.helper.AsmHelper;
 import uwu.narumi.deobfuscator.api.transformer.Transformer;
 
-import java.util.HashSet;
 import java.util.Set;
 
 public class PopUnUsedLocalVariablesTransformer extends Transformer {
@@ -19,32 +16,7 @@ protected void transform() throws Exception {
     scopedClasses().parallelStream().forEach(classWrapper -> classWrapper.methods().parallelStream().forEach(methodNode -> {
       MethodContext methodContext = MethodContext.of(classWrapper, methodNode);
 
-      Set<VarInsnNode> varStoresInUse = new HashSet<>();
-
-      // Find all local variables in use
-      for (AbstractInsnNode insn : methodNode.instructions.toArray()) {
-        if ((insn instanceof VarInsnNode && !insn.isVarStore()) || insn instanceof IincInsnNode) {
-          InsnContext insnContext = methodContext.at(insn);
-
-          Frame<OriginalSourceValue> frame = insnContext.frame();
-          if (frame == null) return;
-
-          int varIndex;
-          if (insn instanceof VarInsnNode varInsnNode) {
-            varIndex = varInsnNode.var;
-          } else {
-            varIndex = ((IincInsnNode) insn).var;
-          }
-
-          OriginalSourceValue localVariableSourceValue = frame.getLocal(varIndex);
-          for (AbstractInsnNode sourceInsn : localVariableSourceValue.insns) {
-            // Save var stores in use
-            if (sourceInsn.isVarStore()) {
-              varStoresInUse.add((VarInsnNode) sourceInsn);
-            }
-          }
-        }
-      }
+      Set<VarInsnNode> varStoresInUse = AsmHelper.getVarUsages(methodContext).keySet();
 
       // Remove all local variables that are not in use
       for (AbstractInsnNode insn : methodNode.instructions.toArray()) {
diff --git a/deobfuscator-transformers/src/main/java/uwu/narumi/deobfuscator/core/other/impl/pool/InlineLocalVariablesTransformer.java b/deobfuscator-transformers/src/main/java/uwu/narumi/deobfuscator/core/other/impl/pool/InlineLocalVariablesTransformer.java
@@ -1,29 +1,23 @@
 package uwu.narumi.deobfuscator.core.other.impl.pool;
 
 import org.objectweb.asm.tree.AbstractInsnNode;
-import org.objectweb.asm.tree.IincInsnNode;
 import org.objectweb.asm.tree.VarInsnNode;
 import org.objectweb.asm.tree.analysis.OriginalSourceValue;
+import uwu.narumi.deobfuscator.api.asm.InsnContext;
 import uwu.narumi.deobfuscator.api.asm.MethodContext;
-import uwu.narumi.deobfuscator.api.asm.matcher.Match;
-import uwu.narumi.deobfuscator.api.asm.matcher.group.SequenceMatch;
-import uwu.narumi.deobfuscator.api.asm.matcher.impl.VarLoadMatch;
+import uwu.narumi.deobfuscator.api.helper.AsmHelper;
 import uwu.narumi.deobfuscator.api.helper.FramedInstructionsStream;
 import uwu.narumi.deobfuscator.api.transformer.Transformer;
 
-import java.util.HashMap;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 
 /**
- * Inlines constant local variables
+ * Inlines local variables (constants and next to each other var load/store pairs).
  */
 public class InlineLocalVariablesTransformer extends Transformer {
-  private static final Match VAR_STORE_LOAD_PAIR = SequenceMatch.of(
-      // Store
-      Match.of(ctx -> ctx.insn().isVarStore()).capture("store"),
-      // Load
-      Match.of(ctx -> ctx.insn().isVarLoad()).capture("load")
-  );
 
   @Override
   protected void transform() throws Exception {
@@ -49,38 +43,68 @@ protected void transform() throws Exception {
           }
         });
 
-    // Find all variable store and load pairs
+    // Inline var load/store pairs that are next to each other and used only once
     scopedClasses().forEach(classWrapper -> classWrapper.methods().forEach(methodNode -> {
+      MethodContext methodContext = MethodContext.of(classWrapper, methodNode);
+
       // Count var usages
-      Map<Integer, Integer> varUsages = new HashMap<>(); // var index -> usage count
+      Map<VarInsnNode, Integer> varUsages = AsmHelper.getVarUsages(methodContext);
+
       for (AbstractInsnNode insn : methodNode.instructions.toArray()) {
-        if (insn instanceof VarInsnNode && insn.isVarLoad()) {
-          varUsages.merge(((VarInsnNode) insn).var, 1, Integer::sum);
-        } else if (insn instanceof IincInsnNode) {
-          varUsages.merge(((IincInsnNode) insn).var, 1, Integer::sum);
-        }
-      }
+        if (insn.isVarStore()) {
+          if (!varUsages.containsKey((VarInsnNode) insn) || varUsages.get((VarInsnNode) insn) > 1) {
+            // If the variable is used more than once, we cannot inline
+            continue;
+          }
 
-      // Now we can find pairs
-      MethodContext methodContext = MethodContext.of(classWrapper, methodNode);
-      VAR_STORE_LOAD_PAIR.findAllMatches(methodContext).forEach(match -> {
-        VarInsnNode varLoadInsn = (VarInsnNode) match.captures().get("load").insn();
-        VarInsnNode varStoreInsn = (VarInsnNode) match.captures().get("store").insn();
-        if (varLoadInsn.var != varStoreInsn.var) {
-          // If they are not the same variable, we cannot inline
-          return;
-        }
-        int varIndex = varLoadInsn.var;
-        if (varUsages.getOrDefault(varIndex, 0) > 1) {
-          // If the variable is used more than once, we cannot inline
-          return;
-        }
+          List<VarInsnNode> varStoreOrder = new ArrayList<>();
+          // Collect all variable indexes in order of their usage
+          AbstractInsnNode end = insn.followJumpsUntil(insn2 -> {
+            if (insn2.isVarStore()) {
+              return AbstractInsnNode.PredicateResult.CONTINUE;
+            }
+            return AbstractInsnNode.PredicateResult.SUCCESS;
+          }, AbstractInsnNode::isVarStore, insn2 -> varStoreOrder.add((VarInsnNode)insn2));
+          if (end == null) {
+            // If we cannot find the end of the sequence, we cannot inline
+            continue;
+          }
+
+          // Check if we have var load instructions in the same order
+          List<VarInsnNode> varLoadOrder = new ArrayList<>();
+          end.followJumpsUntil(insn2 -> {
+            if (insn2.isVarLoad()) {
+              InsnContext insnContext = methodContext.at(insn2);
+              OriginalSourceValue localVariableSourceValue = insnContext.frame().getLocal(((VarInsnNode) insn2).var);
+              if (!localVariableSourceValue.isOneWayProduced()) {
+                // If the variable is not one-way produced, we cannot inline
+                return AbstractInsnNode.PredicateResult.FAILED;
+              }
+              return AbstractInsnNode.PredicateResult.CONTINUE;
+            }
+
+            return AbstractInsnNode.PredicateResult.SUCCESS;
+          }, AbstractInsnNode::isVarLoad, insn2 -> varLoadOrder.add((VarInsnNode)insn2));
+
+          List<Integer> varStoreIndexesOrder = varStoreOrder.stream()
+              .map(varInsn -> varInsn.var)
+              .toList();
+          // In reverse order to match the stack order
+          List<Integer> varLoadIndexesOrder = varLoadOrder.stream()
+              .map(varInsn -> varInsn.var)
+              .sorted(Collections.reverseOrder())
+              .toList();
 
-        // Remove the variable completely and use stack value instead
-        match.removeAll();
+          // Var load and stores are in the same order and next to each other. We can inline them
+          if (varStoreIndexesOrder.equals(varLoadIndexesOrder)) {
+            // Inline (use stack instead of var load/store)
+            varStoreOrder.forEach(varInsn -> methodNode.instructions.remove(varInsn));
+            varLoadOrder.forEach(varInsn -> methodNode.instructions.remove(varInsn));
 
-        markChange();
-      });
+            markChange();
+          }
+        }
+      }
     }));
   }
 }
diff --git a/testData/results/custom-classes/qprotect/sample3/fastcode/IlllIIlIIlllIIIl.dec b/testData/results/custom-classes/qprotect/sample3/fastcode/IlllIIlIIlllIIIl.dec
@@ -365,16 +365,15 @@ public class IlllIIlIIlllIIIl {
                 case 254 -> 134;
                 default -> 29907 - 29696;
             };
-            short var15 = (short)var1;
-            var14 = (~var15 | 0xFF) - ~var15;
-            int var5 = var14 + ~var4 + 1;
+            int var15 = (short)var1;
+            int var5 = (~var15 | 0xFF) - ~var15 + ~var4 + 1;
             if (var5 < 0) {
                 var5 += 256;
             }
 
-            int var17 = (short)var1;
-            var17 = (~var17 | 65535) - ~var17 >>> 8;
-            int var6 = (var17 ^ var4) - 2 * (~var17 & var4);
+            var15 = (short)var1;
+            var14 = (~var15 | 65535) - ~var15 >>> 8;
+            int var6 = (var14 ^ var4) - 2 * (~var14 & var4);
             if (var6 < 0) {
                 var6 += 256;
             }
@@ -383,27 +382,27 @@ public class IlllIIlIIlllIIIl {
                 int var8 = var7 % 2;
                 char var10002 = var3[var7];
                 if (var8 == 0) {
-                    int var19 = var10002;
-                    var3[var7] = (char)((var19 | var5) & ~(var19 & var5));
-                    int var30 = var5 >>> 3;
+                    char var18 = var10002;
+                    var3[var7] = (char)((var18 | var5) & ~(var18 & var5));
+                    int var29 = var5 >>> 3;
                     int var13 = var5 << 5;
-                    var19 = var30;
-                    var30 = (var19 & ~var13) + var13;
-                    char var27 = var3[var7];
-                    var19 = var30;
-                    var19 = var19 & ~var27 | ~var19 & var27;
-                    var5 = (~var19 | 0xFF) - ~var19;
+                    var15 = var29;
+                    var29 = (var15 & ~var13) + var13;
+                    char var26 = var3[var7];
+                    var15 = var29;
+                    var15 = var15 & ~var26 | ~var15 & var26;
+                    var5 = (~var15 | 0xFF) - ~var15;
                 } else {
-                    int var23 = var10002;
-                    var3[var7] = (char)((var23 | var6) - (var23 & var6));
-                    int var32 = var6 >>> 3;
-                    int var28 = var6 << 5;
-                    var23 = var32;
-                    var32 = (var23 & ~var28) + var28;
-                    char var29 = var3[var7];
-                    var23 = var32;
-                    var23 = (var23 | var29) - (var23 & var29);
-                    var6 = (~var23 | 0xFF) - ~var23;
+                    int var22 = var10002;
+                    var3[var7] = (char)((var22 | var6) - (var22 & var6));
+                    int var31 = var6 >>> 3;
+                    int var27 = var6 << 5;
+                    var22 = var31;
+                    var31 = (var22 & ~var27) + var27;
+                    char var28 = var3[var7];
+                    var22 = var31;
+                    var22 = (var22 | var28) - (var22 & var28);
+                    var6 = (~var22 | 0xFF) - ~var22;
                 }
             }
 
