Fix an ugly bug where try/catch offsets weren't being mapped properly.

In dex files, there are several places where one object refers
to another object by its position in the file. The dex merger
is generally very careful to adjust these mappings when combining
dex files.

Unfortunately one of these cases was broken. Each try_item refers
to a corresponding encoded_catch_handler by its byte offset in a
list. Most of the time this byte offset is the same in the input
dex file and the output dex file. But encoded_catch_handlers are
variable-length because they use a variable-length uleb128 encoding
to address the type_idx being caught. When dex files are merged,
some exception types may go from having a small index to having
a large index, increasing the number of bytes required to encode
that index. This breaks our ability to directly copy over offsets
as we were doing previously.

Bug: http://code.google.com/p/android/issues/detail?id=36490
Change-Id: I3bdadf20899fdb5d4d074e69103b33c0404a31f8

4 files changed, 116 insertions, 32 deletions

diff --git a/dx/src/com/android/dx/dex/SizeOf.java b/dx/src/com/android/dx/dex/SizeOf.java
index 476f7bbcd..6ded78282 100644
--- a/dx/src/com/android/dx/dex/SizeOf.java
+++ b/dx/src/com/android/dx/dex/SizeOf.java
@@ -100,4 +100,11 @@ public final class SizeOf {
      * offset uint
      */
     public static final int MAP_ITEM = USHORT + USHORT + UINT + UINT;
+
+    /**
+     * start_addr uint
+     * insn_count ushort
+     * handler_off ushort
+     */
+    public static final int TRY_ITEM = UINT + USHORT + USHORT;
 }
