Implement just-in-time instruction fixup due to string jumboification

dexlib2/src/main/java/org/jf/dexlib2/builder/MethodLocation.java

@@ -54,16 +54,6 @@ public class MethodLocation {
         }
         debugItems.addAll(other.debugItems);
         other.debugItems = debugItems;
-
-        for (int i=debugItems.size()-1; i>=0; i--) {
-            BuilderDebugItem debugItem = debugItems.get(i);
-            debugItem.location = other;
-            other.debugItems.add(0, debugItem);
-        }
-        for (BuilderDebugItem debugItem: debugItems) {
-            debugItem.location = other;
-            other.debugItems.add(0, debugItem);
-        }
     }
 
     @Nonnull

dexlib2/src/main/java/org/jf/dexlib2/builder/MutableMethodImplementation.java

@@ -25,6 +25,10 @@ public class MutableMethodImplementation implements MethodImplementation {
     private final ArrayList<BuilderTryBlock> tryBlocks = Lists.newArrayList();
     private boolean fixInstructions = true;
 
+    public MutableMethodImplementation(@Nonnull MethodImplementation methodImplementation) {
+        throw new UnsupportedOperationException("not implemented yet.");
+    }
+
     public MutableMethodImplementation(int registerCount) {
         this.registerCount = registerCount;
     }
@@ -44,10 +48,16 @@ public class MutableMethodImplementation implements MethodImplementation {
                 if (i >= size()) {
                     throw new IndexOutOfBoundsException();
                 }
+                if (fixInstructions) {
+                    fixInstructions();
+                }
                 return instructionList.get(i).instruction;
             }
 
             @Override public int size() {
+                if (fixInstructions) {
+                    fixInstructions();
+                }
                 // don't include the last MethodLocation, which always has a null instruction
                 return instructionList.size() - 1;
             }
@@ -55,14 +65,24 @@ public class MutableMethodImplementation implements MethodImplementation {
     }
 
     @Nonnull @Override public List<? extends TryBlock<? extends ExceptionHandler>> getTryBlocks() {
+        if (fixInstructions) {
+            fixInstructions();
+        }
         return Collections.unmodifiableList(tryBlocks);
     }
 
     @Nonnull @Override public Iterable<? extends DebugItem> getDebugItems() {
+        if (fixInstructions) {
+            fixInstructions();
+        }
         return Iterables.concat(
                 Iterables.transform(instructionList, new Function<MethodLocation, Iterable<? extends DebugItem>>() {
                     @Nullable @Override public Iterable<? extends DebugItem> apply(@Nullable MethodLocation input) {
                         assert input != null;
+                        if (fixInstructions) {
+                            throw new IllegalStateException("This iterator was invalidated by a change to" +
+                                    " this MutableMethodImplementation.");
+                        }
                         return input.getDebugItems();
                     }
                 }));
@@ -107,7 +127,7 @@ public class MutableMethodImplementation implements MethodImplementation {
                 codeAddress += location.instruction.getCodeUnits();
             } else {
                 // only the last MethodLocation should have a null instruction
-                assert index == instructionList.size()-1;
+                assert i == instructionList.size()-1;
             }
         }
 
@@ -117,6 +137,7 @@ public class MutableMethodImplementation implements MethodImplementation {
     public void addInstruction(@Nonnull BuilderInstruction instruction) {
         MethodLocation last = instructionList.get(instructionList.size()-1);
         last.instruction = instruction;
+        instruction.location = last;
 
         int nextCodeAddress = last.codeAddress + instruction.getCodeUnits();
         instructionList.add(new MethodLocation(null, nextCodeAddress, instructionList.size()));
@@ -149,6 +170,8 @@ public class MutableMethodImplementation implements MethodImplementation {
                 assert i == instructionList.size() - 1;
             }
         }
+
+        this.fixInstructions = true;
     }
 
     public void removeInstruction(int index) {
@@ -164,7 +187,7 @@ public class MutableMethodImplementation implements MethodImplementation {
         instructionList.remove(index);
         int codeAddress = toRemove.codeAddress;
         for (int i=index; i<instructionList.size(); i++) {
-            MethodLocation location = instructionList.get(index);
+            MethodLocation location = instructionList.get(i);
             location.index = i;
             location.codeAddress = codeAddress;
 

dexlib2/src/main/java/org/jf/dexlib2/builder/instruction/BuilderInstruction3rc.java

@@ -36,7 +36,6 @@ import org.jf.dexlib2.Opcode;
 import org.jf.dexlib2.builder.BuilderInstruction;
 import org.jf.dexlib2.iface.instruction.formats.Instruction3rc;
 import org.jf.dexlib2.iface.reference.Reference;
-import org.jf.dexlib2.immutable.reference.ImmutableReferenceFactory;
 import org.jf.dexlib2.util.Preconditions;
 
 import javax.annotation.Nonnull;
@@ -56,7 +55,7 @@ public class BuilderInstruction3rc extends BuilderInstruction implements Instruc
         super(opcode);
         this.startRegister = Preconditions.checkShortRegister(startRegister);
         this.registerCount = Preconditions.checkRegisterRangeCount(registerCount);
-        this.reference = ImmutableReferenceFactory.of(opcode.referenceType, reference);
+        this.reference = reference;
     }
 
     @Override public int getStartRegister() { return startRegister; }

dexlib2/src/main/java/org/jf/dexlib2/builder/instruction/BuilderSwitchElement.java

@@ -24,6 +24,6 @@ public class BuilderSwitchElement implements SwitchElement {
     }
 
     @Override public int getOffset() {
-        return parent.getReferrer().getCodeAddress() - target.getCodeAddress();
+        return target.getCodeAddress() - parent.getReferrer().getCodeAddress();
     }
 }

dexlib2/src/main/java/org/jf/dexlib2/dexbacked/raw/util/DexAnnotator.java

@@ -53,6 +53,8 @@ public class DexAnnotator extends AnnotatedBytes {
 
     static {
         int[] sectionOrder = new int[] {
+                ItemType.MAP_LIST,
+
                 ItemType.HEADER_ITEM,
                 ItemType.STRING_ID_ITEM,
                 ItemType.TYPE_ID_ITEM,
@@ -66,7 +68,6 @@ public class DexAnnotator extends AnnotatedBytes {
                 ItemType.CODE_ITEM,
                 ItemType.DEBUG_INFO_ITEM,
 
-                ItemType.MAP_LIST,
                 ItemType.TYPE_LIST,
                 ItemType.ANNOTATION_SET_REF_LIST,
                 ItemType.ANNOTATION_SET_ITEM,

dexlib2/src/main/java/org/jf/dexlib2/writer/ClassSection.java

@@ -31,6 +31,7 @@
 
 package org.jf.dexlib2.writer;
 
+import org.jf.dexlib2.builder.MutableMethodImplementation;
 import org.jf.dexlib2.iface.ExceptionHandler;
 import org.jf.dexlib2.iface.TryBlock;
 import org.jf.dexlib2.iface.debug.DebugItem;
@@ -78,6 +79,7 @@ public interface ClassSection<StringKey extends CharSequence, TypeKey extends Ch
     @Nullable Iterable<? extends Instruction> getInstructions(@Nonnull MethodKey key);
     @Nonnull List<? extends TryBlock<? extends ExceptionHandler>> getTryBlocks(@Nonnull MethodKey key);
     @Nullable TypeKey getExceptionType(@Nonnull ExceptionHandler handler);
+    @Nonnull MutableMethodImplementation makeMutableMethodImplementation(@Nonnull MethodKey key);
 
     void setEncodedArrayOffset(@Nonnull ClassKey key, int offset);
     int getEncodedArrayOffset(@Nonnull ClassKey key);

dexlib2/src/main/java/org/jf/dexlib2/writer/DexWriter.java

@@ -31,10 +31,16 @@
 
 package org.jf.dexlib2.writer;
 
-import com.google.common.collect.*;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
+import com.google.common.collect.Ordering;
 import org.jf.dexlib2.AccessFlags;
+import org.jf.dexlib2.Opcode;
 import org.jf.dexlib2.ReferenceType;
 import org.jf.dexlib2.base.BaseAnnotation;
+import org.jf.dexlib2.builder.MutableMethodImplementation;
+import org.jf.dexlib2.builder.instruction.BuilderInstruction31c;
 import org.jf.dexlib2.dexbacked.raw.*;
 import org.jf.dexlib2.iface.Annotation;
 import org.jf.dexlib2.iface.ExceptionHandler;
@@ -42,6 +48,7 @@ import org.jf.dexlib2.iface.TryBlock;
 import org.jf.dexlib2.iface.debug.DebugItem;
 import org.jf.dexlib2.iface.debug.LineNumber;
 import org.jf.dexlib2.iface.instruction.Instruction;
+import org.jf.dexlib2.iface.instruction.OneRegisterInstruction;
 import org.jf.dexlib2.iface.instruction.ReferenceInstruction;
 import org.jf.dexlib2.iface.instruction.formats.*;
 import org.jf.dexlib2.iface.reference.*;
@@ -727,10 +734,39 @@ public abstract class DexWriter<
             Iterable<MethodKey> methods = Iterables.concat(directMethods, virtualMethods);
 
             for (MethodKey methodKey: methods) {
-                int debugItemOffset = writeDebugItem(offsetWriter, debugWriter, methodKey);
+                List<? extends TryBlock<? extends ExceptionHandler>> tryBlocks =
-                int codeItemOffset = writeCodeItem(codeWriter, ehBuf, methodKey, debugItemOffset);
+                        classSection.getTryBlocks(methodKey);
+                Iterable<? extends Instruction> instructions = classSection.getInstructions(methodKey);
+                Iterable<? extends DebugItem> debugItems = classSection.getDebugItems(methodKey);
 
-                if (codeItemOffset != NO_OFFSET) {
+                if (instructions != null && stringSection.hasJumboIndexes()) {
+                    boolean needsFix = false;
+                    for (Instruction instruction: instructions) {
+                        if (instruction.getOpcode() == Opcode.CONST_STRING) {
+                            if (stringSection.getItemIndex(
+                                    (StringRef)((ReferenceInstruction)instruction).getReference()) >= 65536) {
+                                needsFix = true;
+                                break;
+                            }
+                        }
+                    }
+
+                    if (needsFix) {
+                        MutableMethodImplementation mutableMethodImplementation =
+                                classSection.makeMutableMethodImplementation(methodKey);
+                        fixInstructions(mutableMethodImplementation);
+
+                        instructions = mutableMethodImplementation.getInstructions();
+                        tryBlocks = mutableMethodImplementation.getTryBlocks();
+                        debugItems = mutableMethodImplementation.getDebugItems();
+                    }
+                }
+
+                int debugItemOffset = writeDebugItem(offsetWriter, debugWriter,
+                        classSection.getParameterNames(methodKey), debugItems);
+                int codeItemOffset = writeCodeItem(codeWriter, ehBuf, methodKey, tryBlocks, instructions, debugItemOffset);
+
+                if (codeItemOffset != -1) {
                     codeOffsets.add(new CodeItemOffset<MethodKey>(methodKey, codeItemOffset));
                 }
             }
@@ -739,6 +775,7 @@ public abstract class DexWriter<
         offsetWriter.align();
         codeSectionOffset = offsetWriter.getPosition();
 
+        codeWriter.close();
         temp.writeTo(offsetWriter);
         temp.close();
 
@@ -747,12 +784,27 @@ public abstract class DexWriter<
         }
     }
 
-    private int writeDebugItem(@Nonnull DexDataWriter writer,
+    private void fixInstructions(@Nonnull MutableMethodImplementation methodImplementation) {
-                                @Nonnull DebugWriter<StringKey, TypeKey> debugWriter,
+        List<Instruction> instructions = methodImplementation.getInstructions();
-                                @Nonnull MethodKey methodKey) throws IOException {
-        Iterable<? extends DebugItem> debugItems = classSection.getDebugItems(methodKey);
-        Iterable<? extends StringKey> parameterNames = classSection.getParameterNames(methodKey);
 
+        for (int i=0; i<instructions.size(); i++) {
+            Instruction instruction = instructions.get(i);
+
+            if (instruction.getOpcode() == Opcode.CONST_STRING) {
+                if (stringSection.getItemIndex(
+                        (StringRef)((ReferenceInstruction)instruction).getReference()) > 65536) {
+                    methodImplementation.replaceInstruction(i, new BuilderInstruction31c(Opcode.CONST_STRING_JUMBO,
+                            ((OneRegisterInstruction)instruction).getRegisterA(),
+                            ((ReferenceInstruction)instruction).getReference()));
+                }
+            }
+        }
+    }
+
+    private int writeDebugItem(@Nonnull DexDataWriter writer,
+                               @Nonnull DebugWriter<StringKey, TypeKey> debugWriter,
+                               @Nullable Iterable<? extends StringKey> parameterNames,
+                               @Nullable Iterable<? extends DebugItem> debugItems) throws IOException {
         int parameterCount = 0;
         if (parameterNames != null) {
             int index = 0;
@@ -809,13 +861,13 @@ public abstract class DexWriter<
     }
 
     private int writeCodeItem(@Nonnull DexDataWriter writer,
-                               @Nonnull ByteArrayOutputStream ehBuf,
+                              @Nonnull ByteArrayOutputStream ehBuf,
-                               @Nonnull MethodKey methodKey,
+                              @Nonnull MethodKey methodKey,
-                               int debugItemOffset) throws IOException {
+                              @Nonnull List<? extends TryBlock<? extends ExceptionHandler>> tryBlocks,
-        Iterable<? extends Instruction> instructions = classSection.getInstructions(methodKey);
+                              @Nullable Iterable<? extends Instruction> instructions,
-
+                              int debugItemOffset) throws IOException {
         if (instructions == null && debugItemOffset == NO_OFFSET) {
-            return NO_OFFSET;
+            return -1;
         }
 
         numCodeItemItems++;
@@ -823,7 +875,6 @@ public abstract class DexWriter<
         writer.align();
 
         int codeItemOffset = writer.getPosition();
-        classSection.setCodeItemOffset(methodKey, writer.getPosition());
 
         writer.writeUshort(classSection.getRegisterCount(methodKey));
 
@@ -831,7 +882,6 @@ public abstract class DexWriter<
         Collection<? extends TypeKey> parameters = typeListSection.getTypes(
                 protoSection.getParameters(methodSection.getPrototype(methodKey)));
 
-        List<? extends TryBlock<? extends ExceptionHandler>> tryBlocks = classSection.getTryBlocks(methodKey);
         writer.writeUshort(MethodUtil.getParameterRegisterCount(parameters, isStatic));
 
         if (instructions != null) {

dexlib2/src/main/java/org/jf/dexlib2/writer/StringSection.java

@@ -37,4 +37,5 @@ import javax.annotation.Nonnull;
 
 public interface StringSection<StringKey, StringRef extends StringReference> extends NullableIndexSection<StringKey> {
     int getItemIndex(@Nonnull StringRef key);
+    boolean hasJumboIndexes();
 }

dexlib2/src/main/java/org/jf/dexlib2/writer/builder/BuilderClassPool.java

@@ -35,6 +35,7 @@ import com.google.common.base.Function;
 import com.google.common.base.Predicate;
 import com.google.common.collect.*;
 import org.jf.dexlib2.DebugItemType;
+import org.jf.dexlib2.builder.MutableMethodImplementation;
 import org.jf.dexlib2.iface.ExceptionHandler;
 import org.jf.dexlib2.iface.Field;
 import org.jf.dexlib2.iface.MethodImplementation;
@@ -313,6 +314,15 @@ public class BuilderClassPool implements ClassSection<BuilderStringReference, Bu
         return checkTypeReference(handler.getExceptionTypeReference());
     }
 
+    @Nonnull @Override
+    public MutableMethodImplementation makeMutableMethodImplementation(@Nonnull BuilderMethod builderMethod) {
+        MethodImplementation impl = builderMethod.getImplementation();
+        if (impl instanceof MutableMethodImplementation) {
+            return (MutableMethodImplementation)impl;
+        }
+        return new MutableMethodImplementation(impl);
+    }
+
     @Override public void setEncodedArrayOffset(@Nonnull BuilderClassDef builderClassDef, int offset) {
         builderClassDef.encodedArrayOffset = offset;
     }

dexlib2/src/main/java/org/jf/dexlib2/writer/builder/BuilderStringPool.java

@@ -69,6 +69,10 @@ class BuilderStringPool implements StringSection<BuilderStringReference, Builder
         return key.index;
     }
 
+    @Override public boolean hasJumboIndexes() {
+        return internedItems.size() > 65536;
+    }
+
     @Nonnull @Override public Collection<? extends Entry<? extends BuilderStringReference, Integer>> getItems() {
         return new BuilderMapEntryCollection<BuilderStringReference>(internedItems.values()) {
             @Override protected int getValue(@Nonnull BuilderStringReference key) {

dexlib2/src/main/java/org/jf/dexlib2/writer/io/MemoryDeferredOutputStream.java

@@ -53,7 +53,7 @@ public class MemoryDeferredOutputStream extends DeferredOutputStream {
         int remaining = remaining();
         int written = 0;
         while (length - written > 0) {
-            int toWrite = Math.min(remaining, length);
+            int toWrite = Math.min(remaining, (length - written));
             System.arraycopy(bytes, offset + written, currentBuffer, currentPosition, toWrite);
             written += toWrite;
             currentPosition += toWrite;

dexlib2/src/main/java/org/jf/dexlib2/writer/pool/ClassPool.java

@@ -36,6 +36,7 @@ import com.google.common.base.Predicate;
 import com.google.common.collect.*;
 import org.jf.dexlib2.DebugItemType;
 import org.jf.dexlib2.ReferenceType;
+import org.jf.dexlib2.builder.MutableMethodImplementation;
 import org.jf.dexlib2.iface.*;
 import org.jf.dexlib2.iface.debug.*;
 import org.jf.dexlib2.iface.instruction.Instruction;
@@ -433,6 +434,11 @@ public class ClassPool implements ClassSection<CharSequence, CharSequence,
         return handler.getExceptionType();
     }
 
+    @Nonnull @Override
+    public MutableMethodImplementation makeMutableMethodImplementation(@Nonnull PoolMethod poolMethod) {
+        return new MutableMethodImplementation(poolMethod.getImplementation());
+    }
+
     @Override public void setEncodedArrayOffset(@Nonnull PoolClassDef classDef, int offset) {
         classDef.encodedArrayOffset = offset;
     }

dexlib2/src/main/java/org/jf/dexlib2/writer/pool/StringPool.java

@@ -56,4 +56,8 @@ public class StringPool extends StringTypeBasePool implements StringSection<Char
         }
         return index;
     }
+
+    @Override public boolean hasJumboIndexes() {
+        return internedItems.size() > 65536;
+    }
 }

dexlib2/src/test/java/org/jf/dexlib2/writer/PayloadAlignmentTest.java

@@ -31,85 +31,224 @@
 
 package org.jf.dexlib2.writer;
 
-// TODO: uncomment/reimplement
+import com.google.common.collect.Lists;
-/*public class PayloadAlignmentTest {
+import org.jf.dexlib2.Opcode;
-    private MockStringIndexProvider mockStringIndexProvider;
+import org.jf.dexlib2.builder.Label;
+import org.jf.dexlib2.builder.MethodImplementationBuilder;
+import org.jf.dexlib2.builder.SwitchLabelElement;
+import org.jf.dexlib2.builder.instruction.*;
+import org.jf.dexlib2.iface.instruction.Instruction;
+import org.jf.dexlib2.iface.instruction.OffsetInstruction;
+import org.jf.dexlib2.iface.instruction.formats.Instruction31t;
+import org.jf.dexlib2.iface.instruction.formats.PackedSwitchPayload;
+import org.jf.dexlib2.iface.instruction.formats.SparseSwitchPayload;
+import org.junit.Assert;
+import org.junit.Test;
 
-    private class InsnWriteUtil extends InstructionWriteUtil<StringReference, Reference> {
+import java.util.List;
-        public InsnWriteUtil(@Nonnull MethodImplementation implementation) {
-            super(implementation.getInstructions(), mockStringIndexProvider, ImmutableInstructionFactory.INSTANCE);
-        }
-    }
 
-    @Before
+public class PayloadAlignmentTest {
-    public void setup() {
+
-        mockStringIndexProvider = new MockStringIndexProvider();
+    @Test
+    public void testPayloadAlignmentRemoveNop() {
+        MethodImplementationBuilder implBuilder = new MethodImplementationBuilder(10);
+
+        implBuilder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+        implBuilder.addInstruction(new BuilderArrayPayload(4, null));
+
+        List<? extends Instruction> instructions =
+                Lists.newArrayList(implBuilder.getMethodImplementation().getInstructions());
+
+        Assert.assertEquals(instructions.size(), 1);
+
+        Instruction instruction = instructions.get(0);
+
+        Assert.assertEquals(instruction.getOpcode(), Opcode.ARRAY_PAYLOAD);
     }
 
     @Test
-    public void testArrayPayloadAlignment() {
+    public void testPayloadAlignmentAddNop() {
-        ArrayList<ImmutableInstruction> instructions = Lists.newArrayList();
+        MethodImplementationBuilder implBuilder = new MethodImplementationBuilder(10);
 
-        // add misaligned array payload
+        implBuilder.addInstruction(new BuilderInstruction12x(Opcode.MOVE, 0, 0));
-        instructions.add(new ImmutableInstruction10x(Opcode.NOP));
+        implBuilder.addInstruction(new BuilderArrayPayload(4, null));
-        instructions.add(new ImmutableArrayPayload(4, null));
 
-        ImmutableMethodImplementation methodImplementation = new ImmutableMethodImplementation(1, instructions, null, null);
+        List<? extends Instruction> instructions =
-        InsnWriteUtil writeUtil = new InsnWriteUtil(methodImplementation);
+                Lists.newArrayList(implBuilder.getMethodImplementation().getInstructions());
 
-        int codeOffset = 0;
+        Assert.assertEquals(instructions.size(), 3);
-        for (Instruction instr: writeUtil.getInstructions()) {
+
-            if (instr.getOpcode().equals(Opcode.ARRAY_PAYLOAD)) {
+        Instruction instruction = instructions.get(0);
-                Assert.assertEquals("array payload was not aligned properly", codeOffset%2, 0);
+        Assert.assertEquals(instruction.getOpcode(), Opcode.MOVE);
-                break;
+
-            }
+        instruction = instructions.get(1);
-            codeOffset += instr.getCodeUnits();
+        Assert.assertEquals(instruction.getOpcode(), Opcode.NOP);
+
+        instruction = instructions.get(2);
+        Assert.assertEquals(instruction.getOpcode(), Opcode.ARRAY_PAYLOAD);
+    }
+
+    @Test
+    public void testPayloadAlignmentRemoveNopWithReferent() {
+        MethodImplementationBuilder implBuilder = new MethodImplementationBuilder(10);
+
+        Label label = implBuilder.getLabel("array_payload");
+        implBuilder.addInstruction(new BuilderInstruction31t(Opcode.FILL_ARRAY_DATA, 0, label));
+        implBuilder.addInstruction(new BuilderInstruction12x(Opcode.MOVE, 0, 0));
+        implBuilder.addInstruction(new BuilderInstruction12x(Opcode.MOVE, 0, 0));
+        implBuilder.addInstruction(new BuilderInstruction12x(Opcode.MOVE, 0, 0));
+        implBuilder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+        implBuilder.addLabel("array_payload");
+        implBuilder.addInstruction(new BuilderArrayPayload(4, null));
+
+        List<Instruction> instructions = Lists.newArrayList(implBuilder.getMethodImplementation().getInstructions());
+
+        checkInstructions(instructions,
+                new Opcode[]{Opcode.FILL_ARRAY_DATA,
+                        Opcode.MOVE,
+                        Opcode.MOVE,
+                        Opcode.MOVE,
+                        Opcode.ARRAY_PAYLOAD});
+
+        Instruction31t referent = (Instruction31t)instructions.get(0);
+        Assert.assertEquals(6, referent.getCodeOffset());
+    }
+
+    @Test
+    public void testPayloadAlignmentAddNopWithReferent() {
+        MethodImplementationBuilder implBuilder = new MethodImplementationBuilder(10);
+
+        Label label = implBuilder.getLabel("array_payload");
+        implBuilder.addInstruction(new BuilderInstruction31t(Opcode.FILL_ARRAY_DATA, 0, label));
+        implBuilder.addInstruction(new BuilderInstruction12x(Opcode.MOVE, 0, 0));
+        implBuilder.addInstruction(new BuilderInstruction12x(Opcode.MOVE, 0, 0));
+        implBuilder.addInstruction(new BuilderInstruction12x(Opcode.MOVE, 0, 0));
+        implBuilder.addInstruction(new BuilderInstruction12x(Opcode.MOVE, 0, 0));
+        implBuilder.addLabel("array_payload");
+        implBuilder.addInstruction(new BuilderArrayPayload(4, null));
+
+        List<Instruction> instructions = Lists.newArrayList(implBuilder.getMethodImplementation().getInstructions());
+
+        checkInstructions(instructions,
+                new Opcode[]{Opcode.FILL_ARRAY_DATA,
+                    Opcode.MOVE,
+                    Opcode.MOVE,
+                    Opcode.MOVE,
+                    Opcode.MOVE,
+                    Opcode.NOP,
+                    Opcode.ARRAY_PAYLOAD});
+
+        Instruction31t referent = (Instruction31t)instructions.get(0);
+        Assert.assertEquals(8, referent.getCodeOffset());
+    }
+
+    private static void checkInstructions(List<Instruction> instructions, Opcode[] expectedOpcodes) {
+        Assert.assertEquals(expectedOpcodes.length, instructions.size());
+
+        for (int i=0; i<expectedOpcodes.length; i++) {
+            Assert.assertEquals(instructions.get(i).getOpcode(), expectedOpcodes[i]);
         }
     }
 
     @Test
     public void testPackedSwitchAlignment() {
-        ArrayList<ImmutableInstruction> instructions = Lists.newArrayList();
+        MethodImplementationBuilder implBuilder = new MethodImplementationBuilder(10);
-        // add misaligned packed switch payload
-        ArrayList<SwitchElement> switchElements = Lists.newArrayList();
-        switchElements.add(new ImmutableSwitchElement(0, 5));
-        instructions.add(new ImmutableInstruction10x(Opcode.NOP));
-        instructions.add(new ImmutablePackedSwitchPayload(switchElements));
 
-        ImmutableMethodImplementation methodImplementation = new ImmutableMethodImplementation(1, instructions, null, null);
+        implBuilder.addLabel("switch_target_1");
-        InsnWriteUtil writeUtil = new InsnWriteUtil(methodImplementation);
+        implBuilder.addInstruction(new BuilderInstruction10t(Opcode.GOTO, implBuilder.getLabel("goto_target")));
 
-        int codeOffset = 0;
+        implBuilder.addLabel("switch_payload");
-        for (Instruction instr: writeUtil.getInstructions()) {
+        implBuilder.addInstruction(new BuilderPackedSwitchPayload(0, Lists.newArrayList(
-            if (instr.getOpcode().equals(Opcode.PACKED_SWITCH_PAYLOAD)) {
+                implBuilder.getLabel("switch_target_1"),
-                Assert.assertEquals("packed switch payload was not aligned properly", codeOffset%2, 0);
+                implBuilder.getLabel("switch_target_2"),
-                break;
+                implBuilder.getLabel("switch_target_3"))));
-            }
+
-            codeOffset += instr.getCodeUnits();
+        implBuilder.addLabel("goto_target");
-        }
+        implBuilder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+        implBuilder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+
+        implBuilder.addLabel("switch_target_2");
+        implBuilder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+
+        implBuilder.addLabel("switch_target_3");
+        implBuilder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+
+        implBuilder.addInstruction(new BuilderInstruction31t(Opcode.PACKED_SWITCH, 0,
+                implBuilder.getLabel("switch_payload")));
+
+        List<Instruction> instructions = Lists.newArrayList(implBuilder.getMethodImplementation().getInstructions());
+
+        checkInstructions(instructions,
+                new Opcode[]{Opcode.GOTO,
+                        Opcode.NOP,
+                        Opcode.PACKED_SWITCH_PAYLOAD,
+                        Opcode.NOP,
+                        Opcode.NOP,
+                        Opcode.NOP,
+                        Opcode.NOP,
+                        Opcode.PACKED_SWITCH});
+
+        OffsetInstruction gotoInstruction = (OffsetInstruction)instructions.get(0);
+        Assert.assertEquals(12, gotoInstruction.getCodeOffset());
+
+        PackedSwitchPayload payload = (PackedSwitchPayload)instructions.get(2);
+        Assert.assertEquals(3, payload.getSwitchElements().size());
+        Assert.assertEquals(-16, payload.getSwitchElements().get(0).getOffset());
+        Assert.assertEquals(-2, payload.getSwitchElements().get(1).getOffset());
+        Assert.assertEquals(-1, payload.getSwitchElements().get(2).getOffset());
+
+        OffsetInstruction referent = (OffsetInstruction)instructions.get(7);
+        Assert.assertEquals(-14, referent.getCodeOffset());
     }
 
     @Test
     public void testSparseSwitchAlignment() {
-        ArrayList<ImmutableInstruction> instructions = Lists.newArrayList();
+        MethodImplementationBuilder implBuilder = new MethodImplementationBuilder(10);
 
-        // add misaligned sparse switch payload
+        implBuilder.addLabel("switch_target_1");
-        ArrayList<SwitchElement> switchElements = Lists.newArrayList();
+        implBuilder.addInstruction(new BuilderInstruction10t(Opcode.GOTO, implBuilder.getLabel("goto_target")));
-        switchElements.add(new ImmutableSwitchElement(0, 5));
 
-        instructions.add(new ImmutableInstruction10x(Opcode.NOP));
+        implBuilder.addLabel("switch_payload");
-        instructions.add(new ImmutableSparseSwitchPayload(switchElements));
+        implBuilder.addInstruction(new BuilderSparseSwitchPayload(Lists.newArrayList(
+                new SwitchLabelElement(0, implBuilder.getLabel("switch_target_1")),
+                new SwitchLabelElement(5, implBuilder.getLabel("switch_target_2")),
+                new SwitchLabelElement(10, implBuilder.getLabel("switch_target_3")))));
 
-        ImmutableMethodImplementation methodImplementation = new ImmutableMethodImplementation(1, instructions, null, null);
+        implBuilder.addLabel("goto_target");
-        InsnWriteUtil writeUtil = new InsnWriteUtil(methodImplementation);
+        implBuilder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+        implBuilder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
 
-        int codeOffset = 0;
+        implBuilder.addLabel("switch_target_2");
-        for (Instruction instr: writeUtil.getInstructions()) {
+        implBuilder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
-            if (instr.getOpcode().equals(Opcode.SPARSE_SWITCH_PAYLOAD)) {
+
-                Assert.assertEquals("packed switch payload was not aligned properly", codeOffset%2, 0);
+        implBuilder.addLabel("switch_target_3");
-                break;
+        implBuilder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
-            }
+
-            codeOffset += instr.getCodeUnits();
+        implBuilder.addInstruction(new BuilderInstruction31t(Opcode.SPARSE_SWITCH, 0,
-        }
+                implBuilder.getLabel("switch_payload")));
+
+        List<Instruction> instructions = Lists.newArrayList(implBuilder.getMethodImplementation().getInstructions());
+
+        checkInstructions(instructions,
+                new Opcode[]{Opcode.GOTO,
+                        Opcode.NOP,
+                        Opcode.SPARSE_SWITCH_PAYLOAD,
+                        Opcode.NOP,
+                        Opcode.NOP,
+                        Opcode.NOP,
+                        Opcode.NOP,
+                        Opcode.SPARSE_SWITCH});
+
+        OffsetInstruction gotoInstruction = (OffsetInstruction)instructions.get(0);
+        Assert.assertEquals(16, gotoInstruction.getCodeOffset());
+
+        SparseSwitchPayload payload = (SparseSwitchPayload)instructions.get(2);
+        Assert.assertEquals(3, payload.getSwitchElements().size());
+        Assert.assertEquals(-20, payload.getSwitchElements().get(0).getOffset());
+        Assert.assertEquals(-2, payload.getSwitchElements().get(1).getOffset());
+        Assert.assertEquals(-1, payload.getSwitchElements().get(2).getOffset());
+
+        OffsetInstruction referent = (OffsetInstruction)instructions.get(7);
+        Assert.assertEquals(-18, referent.getCodeOffset());
     }
-}*/
+}

smali/src/main/antlr3/smaliTreeWalker.g

@@ -538,7 +538,8 @@ catches returns[List<BuilderTryBlock> tryBlocks]
 catch_directive
   : ^(I_CATCH nonvoid_type_descriptor from=label_ref to=label_ref using=label_ref)
   {
-    $method::methodBuilder.addCatch($nonvoid_type_descriptor.type, $from.label, $to.label, $using.label);
+    $method::methodBuilder.addCatch(dexBuilder.internTypeReference($nonvoid_type_descriptor.type),
+        $from.label, $to.label, $using.label);
   };
 
 catchall_directive