[DEX] Smali with .registers or .locals count

REAndroid · Oct 31, 2024 · 2906d0b · 2906d0b
1 parent c89eafa
commit 2906d0b
Show file tree

Hide file tree

Showing 5 changed files with 125 additions and 26 deletions.
diff --git a/src/main/java/com/reandroid/dex/data/CodeItem.java b/src/main/java/com/reandroid/dex/data/CodeItem.java
@@ -325,8 +325,13 @@ public void append(SmaliWriter writer) throws IOException {
         writer.setCurrentRegistersTable(this);
         MethodDef methodDef = getMethodDef();
         writer.newLine();
-        SmaliDirective.LOCALS.append(writer);
-        writer.appendInteger(getLocalRegistersCount());
+        if (writer.isLocalRegistersCount()) {
+            SmaliDirective.LOCALS.append(writer);
+            writer.appendInteger(getLocalRegistersCount());
+        } else {
+            SmaliDirective.REGISTERS.append(writer);
+            writer.appendInteger(getRegistersCount());
+        }
         writer.appendAllWithDoubleNewLine(methodDef.getParameters(true));
         writer.appendAllWithDoubleNewLine(methodDef.getAnnotations(true));
         getInstructionList().append(writer);

diff --git a/src/main/java/com/reandroid/dex/smali/SmaliDirective.java b/src/main/java/com/reandroid/dex/smali/SmaliDirective.java
@@ -35,6 +35,7 @@ public class SmaliDirective implements SmaliFormat {
     public static final SmaliDirective CATCH_ALL;
     public static final SmaliDirective CATCH;
     public static final SmaliDirective LOCALS;
+    public static final SmaliDirective REGISTERS;
 
     public static final SmaliDirective ARRAY_DATA;
     public static final SmaliDirective PACKED_SWITCH;
@@ -67,6 +68,8 @@ public class SmaliDirective implements SmaliFormat {
         ENUM = new SmaliDirective("enum");
 
         LOCALS = new SmaliDirective("locals");
+        REGISTERS = new SmaliDirective("registers");
+
         CATCH = new SmaliDirective("catch", true);
         CATCH_ALL = new SmaliDirective("catchall", true);
         ARRAY_DATA = new SmaliDirective("array-data", true);
@@ -107,7 +110,8 @@ public class SmaliDirective implements SmaliFormat {
                 PROLOGUE,
                 PARAM,
                 END_LOCAL,
-                LOCAL
+                LOCAL,
+                REGISTERS
         };
     }
 

diff --git a/src/main/java/com/reandroid/dex/smali/SmaliWriter.java b/src/main/java/com/reandroid/dex/smali/SmaliWriter.java
@@ -328,6 +328,13 @@ public void setWriterSetting(SmaliWriterSetting writerSetting) {
         this.writerSetting = writerSetting;
     }
 
+    public boolean isLocalRegistersCount() {
+        SmaliWriterSetting setting = getWriterSetting();
+        if (setting != null) {
+            return setting.isLocalRegistersCount();
+        }
+        return true;
+    }
     public SequentialLabelFactory getSequentialLabelFactory() {
         return sequentialLabelFactory;
     }

diff --git a/src/main/java/com/reandroid/dex/smali/SmaliWriterSetting.java b/src/main/java/com/reandroid/dex/smali/SmaliWriterSetting.java
@@ -34,10 +34,12 @@ public class SmaliWriterSetting {
     private List<ClassComment> classCommentList;
     private boolean sequentialLabel;
     private boolean commentUnicodeStrings;
+    private boolean localRegistersCount;
 
     public SmaliWriterSetting() {
         this.sequentialLabel = true;
         this.commentUnicodeStrings = false;
+        this.localRegistersCount = true;
     }
 
     public boolean isSequentialLabel() {
@@ -54,6 +56,13 @@ public void setCommentUnicodeStrings(boolean commentUnicodeStrings) {
         this.commentUnicodeStrings = commentUnicodeStrings;
     }
 
+    public boolean isLocalRegistersCount() {
+        return localRegistersCount;
+    }
+    public void setLocalRegistersCount(boolean localRegistersCount) {
+        this.localRegistersCount = localRegistersCount;
+    }
+
     public void writeResourceIdComment(SmaliWriter writer, long l) throws IOException {
         ResourceIdComment resourceIdComment = getResourceIdComment();
         if(resourceIdComment != null){

diff --git a/src/main/java/com/reandroid/dex/smali/model/SmaliMethod.java b/src/main/java/com/reandroid/dex/smali/model/SmaliMethod.java
@@ -24,25 +24,29 @@
 import com.reandroid.dex.smali.SmaliDirective;
 import com.reandroid.dex.smali.SmaliParseException;
 import com.reandroid.dex.smali.SmaliReader;
+import com.reandroid.dex.smali.SmaliRegion;
 import com.reandroid.dex.smali.SmaliWriter;
+import com.reandroid.dex.smali.SmaliWriterSetting;
 
 import java.io.IOException;
 import java.util.Iterator;
 
 public class SmaliMethod extends SmaliDef implements RegistersTable{
 
     private ProtoKey protoKey;
-    private Integer locals;
 
     private final SmaliParamSet paramSet;
+    private final SmaliRegistersCount smaliRegistersCount;
     private final SmaliCodeSet codeSet;
 
     public SmaliMethod(){
         super();
         this.paramSet = new SmaliParamSet();
+        this.smaliRegistersCount = new SmaliRegistersCount();
         this.codeSet = new SmaliCodeSet();
 
         this.paramSet.setParent(this);
+        this.smaliRegistersCount.setParent(this);
         this.codeSet.setParent(this);
     }
 
@@ -78,11 +82,8 @@ public boolean hasDebugElements(){
     public Iterator<SmaliDebugElement> getDebugElements(){
         return getCodeSet().getDebugElements();
     }
-    public Integer getLocals() {
-        return locals;
-    }
-    public void setLocals(Integer locals) {
-        this.locals = locals;
+    public SmaliRegistersCount getSmaliRegistersCount() {
+        return smaliRegistersCount;
     }
     public ProtoKey getProtoKey() {
         return protoKey;
@@ -127,11 +128,9 @@ public void append(SmaliWriter writer) throws IOException {
         writer.append(getName());
         getProtoKey().append(writer);
         writer.indentPlus();
-        Integer locals = getLocals();
-        if(locals != null){
+        if (hasInstructions()) {
             writer.newLine();
-            SmaliDirective.LOCALS.append(writer);
-            writer.appendInteger(locals);
+            getSmaliRegistersCount().append(writer);
         }
         getParamSet().append(writer);
         if(hasAnnotation()){
@@ -159,8 +158,9 @@ public void parse(SmaliReader reader) throws IOException {
     }
     private boolean parseNoneCode(SmaliReader reader) throws IOException {
         SmaliDirective directive = SmaliDirective.parse(reader, false);
-        if(directive == SmaliDirective.LOCALS){
-            parseLocals(reader);
+        if(directive == SmaliDirective.LOCALS ||
+                directive == SmaliDirective.REGISTERS) {
+            getSmaliRegistersCount().parse(reader);
             return true;
         }
         if(directive == SmaliDirective.ANNOTATION){
@@ -173,11 +173,6 @@ private boolean parseNoneCode(SmaliReader reader) throws IOException {
         }
         return false;
     }
-    private void parseLocals(SmaliReader reader) throws IOException {
-        SmaliParseException.expect(reader, SmaliDirective.LOCALS);
-        reader.skipSpaces();
-        setLocals(reader.readInteger());
-    }
     private void parseName(SmaliReader reader) {
         reader.skipWhitespaces();
         int length = reader.indexOf('(') - reader.position();
@@ -190,7 +185,7 @@ private void parseProto(SmaliReader reader) throws IOException {
 
     @Override
     public int getRegistersCount() {
-        return getLocalRegistersCount() + getParameterRegistersCount();
+        return getSmaliRegistersCount().getRegisters();
     }
 
     @Override
@@ -215,11 +210,7 @@ public boolean ensureLocalRegistersCount(int count) {
     }
     @Override
     public int getLocalRegistersCount() {
-        Integer locals = getLocals();
-        if(locals != null){
-            return locals;
-        }
-        return 0;
+        return getSmaliRegistersCount().getLocals();
     }
 
     @Override
@@ -235,4 +226,87 @@ public String toDebugString() {
         builder.append(getProtoKey());
         return builder.toString();
     }
+
+    public static class SmaliRegistersCount extends Smali implements SmaliRegion {
+
+        private SmaliDirective directive;
+        private int value;
+
+        public SmaliRegistersCount() {
+            this.directive = SmaliDirective.LOCALS;
+        }
+
+        public int getLocals() {
+            int value = getValue();
+            if (!isLocalsMode()) {
+                SmaliMethod method = getParent(SmaliMethod.class);
+                value -= method.getParameterRegistersCount();
+            }
+            return value;
+        }
+        public int getRegisters() {
+            int value = getValue();
+            if (isLocalsMode()) {
+                SmaliMethod method = getParent(SmaliMethod.class);
+                value += method.getParameterRegistersCount();
+            }
+            return value;
+        }
+        public int getValue() {
+            return value;
+        }
+        public void setValue(int value) {
+            this.value = value;
+        }
+
+        public boolean isLocalsMode() {
+            return getSmaliDirective() == SmaliDirective.LOCALS;
+        }
+        public void setLocalsMode(boolean localsMode) {
+            if (localsMode == isLocalsMode()) {
+                return;
+            }
+            SmaliDirective directive;
+            int value;
+            if (localsMode) {
+                value = getLocals();
+                directive = SmaliDirective.LOCALS;
+            } else {
+                value = getRegisters();
+                directive = SmaliDirective.REGISTERS;
+            }
+            setDirective(directive);
+            setValue(value);
+        }
+
+        private void setDirective(SmaliDirective directive) {
+            this.directive = directive;
+        }
+
+        @Override
+        public void parse(SmaliReader reader) throws IOException {
+            SmaliDirective directive = SmaliDirective.parse(reader);
+            if (directive != SmaliDirective.LOCALS && directive != SmaliDirective.REGISTERS) {
+                throw new SmaliParseException("expecting '" + SmaliDirective.LOCALS + "', or '"
+                        + SmaliDirective.REGISTERS + "'", reader);
+            }
+            setDirective(directive);
+            reader.skipSpaces();
+            setValue(reader.readInteger());
+        }
+        @Override
+        public SmaliDirective getSmaliDirective() {
+            return directive;
+        }
+
+        @Override
+        public void append(SmaliWriter writer) throws IOException {
+            SmaliWriterSetting setting = writer.getWriterSetting();
+            if (setting != null) {
+                setLocalsMode(setting.isLocalRegistersCount());
+            }
+            getSmaliDirective().append(writer);
+            writer.appendInteger(getValue());
+        }
+    }
 }