Complete and fix-up the implementation of autofixing instructions

git-svn-id: https://smali.googlecode.com/svn/trunk@511 55b6fa8a-2a1e-11de-a435-ffa8d773f76a

baksmali/src/main/java/org/jf/baksmali/Adaptors/MethodDefinition.java

@@ -526,8 +526,8 @@ public class MethodDefinition {
                 return;
             }
             for (CodeItem.TryItem tryItem: codeItem.getTries()) {
-                int startAddress = tryItem.startAddress;
+                int startAddress = tryItem.getStartAddress();
-                int endAddress = tryItem.startAddress + tryItem.instructionCount;
+                int endAddress = tryItem.getStartAddress() + tryItem.getInstructionCount();
 
                 /**
                  * The end address points to the address immediately after the end of the last

dexlib/src/main/java/org/jf/dexlib/CodeItem.java

@@ -31,6 +31,8 @@ package org.jf.dexlib;
 import org.jf.dexlib.Code.*;
 import org.jf.dexlib.Code.Format.Instruction20t;
 import org.jf.dexlib.Code.Format.Instruction30t;
+import org.jf.dexlib.Code.Format.Instruction21c;
+import org.jf.dexlib.Code.Format.Instruction31c;
 import org.jf.dexlib.Util.*;
 import org.jf.dexlib.Debug.DebugInstructionIterator;
 
@@ -400,43 +402,61 @@ public class CodeItem extends Item<CodeItem> {
      *
      * The above fixes are applied iteratively, until no more fixes have been performed
      */
-    public void fixInstructions() {
+    public void fixInstructions(boolean fixStringConst, boolean fixGoto) {
         boolean didSomething = false;
 
-        int currentCodeOffset = 0;
+        do
-        for (int i=0; i<instructions.length; i++) {
+        {
-            Instruction instruction = instructions[i];
+            didSomething = false;
 
-            if (instruction.opcode == Opcode.GOTO) {
+            int currentCodeOffset = 0;
-                int offset = ((OffsetInstruction)instruction).getOffset();
+            for (int i=0; i<instructions.length; i++) {
+                Instruction instruction = instructions[i];
 
-                if (((byte)offset) != offset) {
+                if (fixGoto && instruction.opcode == Opcode.GOTO) {
-                    //the offset doesn't fit within a byte, we need to upgrade to a goto/16 or goto/32
+                    int offset = ((OffsetInstruction)instruction).getOffset();
 
-                    if ((short)offset == offset) {
+                    if (((byte)offset) != offset) {
-                        //the offset fits in a short, so upgrade to a goto/16            h
+                        //the offset doesn't fit within a byte, we need to upgrade to a goto/16 or goto/32
-                        replaceInstructionAtOffset(currentCodeOffset, new Instruction20t(Opcode.GOTO_16, offset));
+
+                        if ((short)offset == offset) {
+                            //the offset fits in a short, so upgrade to a goto/16            h
+                            replaceInstructionAtOffset(currentCodeOffset, new Instruction20t(Opcode.GOTO_16, offset));
+                        }
+                        else {
+                            //The offset won't fit into a short, we have to upgrade to a goto/32
+                            replaceInstructionAtOffset(currentCodeOffset, new Instruction30t(Opcode.GOTO_32, offset));
+                        }
+                        didSomething = true;
+                        break;
                     }
-                    else {
+                } else if (fixGoto && instruction.opcode == Opcode.GOTO_16) {
-                        //The offset won't fit into a short, we have to upgrade to a goto/32
+                    int offset = ((OffsetInstruction)instruction).getOffset();
+
+                    if (((short)offset) != offset) {
+                        //the offset doesn't fit within a short, we need to upgrade to a goto/32
                         replaceInstructionAtOffset(currentCodeOffset, new Instruction30t(Opcode.GOTO_32, offset));
+                        didSomething = true;
+                        break;
+                    }
+                } else if (fixStringConst && instruction.opcode == Opcode.CONST_STRING) {
+                    Instruction21c constStringInstruction = (Instruction21c)instruction;
+                    if (constStringInstruction.getReferencedItem().getIndex() > 0xFFFF) {
+                        replaceInstructionAtOffset(currentCodeOffset, new Instruction31c(Opcode.CONST_STRING_JUMBO,
+                                (short)constStringInstruction.getRegisterA(),
+                                constStringInstruction.getReferencedItem()));
+                        didSomething = true;
+                        break;
                     }
                 }
-            } else if (instruction.opcode == Opcode.GOTO_16) {
-                int offset = ((OffsetInstruction)instruction).getOffset();
 
-                if (((short)offset) != offset) {
+                currentCodeOffset += instruction.getSize(currentCodeOffset);
-                    //the offset doesn't fit within a short, we need to upgrade to a goto/32
-                    replaceInstructionAtOffset(currentCodeOffset, new Instruction30t(Opcode.GOTO_32, offset));
-                }
             }
-
+        }while(didSomething);
-            currentCodeOffset += instruction.getSize(currentCodeOffset);
-        }
     }
 
     private void replaceInstructionAtOffset(int offset, Instruction replacementInstruction) {
-        Instruction originalInstruction = null;
+         Instruction originalInstruction = null;
 
         int[] originalInstructionOffsets = new int[instructions.length];
         SparseIntArray originalSwitchOffsetByOriginalSwitchDataOffset = new SparseIntArray();
@@ -476,6 +496,7 @@ public class CodeItem extends Item<CodeItem> {
             return;
         }
 
+        //TODO: replace these with a callable delegate
         final SparseIntArray originalOffsetsByNewOffset = new SparseIntArray();
         final SparseIntArray newOffsetsByOriginalOffset = new SparseIntArray();
 
@@ -505,8 +526,11 @@ public class CodeItem extends Item<CodeItem> {
 
                 assert newOffsetsByOriginalOffset.indexOfKey(originalInstructionTarget) >= 0;
                 int newInstructionTarget = newOffsetsByOriginalOffset.get(originalInstructionTarget);
-                if (newInstructionTarget != originalInstructionTarget) {
+
-                    offsetInstruction.updateOffset(newInstructionTarget);
+                int newOffset = (newInstructionTarget - currentCodeOffset) / 2;
+
+                if (newOffset != offsetInstruction.getOffset()) {
+                    offsetInstruction.updateOffset(newOffset);
                 }
             } else if (instruction instanceof MultiOffsetInstruction) {
                 MultiOffsetInstruction multiOffsetInstruction = (MultiOffsetInstruction)instruction;
@@ -520,39 +544,62 @@ public class CodeItem extends Item<CodeItem> {
                     continue;
                 }
 
+                assert newOffsetsByOriginalOffset.indexOfKey(originalSwitchOffset) >= 0;
+                int newSwitchOffset = newOffsetsByOriginalOffset.get(originalSwitchOffset);
+
                 int[] targets = multiOffsetInstruction.getTargets();
                 for (int t=0; t<targets.length; t++) {
-                    int originalTargetOffset = originalSwitchOffset + targets[t];
+                    int originalTargetOffset = originalSwitchOffset + targets[t]*2;
                     assert newOffsetsByOriginalOffset.indexOfKey(originalTargetOffset) >= 0;
                     int newTargetOffset = newOffsetsByOriginalOffset.get(originalTargetOffset);
-                    if (newTargetOffset != originalTargetOffset) {
+                    int newOffset = (newTargetOffset - newSwitchOffset)/2;
-                        multiOffsetInstruction.updateTarget(t, newTargetOffset);
+                    if (newOffset != targets[t]) {
+                        multiOffsetInstruction.updateTarget(t, newOffset);
                     }
                 }
             }
             currentCodeOffset += instruction.getSize(currentCodeOffset);
         }
 
-        final byte[] encodedDebugInfo = debugInfo.getEncodedDebugInfo();
+        if (debugInfo != null) {
+            final byte[] encodedDebugInfo = debugInfo.getEncodedDebugInfo();
 
-        ByteArrayInput debugInput = new ByteArrayInput(encodedDebugInfo);
+            ByteArrayInput debugInput = new ByteArrayInput(encodedDebugInfo);
 
-        DebugInstructionFixer debugInstructionFixer = new DebugInstructionFixer(encodedDebugInfo,
+            DebugInstructionFixer debugInstructionFixer = new DebugInstructionFixer(encodedDebugInfo,
                 newOffsetsByOriginalOffset, originalOffsetsByNewOffset);
-        DebugInstructionIterator.IterateInstructions(debugInput, debugInstructionFixer);
+            DebugInstructionIterator.IterateInstructions(debugInput, debugInstructionFixer);
 
-        debugInfo.setEncodedDebugInfo(debugInstructionFixer.result);
+            assert debugInstructionFixer.result != null;
 
+            debugInfo.setEncodedDebugInfo(debugInstructionFixer.result);
+        }
 
-        for (EncodedCatchHandler encodedCatchHandler: encodedCatchHandlers) {
+        if (encodedCatchHandlers != null) {
-            if (encodedCatchHandler.catchAllHandlerAddress != -1) {
+            for (EncodedCatchHandler encodedCatchHandler: encodedCatchHandlers) {
-                assert newOffsetsByOriginalOffset.indexOfKey(encodedCatchHandler.catchAllHandlerAddress) >= 0;
+                if (encodedCatchHandler.catchAllHandlerAddress != -1) {
-                encodedCatchHandler.catchAllHandlerAddress =
+                    assert newOffsetsByOriginalOffset.indexOfKey(encodedCatchHandler.catchAllHandlerAddress*2) >= 0;
-                        newOffsetsByOriginalOffset.get(encodedCatchHandler.catchAllHandlerAddress);
+                    encodedCatchHandler.catchAllHandlerAddress =
+                            newOffsetsByOriginalOffset.get(encodedCatchHandler.catchAllHandlerAddress*2)/2;
+                }
+
+                for (EncodedTypeAddrPair handler: encodedCatchHandler.handlers) {
+                    assert newOffsetsByOriginalOffset.indexOfKey(handler.handlerAddress*2) >= 0;
+                    handler.handlerAddress = newOffsetsByOriginalOffset.get(handler.handlerAddress*2)/2;
+                }
             }
+        }
 
-            for (EncodedTypeAddrPair handler: encodedCatchHandler.handlers) {
+        if (this.tries != null) {
-                handler.handlerAddress = newOffsetsByOriginalOffset.get(handler.handlerAddress);
+            for (TryItem tryItem: tries) {
+                int startAddress = tryItem.startAddress;
+                int endAddress = tryItem.startAddress + tryItem.instructionCount;
+
+                assert newOffsetsByOriginalOffset.indexOfKey(startAddress * 2) >= 0;
+                tryItem.startAddress = newOffsetsByOriginalOffset.get(startAddress * 2)/2;
+
+                assert newOffsetsByOriginalOffset.indexOfKey(endAddress * 2) >= 0;
+                tryItem.instructionCount = newOffsetsByOriginalOffset.get(endAddress * 2)/2 - tryItem.startAddress;
             }
         }
     }
@@ -560,62 +607,62 @@ public class CodeItem extends Item<CodeItem> {
     private class DebugInstructionFixer extends DebugInstructionIterator.ProcessRawDebugInstructionDelegate {
         private int address = 0;
         private SparseIntArray newOffsetsByOriginalOffset;
-        private SparseIntArray originalOffsetByNewOffset;
+        private SparseIntArray originalOffsetsByNewOffset;
         private final byte[] originalEncodedDebugInfo;
         public byte[] result = null;
 
-        public DebugInstructionFixer(byte[] originalEncodedDebugInfo, SparseIntArray newOffsetsbyOriginalOffset,
+        public DebugInstructionFixer(byte[] originalEncodedDebugInfo, SparseIntArray newOffsetsByOriginalOffset,
                                      SparseIntArray originalOffsetsByNewOffset) {
             this.newOffsetsByOriginalOffset = newOffsetsByOriginalOffset;
-            this.originalOffsetByNewOffset = originalOffsetByNewOffset;
+            this.originalOffsetsByNewOffset = originalOffsetsByNewOffset;
             this.originalEncodedDebugInfo = originalEncodedDebugInfo;
         }
 
 
         @Override
-        public void ProcessAdvancePC(int startOffset, int length, int addressDiff) {
+        public void ProcessAdvancePC(int startOffset, int length, int addressDelta) {
-            if (result != null) {
-                return;
-            }
-
-            int newOffset = newOffsetsByOriginalOffset.get((address + addressDiff)*2, -1);
-            assert newOffset != -1;
-            newOffset = newOffset / 2;
-
-            if (newOffset != address) {
-                int newAddressDiff = addressDiff + newOffset - address;
-                assert newAddressDiff > 0;
-                int addressDiffSize = Leb128Utils.unsignedLeb128Size(newAddressDiff);
-
-                result = new byte[originalEncodedDebugInfo.length + addressDiffSize - (length - 1)];
-
-                System.arraycopy(originalEncodedDebugInfo, 0, result, 0, startOffset);
-
-                originalEncodedDebugInfo[startOffset] = 0x01; //DBG_ADVANCE_PC debug opcode
-                Leb128Utils.writeUnsignedLeb128(newAddressDiff, originalEncodedDebugInfo, startOffset+1);
-
-                System.arraycopy(originalEncodedDebugInfo, startOffset+length, originalEncodedDebugInfo,
-                        startOffset + addressDiffSize + 1,
-                        originalEncodedDebugInfo.length - (startOffset + addressDiffSize + 1));
-            }
-
-            address += addressDiff;
-        }
-
-        @Override
-        public void ProcessSpecialOpcode(int startOffset, int debugOpcode, int lineDelta,
-                                         int addressDelta) {
-            if (result != null) {
-                return;
-            }
-
             address += addressDelta;
+
+            if (result != null) {
+                return;
+            }
+
             int newOffset = newOffsetsByOriginalOffset.get(address*2, -1);
             assert newOffset != -1;
             newOffset = newOffset / 2;
 
             if (newOffset != address) {
-                int newAddressDelta = addressDelta + newOffset - address;
+                int newAddressDelta = newOffset - (address - addressDelta);
+                assert newAddressDelta > 0;
+                int addressDiffSize = Leb128Utils.unsignedLeb128Size(newAddressDelta);
+
+                result = new byte[originalEncodedDebugInfo.length + addressDiffSize - (length - 1)];
+
+                System.arraycopy(originalEncodedDebugInfo, 0, result, 0, startOffset);
+
+                result[startOffset] = 0x01; //DBG_ADVANCE_PC debug opcode
+                Leb128Utils.writeUnsignedLeb128(newAddressDelta, result, startOffset+1);
+
+                System.arraycopy(originalEncodedDebugInfo, startOffset+length, result,
+                        startOffset + addressDiffSize + 1,
+                        originalEncodedDebugInfo.length - (startOffset + addressDiffSize + 1));
+            }
+        }
+
+        @Override
+        public void ProcessSpecialOpcode(int startOffset, int debugOpcode, int lineDelta,
+                                         int addressDelta) {
+            address += addressDelta;
+            if (result != null) {
+                return;
+            }
+
+            int newOffset = newOffsetsByOriginalOffset.get(address*2, -1);
+            assert newOffset != -1;
+            newOffset = newOffset / 2;
+
+            if (newOffset != address) {
+                int newAddressDelta = newOffset - (address - addressDelta);
                 assert newAddressDelta > 0;
 
                 //if the new address delta won't fit in the special opcode, we need to insert
@@ -624,7 +671,7 @@ public class CodeItem extends Item<CodeItem> {
                     int additionalAddressDelta = newOffset - address;
                     int additionalAddressDeltaSize = Leb128Utils.signedLeb128Size(additionalAddressDelta);
 
-                    result = new byte[result.length + additionalAddressDeltaSize + 1];
+                    result = new byte[originalEncodedDebugInfo.length + additionalAddressDeltaSize + 1];
 
                     System.arraycopy(originalEncodedDebugInfo, 0, result, 0, startOffset);
                     result[startOffset] = 0x01; //DBG_ADVANCE_PC
@@ -633,7 +680,7 @@ public class CodeItem extends Item<CodeItem> {
                             startOffset+additionalAddressDeltaSize+1,
                             result.length - (startOffset+additionalAddressDeltaSize+1));
                 } else {
-                    result = new byte[result.length];
+                    result = new byte[originalEncodedDebugInfo.length];
                     System.arraycopy(originalEncodedDebugInfo, 0, result, 0, result.length);
                     result[startOffset] = DebugInfoBuilder.calculateSpecialOpcode(lineDelta,
                             newAddressDelta);
@@ -646,12 +693,12 @@ public class CodeItem extends Item<CodeItem> {
         /**
          * The address (in 2-byte words) within the code where the try block starts
          */
-        public final int startAddress;
+        private int startAddress;
 
         /**
          * The number of 2-byte words that the try block covers
          */
-        public final int instructionCount;
+        private int instructionCount;
 
         /**
          * The associated exception handler
@@ -703,6 +750,20 @@ public class CodeItem extends Item<CodeItem> {
             out.writeShort(instructionCount);
             out.writeShort(encodedCatchHandler.getOffsetInList());
         }
+
+        /**
+         * @return The address (in 2-byte words) within the code where the try block starts
+         */
+        public int getStartAddress() {
+            return startAddress;
+        }
+
+        /**
+         * @return The number of 2-byte words that the try block covers
+         */
+        public int getInstructionCount() {
+            return instructionCount;
+        }
     }
 
     public static class EncodedCatchHandler {

dexlib/src/main/java/org/jf/dexlib/Util/DeodexUtil.java

@@ -166,8 +166,8 @@ public class DeodexUtil {
                     handlers[i] = insnsMap.get(tryItem.encodedCatchHandler.handlers[i].getHandlerAddress());
                 }
 
-                int insnoffset = tryItem.startAddress;
+                int insnoffset = tryItem.getStartAddress();
-                while (insnoffset < tryItem.startAddress + tryItem.instructionCount) {
+                while (insnoffset < tryItem.getStartAddress() + tryItem.getInstructionCount()) {
                     insn i = insnsMap.get(insnoffset);
 
                     i.exceptionHandlers = handlers;

smali-integration-tests/src/test/smali/junit-tests/TestSuite.smali

@@ -47,6 +47,7 @@
                 LFormat32x;,
                 LFormat35c;,
                 LFormat3rc;,
-                LFormat51l;
+                LFormat51l;,
+                LGotoTest;
             }
 .end annotation

smali/src/main/antlr3/org/jf/smali/smaliTreeWalker.g

@@ -880,11 +880,7 @@ instruction[int totalMethodRegisters, int methodParameterRegisters, List<Instruc
 			
 			int addressOffset = $offset_or_label.offsetValue;
 
-			if (addressOffset < Byte.MIN_VALUE || addressOffset > Byte.MAX_VALUE) {
+			$instructions.add(new Instruction10t(opcode, addressOffset));
-				throw new SemanticException(input, "The offset/label is out of range. The offset is " + Integer.toString(addressOffset) + " and the range for this opcode is [-128, 127].");
-			}
-			
-			$instructions.add(new Instruction10t(opcode, (byte)addressOffset));
 		}
 	|	//e.g. return
 		^(I_STATEMENT_FORMAT10x INSTRUCTION_FORMAT10x)
@@ -927,11 +923,7 @@ instruction[int totalMethodRegisters, int methodParameterRegisters, List<Instruc
 			
 			int addressOffset = $offset_or_label.offsetValue;
 
-			if (addressOffset < Short.MIN_VALUE || addressOffset > Short.MAX_VALUE) {
+			$instructions.add(new Instruction20t(opcode, addressOffset));
-				throw new SemanticException(input, "The offset/label is out of range. The offset is " + Integer.toString(addressOffset) + " and the range for this opcode is [-32768, 32767].");
-			}
-			
-			$instructions.add(new Instruction20t(opcode, (short)addressOffset));
 		}
 	|	//e.g. sget_object v0 java/lang/System/out LJava/io/PrintStream;
 		^(I_STATEMENT_FORMAT21c_FIELD INSTRUCTION_FORMAT21c_FIELD REGISTER fully_qualified_field)

smali/src/main/java/org/jf/smali/main.java

@@ -18,6 +18,10 @@ package org.jf.smali;
 
 import org.apache.commons.cli.*;
 import org.jf.dexlib.DexFile;
+import org.jf.dexlib.CodeItem;
+import org.jf.dexlib.Code.InstructionIterator;
+import org.jf.dexlib.Code.Opcode;
+import org.jf.dexlib.Code.Format.Format;
 import org.jf.dexlib.Util.ByteArrayAnnotatedOutput;
 import org.jf.dexlib.Util.FileUtils;
 import org.antlr.runtime.ANTLRInputStream;
@@ -80,6 +84,8 @@ public class main {
 
         boolean sort = false;
         boolean rewriteLabels = false;
+        boolean fixStringConst = true;
+        boolean fixGoto = true;
 
         String outputDexFile = "out.dex";
         String dumpFileName = null;
@@ -129,6 +135,14 @@ public class main {
             }
         }
 
+        if (commandLine.hasOption("c")) {
+            fixStringConst = false;
+        }
+
+        if (commandLine.hasOption("g")) {
+            fixGoto = false;
+        }
+
 
         try {
 
@@ -174,6 +188,10 @@ public class main {
                 dexFile.setSortAllItems(true);
             }
 
+            if (fixStringConst || fixGoto) {
+                fixInstructions(dexFile, fixStringConst, fixGoto);
+            }
+
             dexFile.place();
 
             ByteArrayAnnotatedOutput out = new ByteArrayAnnotatedOutput();
@@ -222,6 +240,16 @@ public class main {
         }
     }
 
+    private static void fixInstructions(DexFile dexFile, boolean fixStringConst, boolean fixGoto) {
+        dexFile.place();
+
+        byte[] newInsns = null;
+
+        for (CodeItem codeItem: dexFile.CodeItemsSection.getItems()) {
+            codeItem.fixInstructions(fixStringConst, fixGoto);
+        }
+    }
+
     private static boolean doRewriteLabels(Set<File> files)
         throws Exception {
         boolean errorOccurred = false;
@@ -345,11 +373,21 @@ public class main {
                 .withDescription("rewrite the input smali files, converting any labels in the old (pre .97) format to the new format")
                 .create("r");
 
+        Option noFixStringConstOption = OptionBuilder.withLongOpt("no-fix-string-const")
+                .withDescription("Don't replace string-const instructions with string-const/jumbo where appropriate")
+                .create("c");
+
+        Option noFixGotoOption = OptionBuilder.withLongOpt("no-fix-goto")
+                .withDescription("Don't replace goto type instructions with a larger version where appropriate")
+                .create("g");
+
         options.addOption(versionOption);
         options.addOption(helpOption);
         options.addOption(dumpOption);
         options.addOption(outputOption);
         options.addOption(sortOption);
         options.addOption(rewriteLabelOption);
+        options.addOption(noFixStringConstOption);
+        options.addOption(noFixGotoOption);
     }
 }