diff --git a/dx/src/com/android/dx/io/Code.java b/dx/src/com/android/dx/io/Code.java
index ba95d1b92..82da86267 100644
--- a/dx/src/com/android/dx/io/Code.java
+++ b/dx/src/com/android/dx/io/Code.java
@@ -67,12 +67,12 @@ public final class Code {
     public static class Try {
         final int startAddress;
         final int instructionCount;
-        final int handlerOffset;
+        final int catchHandlerIndex;
 
-        Try(int startAddress, int instructionCount, int handlerOffset) {
+        Try(int startAddress, int instructionCount, int catchHandlerIndex) {
             this.startAddress = startAddress;
             this.instructionCount = instructionCount;
-            this.handlerOffset = handlerOffset;
+            this.catchHandlerIndex = catchHandlerIndex;
         }
 
         public int getStartAddress() {
@@ -83,8 +83,12 @@ public final class Code {
             return instructionCount;
         }
 
-        public int getHandlerOffset() {
-            return handlerOffset;
+        /**
+         * Returns this try's catch handler <strong>index</strong>. Note that
+         * this is distinct from the its catch handler <strong>offset</strong>.
+         */
+        public int getCatchHandlerIndex() {
+            return catchHandlerIndex;
         }
     }
 
@@ -92,11 +96,13 @@ public final class Code {
         final int[] typeIndexes;
         final int[] addresses;
         final int catchAllAddress;
+        final int offset;
 
-        public CatchHandler(int[] typeIndexes, int[] addresses, int catchAllAddress) {
+        public CatchHandler(int[] typeIndexes, int[] addresses, int catchAllAddress, int offset) {
             this.typeIndexes = typeIndexes;
             this.addresses = addresses;
             this.catchAllAddress = catchAllAddress;
+            this.offset = offset;
         }
 
         public int[] getTypeIndexes() {
@@ -110,5 +116,9 @@ public final class Code {
         public int getCatchAllAddress() {
             return catchAllAddress;
         }
+
+        public int getOffset() {
+            return offset;
+        }
     }
 }
diff --git a/dx/src/com/android/dx/io/DexBuffer.java b/dx/src/com/android/dx/io/DexBuffer.java
index 39e5858e6..9fbc78cac 100644
--- a/dx/src/com/android/dx/io/DexBuffer.java
+++ b/dx/src/com/android/dx/io/DexBuffer.java
@@ -19,6 +19,8 @@ package com.android.dx.io;
 import com.android.dx.dex.DexFormat;
 import com.android.dx.dex.SizeOf;
 import com.android.dx.dex.TableOfContents;
+import com.android.dx.io.Code.CatchHandler;
+import com.android.dx.io.Code.Try;
 import com.android.dx.merge.TypeList;
 import com.android.dx.util.ByteInput;
 import com.android.dx.util.ByteOutput;
@@ -443,31 +445,64 @@ public final class DexBuffer {
             int debugInfoOffset = readInt();
             int instructionsSize = readInt();
             short[] instructions = readShortArray(instructionsSize);
-            Code.Try[] tries = new Code.Try[triesSize];
-            Code.CatchHandler[] catchHandlers = new Code.CatchHandler[0];
+            Try[] tries;
+            CatchHandler[] catchHandlers;
             if (triesSize > 0) {
                 if (instructions.length % 2 == 1) {
                     readShort(); // padding
                 }
 
-                for (int i = 0; i < triesSize; i++) {
-                    int startAddress = readInt();
-                    int instructionCount = readUnsignedShort();
-                    int handlerOffset = readUnsignedShort();
-                    tries[i] = new Code.Try(startAddress, instructionCount, handlerOffset);
-                }
-
-                int catchHandlersSize = readUleb128();
-                catchHandlers = new Code.CatchHandler[catchHandlersSize];
-                for (int i = 0; i < catchHandlersSize; i++) {
-                    catchHandlers[i] = readCatchHandler();
-                }
+                /*
+                 * We can't read the tries until we've read the catch handlers.
+                 * Unfortunately they're in the opposite order in the dex file
+                 * so we need to read them out-of-order.
+                 */
+                Section triesSection = open(position);
+                skip(triesSize * SizeOf.TRY_ITEM);
+                catchHandlers = readCatchHandlers();
+                tries = triesSection.readTries(triesSize, catchHandlers);
+            } else {
+                tries = new Try[0];
+                catchHandlers = new CatchHandler[0];
             }
             return new Code(registersSize, insSize, outsSize, debugInfoOffset, instructions,
                     tries, catchHandlers);
         }
 
-        private Code.CatchHandler readCatchHandler() {
+        private CatchHandler[] readCatchHandlers() {
+            int baseOffset = position;
+            int catchHandlersSize = readUleb128();
+            CatchHandler[] result = new CatchHandler[catchHandlersSize];
+            for (int i = 0; i < catchHandlersSize; i++) {
+                int offset = position - baseOffset;
+                result[i] = readCatchHandler(offset);
+            }
+            return result;
+        }
+
+        private Try[] readTries(int triesSize, CatchHandler[] catchHandlers) {
+            Try[] result = new Try[triesSize];
+            for (int i = 0; i < triesSize; i++) {
+                int startAddress = readInt();
+                int instructionCount = readUnsignedShort();
+                int handlerOffset = readUnsignedShort();
+                int catchHandlerIndex = findCatchHandlerIndex(catchHandlers, handlerOffset);
+                result[i] = new Try(startAddress, instructionCount, catchHandlerIndex);
+            }
+            return result;
+        }
+
+        private int findCatchHandlerIndex(CatchHandler[] catchHandlers, int offset) {
+            for (int i = 0; i < catchHandlers.length; i++) {
+                CatchHandler catchHandler = catchHandlers[i];
+                if (catchHandler.getOffset() == offset) {
+                    return i;
+                }
+            }
+            throw new IllegalArgumentException();
+        }
+
+        private CatchHandler readCatchHandler(int offset) {
             int size = readSleb128();
             int handlersCount = Math.abs(size);
             int[] typeIndexes = new int[handlersCount];
@@ -477,7 +512,7 @@ public final class DexBuffer {
                 addresses[i] = readUleb128();
             }
             int catchAllAddress = size <= 0 ? readUleb128() : -1;
-            return new Code.CatchHandler(typeIndexes, addresses, catchAllAddress);
+            return new CatchHandler(typeIndexes, addresses, catchAllAddress, offset);
         }
 
         private ClassData readClassData() {
@@ -548,6 +583,14 @@ public final class DexBuffer {
             }
         }
 
+        public void skip(int count) {
+            if (count < 0) {
+                throw new IllegalArgumentException();
+            }
+            ensureCapacity(count);
+            position += count;
+        }
+
         /**
          * Writes 0x00 until the position is aligned to a multiple of 4.
          */
diff --git a/dx/src/com/android/dx/merge/DexMerger.java b/dx/src/com/android/dx/merge/DexMerger.java
index c48f436f4..87a59fcae 100644
--- a/dx/src/com/android/dx/merge/DexMerger.java
+++ b/dx/src/com/android/dx/merge/DexMerger.java
@@ -790,6 +790,7 @@ public final class DexMerger {
         codeOut.writeUnsignedShort(code.getOutsSize());
 
         Code.Try[] tries = code.getTries();
+        Code.CatchHandler[] catchHandlers = code.getCatchHandlers();
         codeOut.writeUnsignedShort(tries.length);
 
         int debugInfoOffset = code.getDebugInfoOffset();
@@ -812,16 +813,39 @@ public final class DexMerger {
             if (newInstructions.length % 2 == 1) {
                 codeOut.writeShort((short) 0); // padding
             }
-            for (Code.Try tryItem : tries) {
-                codeOut.writeInt(tryItem.getStartAddress());
-                codeOut.writeUnsignedShort(tryItem.getInstructionCount());
-                codeOut.writeUnsignedShort(tryItem.getHandlerOffset());
-            }
-            Code.CatchHandler[] catchHandlers = code.getCatchHandlers();
-            codeOut.writeUleb128(catchHandlers.length);
-            for (Code.CatchHandler catchHandler : catchHandlers) {
-                transformEncodedCatchHandler(catchHandler, indexMap);
-            }
+
+            /*
+             * We can't write the tries until we've written the catch handlers.
+             * Unfortunately they're in the opposite order in the dex file so we
+             * need to transform them out-of-order.
+             */
+            DexBuffer.Section triesSection = dexOut.open(codeOut.getPosition());
+            codeOut.skip(tries.length * SizeOf.TRY_ITEM);
+            int[] offsets = transformCatchHandlers(indexMap, catchHandlers);
+            transformTries(triesSection, tries, offsets);
+        }
+    }
+
+    /**
+     * Writes the catch handlers to {@code codeOut} and returns their indices.
+     */
+    private int[] transformCatchHandlers(IndexMap indexMap, Code.CatchHandler[] catchHandlers) {
+        int baseOffset = codeOut.getPosition();
+        codeOut.writeUleb128(catchHandlers.length);
+        int[] offsets = new int[catchHandlers.length];
+        for (int i = 0; i < catchHandlers.length; i++) {
+            offsets[i] = codeOut.getPosition() - baseOffset;
+            transformEncodedCatchHandler(catchHandlers[i], indexMap);
+        }
+        return offsets;
+    }
+
+    private void transformTries(DexBuffer.Section out, Code.Try[] tries,
+            int[] catchHandlerOffsets) {
+        for (Code.Try tryItem : tries) {
+            out.writeInt(tryItem.getStartAddress());
+            out.writeUnsignedShort(tryItem.getInstructionCount());
+            out.writeUnsignedShort(catchHandlerOffsets[tryItem.getCatchHandlerIndex()]);
         }
     }