Implement support for adding new labels to a MutableMethodImplementation

2 files changed, 161 insertions, 0 deletions

diff --git a/dexlib2/src/main/java/org/jf/dexlib2/builder/MutableMethodImplementation.java b/dexlib2/src/main/java/org/jf/dexlib2/builder/MutableMethodImplementation.java
index 84981221..b1e5dbbf 100644
--- a/dexlib2/src/main/java/org/jf/dexlib2/builder/MutableMethodImplementation.java
+++ b/dexlib2/src/main/java/org/jf/dexlib2/builder/MutableMethodImplementation.java
@@ -482,6 +482,49 @@ public class MutableMethodImplementation implements MethodImplementation {
         } while (true);
     }
 
+    private int mapCodeAddressToIndex(int codeAddress) {
+        float avgCodeUnitsPerInstruction = 1.9f;
+
+        int index = (int)(codeAddress/avgCodeUnitsPerInstruction);
+        if (index >= instructionList.size()) {
+            index = instructionList.size() - 1;
+        }
+
+        MethodLocation guessedLocation = instructionList.get(index);
+
+        if (guessedLocation.codeAddress == codeAddress) {
+            return index;
+        } else if (guessedLocation.codeAddress > codeAddress) {
+            do {
+                index--;
+            } while (instructionList.get(index).codeAddress > codeAddress);
+            return index;
+        } else {
+            do {
+                index++;
+            } while (index < instructionList.size() && instructionList.get(index).codeAddress <= codeAddress);
+            return index-1;
+        }
+    }
+
+    @Nonnull
+    public Label newLabelForAddress(int codeAddress) {
+        if (codeAddress < 0 || codeAddress > instructionList.get(instructionList.size()-1).codeAddress) {
+            throw new IndexOutOfBoundsException(String.format("codeAddress %d out of bounds", codeAddress));
+        }
+        MethodLocation referent = instructionList.get(mapCodeAddressToIndex(codeAddress));
+        return referent.addNewLabel();
+    }
+
+    @Nonnull
+    public Label newLabelForIndex(int instructionIndex) {
+        if (instructionIndex < 0 || instructionIndex >= instructionList.size()) {
+            throw new IndexOutOfBoundsException(String.format("instruction index %d out of bounds", instructionIndex));
+        }
+        MethodLocation referent = instructionList.get(instructionIndex);
+        return referent.addNewLabel();
+    }
+
     @Nonnull
     private Label newLabel(@Nonnull int[] codeAddressToIndex, int codeAddress) {
         MethodLocation referent = instructionList.get(mapCodeAddressToIndex(codeAddressToIndex, codeAddress));
diff --git a/dexlib2/src/test/java/org/jf/dexlib2/builder/MutableMethodImplementationTest.java b/dexlib2/src/test/java/org/jf/dexlib2/builder/MutableMethodImplementationTest.java
new file mode 100644
index 00000000..0df5ab3d
--- /dev/null
+++ b/dexlib2/src/test/java/org/jf/dexlib2/builder/MutableMethodImplementationTest.java
@@ -0,0 +1,118 @@
+/*
+ * Copyright 2015, Google Inc.
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *
+ * Redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer.
+ * Redistributions in binary form must reproduce the above
+ * copyright notice, this list of conditions and the following disclaimer
+ * in the documentation and/or other materials provided with the
+ * distribution.
+ * Neither the name of Google Inc. nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+package org.jf.dexlib2.builder;
+
+import org.jf.dexlib2.Opcode;
+import org.jf.dexlib2.builder.instruction.BuilderInstruction10x;
+import org.jf.dexlib2.builder.instruction.BuilderInstruction32x;
+import org.jf.dexlib2.iface.MethodImplementation;
+import org.junit.Assert;
+import org.junit.Test;
+
+public class MutableMethodImplementationTest {
+
+    @Test
+    public void testTryEndAtEndOfMethod() {
+        MethodImplementationBuilder builder = new MethodImplementationBuilder(10);
+
+        Label startLabel = builder.addLabel("start");
+        builder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+        builder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+        builder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+        builder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+        builder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+        builder.addInstruction(new BuilderInstruction32x(Opcode.MOVE_16, 0, 0));
+        Label endLabel = builder.addLabel("end");
+
+        builder.addCatch(startLabel, endLabel, startLabel);
+
+        MethodImplementation methodImplementation = builder.getMethodImplementation();
+
+        Assert.assertEquals(0, methodImplementation.getTryBlocks().get(0).getStartCodeAddress());
+        Assert.assertEquals(8, methodImplementation.getTryBlocks().get(0).getCodeUnitCount());
+
+        methodImplementation = new MutableMethodImplementation(methodImplementation);
+
+        Assert.assertEquals(0, methodImplementation.getTryBlocks().get(0).getStartCodeAddress());
+        Assert.assertEquals(8, methodImplementation.getTryBlocks().get(0).getCodeUnitCount());
+    }
+
+    @Test
+    public void testNewLabelByAddress() {
+        MethodImplementationBuilder builder = new MethodImplementationBuilder(10);
+
+        builder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+        builder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+        builder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+        builder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+        builder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+        builder.addInstruction(new BuilderInstruction32x(Opcode.MOVE_16, 0, 0));
+
+        MutableMethodImplementation mutableMethodImplementation =
+                new MutableMethodImplementation(builder.getMethodImplementation());
+
+        mutableMethodImplementation.addCatch(
+                mutableMethodImplementation.newLabelForAddress(0),
+                mutableMethodImplementation.newLabelForAddress(8),
+                mutableMethodImplementation.newLabelForAddress(1));
+
+        Assert.assertEquals(0, mutableMethodImplementation.getTryBlocks().get(0).getStartCodeAddress());
+        Assert.assertEquals(8, mutableMethodImplementation.getTryBlocks().get(0).getCodeUnitCount());
+        Assert.assertEquals(1, mutableMethodImplementation.getTryBlocks().get(0).getExceptionHandlers().get(0)
+                .getHandlerCodeAddress());
+    }
+
+    @Test
+    public void testNewLabelByIndex() {
+        MethodImplementationBuilder builder = new MethodImplementationBuilder(10);
+
+        builder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+        builder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+        builder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+        builder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+        builder.addInstruction(new BuilderInstruction10x(Opcode.NOP));
+        builder.addInstruction(new BuilderInstruction32x(Opcode.MOVE_16, 0, 0));
+
+        MutableMethodImplementation mutableMethodImplementation =
+                new MutableMethodImplementation(builder.getMethodImplementation());
+
+        mutableMethodImplementation.addCatch(
+                mutableMethodImplementation.newLabelForIndex(0),
+                mutableMethodImplementation.newLabelForIndex(6),
+                mutableMethodImplementation.newLabelForIndex(1));
+
+        Assert.assertEquals(0, mutableMethodImplementation.getTryBlocks().get(0).getStartCodeAddress());
+        Assert.assertEquals(8, mutableMethodImplementation.getTryBlocks().get(0).getCodeUnitCount());
+        Assert.assertEquals(1, mutableMethodImplementation.getTryBlocks().get(0).getExceptionHandlers().get(0)
+                .getHandlerCodeAddress());
+    }
+}