Deduplicate type lists on merged .dex files.

Change-Id: Ic089fecebb66544dfdb422421cce47a6f203a31a
http://b/3447216

6 files changed, 187 insertions, 105 deletions

diff --git a/dx/src/com/android/dx/io/ClassDef.java b/dx/src/com/android/dx/io/ClassDef.java
index a29ea012e..5c8d10b88 100644
--- a/dx/src/com/android/dx/io/ClassDef.java
+++ b/dx/src/com/android/dx/io/ClassDef.java
@@ -27,14 +27,13 @@ public final class ClassDef {
     private final int accessFlags;
     private final int supertypeIndex;
     private final int interfacesOffset;
-    private final short[] interfaces;
     private final int sourceFileIndex;
     private final int annotationsOffset;
     private final int classDataOffset;
     private final int staticValuesOffset;
 
     public ClassDef(DexBuffer buffer, int offset, int typeIndex, int accessFlags,
-            int supertypeIndex, int interfacesOffset, short[] interfaces, int sourceFileIndex,
+            int supertypeIndex, int interfacesOffset, int sourceFileIndex,
             int annotationsOffset, int classDataOffset, int staticValuesOffset) {
         this.buffer = buffer;
         this.offset = offset;
@@ -42,7 +41,6 @@ public final class ClassDef {
         this.accessFlags = accessFlags;
         this.supertypeIndex = supertypeIndex;
         this.interfacesOffset = interfacesOffset;
-        this.interfaces = interfaces;
         this.sourceFileIndex = sourceFileIndex;
         this.annotationsOffset = annotationsOffset;
         this.classDataOffset = classDataOffset;
@@ -66,7 +64,7 @@ public final class ClassDef {
     }
 
     public short[] getInterfaces() {
-        return interfaces;
+        return buffer.readTypeList(interfacesOffset).getTypes();
     }
 
     public int getAccessFlags() {
diff --git a/dx/src/com/android/dx/io/DexBuffer.java b/dx/src/com/android/dx/io/DexBuffer.java
index e971ae660..a281fe5ee 100644
--- a/dx/src/com/android/dx/io/DexBuffer.java
+++ b/dx/src/com/android/dx/io/DexBuffer.java
@@ -18,6 +18,7 @@ package com.android.dx.io;
 
 import com.android.dx.dex.SizeOf;
 import com.android.dx.dex.TableOfContents;
+import com.android.dx.merge.TypeList;
 import com.android.dx.util.DexException;
 import com.android.dx.util.Leb128Utils;
 import com.android.dx.util.Mutf8;
@@ -47,9 +48,8 @@ public final class DexBuffer {
     private final List<String> strings = new AbstractList<String>() {
         @Override public String get(int index) {
             checkBounds(index, tableOfContents.stringIds.size);
-            int offset = open(tableOfContents.stringIds.off + (index * SizeOf.STRING_ID_ITEM))
-                    .readInt();
-            return open(offset).readStringDataItem();
+            return open(tableOfContents.stringIds.off + (index * SizeOf.STRING_ID_ITEM))
+                    .readString();
         }
         @Override public int size() {
             return tableOfContents.stringIds.size;
@@ -224,6 +224,13 @@ public final class DexBuffer {
         };
     }
 
+    public TypeList readTypeList(int offset) {
+        if (offset == 0) {
+            return TypeList.EMPTY;
+        }
+        return open(offset).readTypeList();
+    }
+
     public ClassData readClassData(ClassDef classDef) {
         int offset = classDef.getClassDataOffset();
         if (offset == 0) {
@@ -315,22 +322,20 @@ public final class DexBuffer {
             }
         }
 
-        public short[] readTypeList(int offset) {
-            if (offset == 0) {
-                return new short[0];
-            }
-            int savedPosition = position;
-            position = offset;
+        public TypeList readTypeList() {
             int size = readInt();
-            short[] parameters = new short[size];
+            short[] types = new short[size];
             for (int i = 0; i < size; i++) {
-                parameters[i] = readShort();
+                types[i] = readShort();
             }
-            position = savedPosition;
-            return parameters;
+            alignToFourBytes();
+            return new TypeList(DexBuffer.this, types);
         }
 
-        public String readStringDataItem() {
+        public String readString() {
+            int offset = readInt();
+            int savedPosition = position;
+            position = offset;
             try {
                 int expectedLength = readUleb128();
                 String result = Mutf8.decode(asDataInput, new char[expectedLength]);
@@ -341,6 +346,8 @@ public final class DexBuffer {
                 return result;
             } catch (IOException e) {
                 throw new DexException(e);
+            } finally {
+                position = savedPosition;
             }
         }
 
@@ -361,9 +368,8 @@ public final class DexBuffer {
         public ProtoId readProtoId() {
             int shortyIndex = readInt();
             int returnTypeIndex = readInt();
-            int parametersOff = readInt();
-            short[] parameters = readTypeList(parametersOff);
-            return new ProtoId(DexBuffer.this, shortyIndex, returnTypeIndex, parameters);
+            int parametersOffset = readInt();
+            return new ProtoId(DexBuffer.this, shortyIndex, returnTypeIndex, parametersOffset);
         }
 
         public ClassDef readClassDef() {
@@ -372,14 +378,13 @@ public final class DexBuffer {
             int accessFlags = readInt();
             int supertype = readInt();
             int interfacesOffset = readInt();
-            short[] interfaces = readTypeList(interfacesOffset);
             int sourceFileIndex = readInt();
             int annotationsOffset = readInt();
             int classDataOffset = readInt();
             int staticValuesOffset = readInt();
             return new ClassDef(DexBuffer.this, offset, type, accessFlags, supertype,
-                    interfacesOffset, interfaces, sourceFileIndex, annotationsOffset,
-                    classDataOffset, staticValuesOffset);
+                    interfacesOffset, sourceFileIndex, annotationsOffset, classDataOffset,
+                    staticValuesOffset);
         }
 
         private Code readCode() {
@@ -538,6 +543,15 @@ public final class DexBuffer {
                 throw new AssertionError();
             }
         }
+
+        public void writeTypeList(TypeList typeList) {
+            short[] types = typeList.getTypes();
+            writeInt(types.length);
+            for (short type : types) {
+                writeShort(type);
+            }
+            alignToFourBytes();
+        }
     }
 
     private static class DataInputStub implements DataInput {
diff --git a/dx/src/com/android/dx/io/ProtoId.java b/dx/src/com/android/dx/io/ProtoId.java
index da64c44a3..98c07771f 100644
--- a/dx/src/com/android/dx/io/ProtoId.java
+++ b/dx/src/com/android/dx/io/ProtoId.java
@@ -17,31 +17,25 @@
 package com.android.dx.io;
 
 import com.android.dx.util.Unsigned;
-import java.util.Arrays;
 
 public final class ProtoId implements Comparable<ProtoId> {
     private final DexBuffer buffer;
     private final int shortyIndex;
     private final int returnTypeIndex;
-    private final short[] parameters;
+    private final int parametersOffset;
 
-    public ProtoId(DexBuffer buffer, int shortyIndex, int returnTypeIndex, short[] parameters) {
+    public ProtoId(DexBuffer buffer, int shortyIndex, int returnTypeIndex, int parametersOffset) {
         this.buffer = buffer;
         this.shortyIndex = shortyIndex;
         this.returnTypeIndex = returnTypeIndex;
-        this.parameters = parameters;
+        this.parametersOffset = parametersOffset;
     }
 
     public int compareTo(ProtoId other) {
         if (returnTypeIndex != other.returnTypeIndex) {
             return Unsigned.compare(returnTypeIndex, other.returnTypeIndex);
         }
-        for (int i = 0; i < parameters.length && i < other.parameters.length; i++) {
-            if (parameters[i] != other.parameters[i]) {
-                return Unsigned.compare(parameters[i], other.parameters[i]);
-            }
-        }
-        return Unsigned.compare(parameters.length, other.parameters.length);
+        return Unsigned.compare(parametersOffset, other.parametersOffset);
     }
 
     public int getShortyIndex() {
@@ -52,35 +46,23 @@ public final class ProtoId implements Comparable<ProtoId> {
         return returnTypeIndex;
     }
 
-    public short[] getParameters() {
-        return parameters;
+    public int getParametersOffset() {
+        return parametersOffset;
     }
 
-    public void writeTo(DexBuffer.Section out, int typeListOffset) {
+    public void writeTo(DexBuffer.Section out) {
         out.writeInt(shortyIndex);
         out.writeInt(returnTypeIndex);
-        out.writeInt(typeListOffset);
+        out.writeInt(parametersOffset);
     }
 
     @Override public String toString() {
         if (buffer == null) {
-            return shortyIndex + " " + returnTypeIndex + " " + Arrays.toString(parameters);
+            return shortyIndex + " " + returnTypeIndex + " " + parametersOffset;
         }
 
-        StringBuilder result = new StringBuilder()
-                .append(buffer.strings().get(shortyIndex))
-                .append(": ")
-                .append(buffer.typeNames().get(returnTypeIndex))
-                .append(" (");
-        int j = 0;
-        for (short parameter : parameters) {
-            if (j > 0) {
-                result.append(", ");
-            }
-            result.append(buffer.typeNames().get(parameter));
-            j++;
-        }
-        result.append(")");
-        return result.toString();
+        return buffer.strings().get(shortyIndex)
+                + ": " + buffer.typeNames().get(returnTypeIndex)
+                + " " + buffer.readTypeList(parametersOffset);
     }
 }
diff --git a/dx/src/com/android/dx/merge/DexMerger.java b/dx/src/com/android/dx/merge/DexMerger.java
index fce0091c5..d1dda9420 100644
--- a/dx/src/com/android/dx/merge/DexMerger.java
+++ b/dx/src/com/android/dx/merge/DexMerger.java
@@ -95,8 +95,6 @@ public final class DexMerger {
         /*
          * TODO: several of these sections are far too large than they need to be.
          *
-         * typeList: we don't deduplicate identical type lists. This should be fixed.
-         *
          * classDataWriter: uleb references to code items are larger than
          *     expected. We should use old & new code_item section offsets to
          *     pick an appropriate blow up size
@@ -109,7 +107,7 @@ public final class DexMerger {
         contentsOut.typeLists.off = dexWriter.getLength();
         contentsOut.typeLists.size = 0;
         int maxTypeListBytes = aContents.typeLists.byteCount + bContents.typeLists.byteCount;
-        typeListWriter = dexWriter.appendSection(maxTypeListBytes * 5, "type list");
+        typeListWriter = dexWriter.appendSection(maxTypeListBytes, "type list");
 
         contentsOut.annotationSetRefLists.off = dexWriter.getLength();
         contentsOut.annotationSetRefLists.size = 0;
@@ -169,6 +167,7 @@ public final class DexMerger {
 
         mergeStringIds();
         mergeTypeIds();
+        mergeTypeLists();
         mergeProtoIds();
         mergeFieldIds();
         mergeMethodIds();
@@ -208,6 +207,8 @@ public final class DexMerger {
             TableOfContents.Section bSection = getSection(dexB.getTableOfContents());
             getSection(contentsOut).off = idsDefsWriter.getPosition();
 
+            DexBuffer.Section inA = dexA.open(aSection.off);
+            DexBuffer.Section inB = dexB.open(bSection.off);
             int aIndex = 0;
             int bIndex = 0;
             int outCount = 0;
@@ -215,11 +216,13 @@ public final class DexMerger {
             T b = null;
 
             while (true) {
+                int aOffset = inA.getPosition();
                 if (a == null && aIndex < aSection.size) {
-                    a = read(dexA, aIndexMap, aIndex);
+                    a = read(inA, aIndexMap, aIndex);
                 }
+                int bOffset = inB.getPosition();
                 if (b == null && bIndex < bSection.size) {
-                    b = read(dexB, bIndexMap, bIndex);
+                    b = read(inB, bIndexMap, bIndex);
                 }
 
                 // Write the smaller of a and b. If they're equal, write only once
@@ -237,12 +240,12 @@ public final class DexMerger {
                 T toWrite = null;
                 if (advanceA) {
                     toWrite = a;
-                    updateIndex(aIndexMap, aIndex++, outCount);
+                    updateIndex(aOffset, aIndexMap, aIndex++, outCount);
                     a = null;
                 }
                 if (advanceB) {
                     toWrite = b;
-                    updateIndex(bIndexMap, bIndex++, outCount);
+                    updateIndex(bOffset, bIndexMap, bIndex++, outCount);
                     b = null;
                 }
                 if (toWrite == null) {
@@ -256,8 +259,8 @@ public final class DexMerger {
         }
 
         abstract TableOfContents.Section getSection(TableOfContents tableOfContents);
-        abstract T read(DexBuffer dexBuffer, IndexMap indexMap, int index);
-        abstract void updateIndex(IndexMap indexMap, int oldIndex, int newIndex);
+        abstract T read(DexBuffer.Section in, IndexMap indexMap, int index);
+        abstract void updateIndex(int offset, IndexMap indexMap, int oldIndex, int newIndex);
         abstract void write(T value);
     }
 
@@ -267,11 +270,11 @@ public final class DexMerger {
                 return tableOfContents.stringIds;
             }
 
-            @Override String read(DexBuffer dexBuffer, IndexMap indexMap, int index) {
-                return dexBuffer.strings().get(index);
+            @Override String read(DexBuffer.Section in, IndexMap indexMap, int index) {
+                return in.readString();
             }
 
-            @Override void updateIndex(IndexMap indexMap, int oldIndex, int newIndex) {
+            @Override void updateIndex(int offset, IndexMap indexMap, int oldIndex, int newIndex) {
                 indexMap.stringIds[oldIndex] = newIndex;
             }
 
@@ -289,12 +292,12 @@ public final class DexMerger {
                 return tableOfContents.typeIds;
             }
 
-            @Override Integer read(DexBuffer dexBuffer, IndexMap indexMap, int index) {
-                Integer stringIndex = dexBuffer.typeIds().get(index);
+            @Override Integer read(DexBuffer.Section in, IndexMap indexMap, int index) {
+                int stringIndex = in.readInt();
                 return indexMap.adjustString(stringIndex);
             }
 
-            @Override void updateIndex(IndexMap indexMap, int oldIndex, int newIndex) {
+            @Override void updateIndex(int offset, IndexMap indexMap, int oldIndex, int newIndex) {
                 indexMap.typeIds[oldIndex] = (short) newIndex;
             }
 
@@ -310,17 +313,16 @@ public final class DexMerger {
                 return tableOfContents.protoIds;
             }
 
-            @Override ProtoId read(DexBuffer dexBuffer, IndexMap indexMap, int index) {
-                return indexMap.adjust(dexBuffer.protoIds().get(index));
+            @Override ProtoId read(DexBuffer.Section in, IndexMap indexMap, int index) {
+                return indexMap.adjust(in.readProtoId());
             }
 
-            @Override void updateIndex(IndexMap indexMap, int oldIndex, int newIndex) {
+            @Override void updateIndex(int offset, IndexMap indexMap, int oldIndex, int newIndex) {
                 indexMap.protoIds[oldIndex] = (short) newIndex;
             }
 
             @Override void write(ProtoId value) {
-                int typeListPosition = writeTypeList(value.getParameters());
-                value.writeTo(idsDefsWriter, typeListPosition);
+                value.writeTo(idsDefsWriter);
             }
         }.merge();
     }
@@ -331,11 +333,11 @@ public final class DexMerger {
                 return tableOfContents.fieldIds;
             }
 
-            @Override FieldId read(DexBuffer dexBuffer, IndexMap indexMap, int index) {
-                return indexMap.adjust(dexBuffer.fieldIds().get(index));
+            @Override FieldId read(DexBuffer.Section in, IndexMap indexMap, int index) {
+                return indexMap.adjust(in.readFieldId());
             }
 
-            @Override void updateIndex(IndexMap indexMap, int oldIndex, int newIndex) {
+            @Override void updateIndex(int offset, IndexMap indexMap, int oldIndex, int newIndex) {
                 indexMap.fieldIds[oldIndex] = (short) newIndex;
             }
 
@@ -351,11 +353,11 @@ public final class DexMerger {
                 return tableOfContents.methodIds;
             }
 
-            @Override MethodId read(DexBuffer dexBuffer, IndexMap indexMap, int index) {
-                return indexMap.adjust(dexBuffer.methodIds().get(index));
+            @Override MethodId read(DexBuffer.Section in, IndexMap indexMap, int index) {
+                return indexMap.adjust(in.readMethodId());
             }
 
-            @Override void updateIndex(IndexMap indexMap, int oldIndex, int newIndex) {
+            @Override void updateIndex(int offset, IndexMap indexMap, int oldIndex, int newIndex) {
                 indexMap.methodIds[oldIndex] = (short) newIndex;
             }
 
@@ -365,6 +367,26 @@ public final class DexMerger {
         }.merge();
     }
 
+    private void mergeTypeLists() {
+        new IdMerger<TypeList>() {
+            @Override TableOfContents.Section getSection(TableOfContents tableOfContents) {
+                return tableOfContents.typeLists;
+            }
+
+            @Override TypeList read(DexBuffer.Section in, IndexMap indexMap, int index) {
+                return indexMap.adjustTypeList(in.readTypeList());
+            }
+
+            @Override void updateIndex(int offset, IndexMap indexMap, int oldIndex, int newIndex) {
+                indexMap.typeListOffsets.put(offset, typeListWriter.getPosition());
+            }
+
+            @Override void write(TypeList value) {
+                typeListWriter.writeTypeList(value);
+            }
+        }.merge();
+    }
+
     private void mergeClassDefs() {
         SortableType[] types = getSortedTypes();
         contentsOut.classDefs.off = idsDefsWriter.getPosition();
@@ -439,10 +461,7 @@ public final class DexMerger {
         idsDefsWriter.writeInt(classDef.getTypeIndex());
         idsDefsWriter.writeInt(classDef.getAccessFlags());
         idsDefsWriter.writeInt(classDef.getSupertypeIndex());
-
-        short[] interfaces = classDef.getInterfaces();
-        int typeListPosition = writeTypeList(interfaces);
-        idsDefsWriter.writeInt(typeListPosition);
+        idsDefsWriter.writeInt(classDef.getInterfacesOffset());
 
         int sourceFileIndex = indexMap.adjustString(
                 classDef.getSourceFileIndex()); // source file idx
@@ -477,18 +496,6 @@ public final class DexMerger {
         }
     }
 
-    private int writeTypeList(short[] interfaces) {
-        if (interfaces.length == 0) {
-            return 0;
-        }
-        contentsOut.typeLists.size++;
-        typeListWriter.alignToFourBytes();
-        int cursor = typeListWriter.getPosition();
-        typeListWriter.writeInt(interfaces.length);
-        typeListWriter.write(interfaces);
-        return cursor;
-    }
-
     private void transformAnnotations(DexBuffer.Section in, IndexMap indexMap) {
         contentsOut.annotationsDirectories.size++;
 
diff --git a/dx/src/com/android/dx/merge/IndexMap.java b/dx/src/com/android/dx/merge/IndexMap.java
index f12adc790..72e7c96b6 100644
--- a/dx/src/com/android/dx/merge/IndexMap.java
+++ b/dx/src/com/android/dx/merge/IndexMap.java
@@ -22,6 +22,7 @@ import com.android.dx.io.DexBuffer;
 import com.android.dx.io.FieldId;
 import com.android.dx.io.MethodId;
 import com.android.dx.io.ProtoId;
+import java.util.HashMap;
 
 /**
  * Maps the index offsets from one dex file to those in another. For example, if
@@ -35,6 +36,7 @@ public final class IndexMap {
     public final short[] protoIds;
     public final short[] fieldIds;
     public final short[] methodIds;
+    public final HashMap<Integer, Integer> typeListOffsets;
 
     public IndexMap(DexBuffer target, TableOfContents tableOfContents) {
         this.target = target;
@@ -43,6 +45,13 @@ public final class IndexMap {
         this.protoIds = new short[tableOfContents.protoIds.size];
         this.fieldIds = new short[tableOfContents.fieldIds.size];
         this.methodIds = new short[tableOfContents.methodIds.size];
+        this.typeListOffsets = new HashMap<Integer, Integer>();
+
+        /*
+         * A type list at offset 0 is always the empty type list. Always map
+         * this to itself.
+         */
+        this.typeListOffsets.put(0, 0);
     }
 
     public int adjustString(int stringIndex) {
@@ -53,12 +62,15 @@ public final class IndexMap {
         return (typeIndex == ClassDef.NO_INDEX) ? ClassDef.NO_INDEX : typeIds[typeIndex];
     }
 
-    public short[] adjustTypeList(short[] typeList) {
-        short[] result = new short[typeList.length];
-        for (int i = 0; i < typeList.length; i++) {
-            result[i] = adjustType(typeList[i]);
+    public TypeList adjustTypeList(TypeList typeList) {
+        if (typeList == TypeList.EMPTY) {
+            return typeList;
+        }
+        short[] types = typeList.getTypes().clone();
+        for (int i = 0; i < types.length; i++) {
+            types[i] = adjustType(types[i]);
         }
-        return result;
+        return new TypeList(target, types);
     }
 
     public short adjustProto(int protoIndex) {
@@ -73,6 +85,10 @@ public final class IndexMap {
         return methodIds[methodIndex];
     }
 
+    public int adjustTypeListOffset(int typeListOffset) {
+        return typeListOffsets.get(typeListOffset);
+    }
+
     public MethodId adjust(MethodId methodId) {
         return new MethodId(target,
                 adjustType(methodId.getDeclaringClassIndex()),
@@ -92,15 +108,15 @@ public final class IndexMap {
         return new ProtoId(target,
                 adjustString(protoId.getShortyIndex()),
                 adjustType(protoId.getReturnTypeIndex()),
-                adjustTypeList(protoId.getParameters()));
+                adjustTypeListOffset(protoId.getParametersOffset()));
     }
 
     public ClassDef adjust(ClassDef classDef) {
         return new ClassDef(target, classDef.getOffset(), adjustType(classDef.getTypeIndex()),
                 classDef.getAccessFlags(), adjustType(classDef.getSupertypeIndex()),
-                classDef.getInterfacesOffset(), adjustTypeList(classDef.getInterfaces()),
-                classDef.getSourceFileIndex(), classDef.getAnnotationsOffset(),
-                classDef.getClassDataOffset(), classDef.getStaticValuesOffset());
+                adjustTypeListOffset(classDef.getInterfacesOffset()), classDef.getSourceFileIndex(),
+                classDef.getAnnotationsOffset(), classDef.getClassDataOffset(),
+                classDef.getStaticValuesOffset());
     }
 
     public SortableType adjust(SortableType sortableType) {
diff --git a/dx/src/com/android/dx/merge/TypeList.java b/dx/src/com/android/dx/merge/TypeList.java
new file mode 100644
index 000000000..1619f1912
--- /dev/null
+++ b/dx/src/com/android/dx/merge/TypeList.java
@@ -0,0 +1,65 @@
+/*
+ * Copyright (C) 2011 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.dx.merge;
+
+import com.android.dx.io.DexBuffer;
+import com.android.dx.util.Unsigned;
+import java.util.Arrays;
+
+public final class TypeList implements Comparable<TypeList> {
+
+    public static final TypeList EMPTY = new TypeList(null, new short[0]);
+
+    private final DexBuffer buffer;
+    private final short[] types;
+
+    public TypeList(DexBuffer buffer, short[] types) {
+        this.buffer = buffer;
+        this.types = types;
+    }
+
+    public short[] getTypes() {
+        return types;
+    }
+
+    public int compareTo(TypeList other) {
+        for (int i = 0; i < types.length && i < other.types.length; i++) {
+            if (types[i] != other.types[i]) {
+                return Unsigned.compare(types[i], other.types[i]);
+            }
+        }
+        return Unsigned.compare(types.length, other.types.length);
+    }
+
+    @Override public String toString() {
+        if (buffer == null) {
+            return Arrays.toString(types);
+        }
+
+        StringBuilder result = new StringBuilder();
+        result.append("[");
+        for (int i = 0, typesLength = types.length; i < typesLength; i++) {
+            short parameter = types[i];
+            if (i > 0) {
+                result.append(", ");
+            }
+            result.append(buffer.typeNames().get(parameter));
+        }
+        result.append("]");
+        return result.toString();
+    }
+}