1 /*
2  * Copyright (C) 2011 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.dx.merge;
18 
19 import com.android.dex.Annotation;
20 import com.android.dex.ClassData;
21 import com.android.dex.ClassDef;
22 import com.android.dex.Code;
23 import com.android.dex.Dex;
24 import com.android.dex.DexException;
25 import com.android.dex.DexIndexOverflowException;
26 import com.android.dex.FieldId;
27 import com.android.dex.MethodId;
28 import com.android.dex.ProtoId;
29 import com.android.dex.SizeOf;
30 import com.android.dex.TableOfContents;
31 import com.android.dex.TypeList;
32 
33 import java.io.File;
34 import java.io.IOException;
35 import java.util.ArrayList;
36 import java.util.Arrays;
37 import java.util.Collections;
38 import java.util.List;
39 
40 /**
41  * Combine two dex files into one.
42  */
43 public final class DexMerger {
44     private final Dex dexA;
45     private final Dex dexB;
46     private final CollisionPolicy collisionPolicy;
47     private final WriterSizes writerSizes;
48 
49     private final Dex dexOut;
50 
51     private final Dex.Section headerOut;
52 
53     /** All IDs and definitions sections */
54     private final Dex.Section idsDefsOut;
55 
56     private final Dex.Section mapListOut;
57 
58     private final Dex.Section typeListOut;
59 
60     private final Dex.Section classDataOut;
61 
62     private final Dex.Section codeOut;
63 
64     private final Dex.Section stringDataOut;
65 
66     private final Dex.Section debugInfoOut;
67 
68     private final Dex.Section encodedArrayOut;
69 
70     /** annotations directory on a type */
71     private final Dex.Section annotationsDirectoryOut;
72 
73     /** sets of annotations on a member, parameter or type */
74     private final Dex.Section annotationSetOut;
75 
76     /** parameter lists */
77     private final Dex.Section annotationSetRefListOut;
78 
79     /** individual annotations, each containing zero or more fields */
80     private final Dex.Section annotationOut;
81 
82     private final TableOfContents contentsOut;
83 
84     private final IndexMap aIndexMap;
85     private final IndexMap bIndexMap;
86     private final InstructionTransformer aInstructionTransformer;
87     private final InstructionTransformer bInstructionTransformer;
88 
89     /** minimum number of wasted bytes before it's worthwhile to compact the result */
90     private int compactWasteThreshold = 1024 * 1024; // 1MiB
91 
DexMerger(Dex dexA, Dex dexB, CollisionPolicy collisionPolicy)92     public DexMerger(Dex dexA, Dex dexB, CollisionPolicy collisionPolicy)
93             throws IOException {
94         this(dexA, dexB, collisionPolicy, new WriterSizes(dexA, dexB));
95     }
96 
DexMerger(Dex dexA, Dex dexB, CollisionPolicy collisionPolicy, WriterSizes writerSizes)97     private DexMerger(Dex dexA, Dex dexB, CollisionPolicy collisionPolicy,
98             WriterSizes writerSizes) throws IOException {
99         this.dexA = dexA;
100         this.dexB = dexB;
101         this.collisionPolicy = collisionPolicy;
102         this.writerSizes = writerSizes;
103 
104         dexOut = new Dex(writerSizes.size());
105 
106         TableOfContents aContents = dexA.getTableOfContents();
107         TableOfContents bContents = dexB.getTableOfContents();
108         aIndexMap = new IndexMap(dexOut, aContents);
109         bIndexMap = new IndexMap(dexOut, bContents);
110         aInstructionTransformer = new InstructionTransformer(aIndexMap);
111         bInstructionTransformer = new InstructionTransformer(bIndexMap);
112 
113         headerOut = dexOut.appendSection(writerSizes.header, "header");
114         idsDefsOut = dexOut.appendSection(writerSizes.idsDefs, "ids defs");
115 
116         contentsOut = dexOut.getTableOfContents();
117         contentsOut.dataOff = dexOut.getNextSectionStart();
118 
119         contentsOut.mapList.off = dexOut.getNextSectionStart();
120         contentsOut.mapList.size = 1;
121         mapListOut = dexOut.appendSection(writerSizes.mapList, "map list");
122 
123         contentsOut.typeLists.off = dexOut.getNextSectionStart();
124         typeListOut = dexOut.appendSection(writerSizes.typeList, "type list");
125 
126         contentsOut.annotationSetRefLists.off = dexOut.getNextSectionStart();
127         annotationSetRefListOut = dexOut.appendSection(
128                 writerSizes.annotationsSetRefList, "annotation set ref list");
129 
130         contentsOut.annotationSets.off = dexOut.getNextSectionStart();
131         annotationSetOut = dexOut.appendSection(writerSizes.annotationsSet, "annotation sets");
132 
133         contentsOut.classDatas.off = dexOut.getNextSectionStart();
134         classDataOut = dexOut.appendSection(writerSizes.classData, "class data");
135 
136         contentsOut.codes.off = dexOut.getNextSectionStart();
137         codeOut = dexOut.appendSection(writerSizes.code, "code");
138 
139         contentsOut.stringDatas.off = dexOut.getNextSectionStart();
140         stringDataOut = dexOut.appendSection(writerSizes.stringData, "string data");
141 
142         contentsOut.debugInfos.off = dexOut.getNextSectionStart();
143         debugInfoOut = dexOut.appendSection(writerSizes.debugInfo, "debug info");
144 
145         contentsOut.annotations.off = dexOut.getNextSectionStart();
146         annotationOut = dexOut.appendSection(writerSizes.annotation, "annotation");
147 
148         contentsOut.encodedArrays.off = dexOut.getNextSectionStart();
149         encodedArrayOut = dexOut.appendSection(writerSizes.encodedArray, "encoded array");
150 
151         contentsOut.annotationsDirectories.off = dexOut.getNextSectionStart();
152         annotationsDirectoryOut = dexOut.appendSection(
153                 writerSizes.annotationsDirectory, "annotations directory");
154 
155         contentsOut.dataSize = dexOut.getNextSectionStart() - contentsOut.dataOff;
156     }
157 
setCompactWasteThreshold(int compactWasteThreshold)158     public void setCompactWasteThreshold(int compactWasteThreshold) {
159         this.compactWasteThreshold = compactWasteThreshold;
160     }
161 
mergeDexes()162     private Dex mergeDexes() throws IOException {
163         mergeStringIds();
164         mergeTypeIds();
165         mergeTypeLists();
166         mergeProtoIds();
167         mergeFieldIds();
168         mergeMethodIds();
169         mergeAnnotations();
170         unionAnnotationSetsAndDirectories();
171         mergeClassDefs();
172 
173         // write the header
174         contentsOut.header.off = 0;
175         contentsOut.header.size = 1;
176         contentsOut.fileSize = dexOut.getLength();
177         contentsOut.computeSizesFromOffsets();
178         contentsOut.writeHeader(headerOut);
179         contentsOut.writeMap(mapListOut);
180 
181         // generate and write the hashes
182         dexOut.writeHashes();
183 
184         return dexOut;
185     }
186 
merge()187     public Dex merge() throws IOException {
188         long start = System.nanoTime();
189         Dex result = mergeDexes();
190 
191         /*
192          * We use pessimistic sizes when merging dex files. If those sizes
193          * result in too many bytes wasted, compact the result. To compact,
194          * simply merge the result with itself.
195          */
196         WriterSizes compactedSizes = new WriterSizes(this);
197         int wastedByteCount = writerSizes.size() - compactedSizes.size();
198         if (wastedByteCount >  + compactWasteThreshold) {
199             DexMerger compacter = new DexMerger(
200                     dexOut, new Dex(0), CollisionPolicy.FAIL, compactedSizes);
201             result = compacter.mergeDexes();
202             System.out.printf("Result compacted from %.1fKiB to %.1fKiB to save %.1fKiB%n",
203                     dexOut.getLength() / 1024f,
204                     result.getLength() / 1024f,
205                     wastedByteCount / 1024f);
206         }
207 
208         long elapsed = System.nanoTime() - start;
209         System.out.printf("Merged dex A (%d defs/%.1fKiB) with dex B "
210                 + "(%d defs/%.1fKiB). Result is %d defs/%.1fKiB. Took %.1fs%n",
211                 dexA.getTableOfContents().classDefs.size,
212                 dexA.getLength() / 1024f,
213                 dexB.getTableOfContents().classDefs.size,
214                 dexB.getLength() / 1024f,
215                 result.getTableOfContents().classDefs.size,
216                 result.getLength() / 1024f,
217                 elapsed / 1000000000f);
218 
219         return result;
220     }
221 
222     /**
223      * Reads an IDs section of two dex files and writes an IDs section of a
224      * merged dex file. Populates maps from old to new indices in the process.
225      */
226     abstract class IdMerger<T extends Comparable<T>> {
227         private final Dex.Section out;
228 
IdMerger(Dex.Section out)229         protected IdMerger(Dex.Section out) {
230             this.out = out;
231         }
232 
233         /**
234          * Merges already-sorted sections, reading only two values into memory
235          * at a time.
236          */
mergeSorted()237         public final void mergeSorted() {
238             TableOfContents.Section aSection = getSection(dexA.getTableOfContents());
239             TableOfContents.Section bSection = getSection(dexB.getTableOfContents());
240             getSection(contentsOut).off = out.getPosition();
241 
242             Dex.Section inA = aSection.exists() ? dexA.open(aSection.off) : null;
243             Dex.Section inB = bSection.exists() ? dexB.open(bSection.off) : null;
244             int aOffset = -1;
245             int bOffset = -1;
246             int aIndex = 0;
247             int bIndex = 0;
248             int outCount = 0;
249             T a = null;
250             T b = null;
251 
252             while (true) {
253                 if (a == null && aIndex < aSection.size) {
254                     aOffset = inA.getPosition();
255                     a = read(inA, aIndexMap, aIndex);
256                 }
257                 if (b == null && bIndex < bSection.size) {
258                     bOffset = inB.getPosition();
259                     b = read(inB, bIndexMap, bIndex);
260                 }
261 
262                 // Write the smaller of a and b. If they're equal, write only once
263                 boolean advanceA;
264                 boolean advanceB;
265                 if (a != null && b != null) {
266                     int compare = a.compareTo(b);
267                     advanceA = compare <= 0;
268                     advanceB = compare >= 0;
269                 } else {
270                     advanceA = (a != null);
271                     advanceB = (b != null);
272                 }
273 
274                 T toWrite = null;
275                 if (advanceA) {
276                     toWrite = a;
277                     updateIndex(aOffset, aIndexMap, aIndex++, outCount);
278                     a = null;
279                     aOffset = -1;
280                 }
281                 if (advanceB) {
282                     toWrite = b;
283                     updateIndex(bOffset, bIndexMap, bIndex++, outCount);
284                     b = null;
285                     bOffset = -1;
286                 }
287                 if (toWrite == null) {
288                     break; // advanceA == false && advanceB == false
289                 }
290                 write(toWrite);
291                 outCount++;
292             }
293 
294             getSection(contentsOut).size = outCount;
295         }
296 
297         /**
298          * Merges unsorted sections by reading them completely into memory and
299          * sorting in memory.
300          */
mergeUnsorted()301         public final void mergeUnsorted() {
302             getSection(contentsOut).off = out.getPosition();
303 
304             List<UnsortedValue> all = new ArrayList<UnsortedValue>();
305             all.addAll(readUnsortedValues(dexA, aIndexMap));
306             all.addAll(readUnsortedValues(dexB, bIndexMap));
307             Collections.sort(all);
308 
309             int outCount = 0;
310             for (int i = 0; i < all.size(); ) {
311                 UnsortedValue e1 = all.get(i++);
312                 updateIndex(e1.offset, getIndexMap(e1.source), e1.index, outCount - 1);
313 
314                 while (i < all.size() && e1.compareTo(all.get(i)) == 0) {
315                     UnsortedValue e2 = all.get(i++);
316                     updateIndex(e2.offset, getIndexMap(e2.source), e2.index, outCount - 1);
317                 }
318 
319                 write(e1.value);
320                 outCount++;
321             }
322 
323             getSection(contentsOut).size = outCount;
324         }
325 
readUnsortedValues(Dex source, IndexMap indexMap)326         private List<UnsortedValue> readUnsortedValues(Dex source, IndexMap indexMap) {
327             TableOfContents.Section section = getSection(source.getTableOfContents());
328             if (!section.exists()) {
329                 return Collections.emptyList();
330             }
331 
332             List<UnsortedValue> result = new ArrayList<UnsortedValue>();
333             Dex.Section in = source.open(section.off);
334             for (int i = 0; i < section.size; i++) {
335                 int offset = in.getPosition();
336                 T value = read(in, indexMap, 0);
337                 result.add(new UnsortedValue(source, indexMap, value, i, offset));
338             }
339             return result;
340         }
341 
getSection(TableOfContents tableOfContents)342         abstract TableOfContents.Section getSection(TableOfContents tableOfContents);
read(Dex.Section in, IndexMap indexMap, int index)343         abstract T read(Dex.Section in, IndexMap indexMap, int index);
updateIndex(int offset, IndexMap indexMap, int oldIndex, int newIndex)344         abstract void updateIndex(int offset, IndexMap indexMap, int oldIndex, int newIndex);
write(T value)345         abstract void write(T value);
346 
347         class UnsortedValue implements Comparable<UnsortedValue> {
348             final Dex source;
349             final IndexMap indexMap;
350             final T value;
351             final int index;
352             final int offset;
353 
UnsortedValue(Dex source, IndexMap indexMap, T value, int index, int offset)354             UnsortedValue(Dex source, IndexMap indexMap, T value, int index, int offset) {
355                 this.source = source;
356                 this.indexMap = indexMap;
357                 this.value = value;
358                 this.index = index;
359                 this.offset = offset;
360             }
361 
compareTo(UnsortedValue unsortedValue)362             public int compareTo(UnsortedValue unsortedValue) {
363                 return value.compareTo(unsortedValue.value);
364             }
365         }
366     }
367 
getIndexMap(Dex dex)368     private IndexMap getIndexMap(Dex dex) {
369         if (dex == dexA) {
370             return aIndexMap;
371         } else if (dex == dexB) {
372             return bIndexMap;
373         } else {
374             throw new IllegalArgumentException();
375         }
376     }
377 
mergeStringIds()378     private void mergeStringIds() {
379         new IdMerger<String>(idsDefsOut) {
380             @Override TableOfContents.Section getSection(TableOfContents tableOfContents) {
381                 return tableOfContents.stringIds;
382             }
383 
384             @Override String read(Dex.Section in, IndexMap indexMap, int index) {
385                 return in.readString();
386             }
387 
388             @Override void updateIndex(int offset, IndexMap indexMap, int oldIndex, int newIndex) {
389                 indexMap.stringIds[oldIndex] = newIndex;
390             }
391 
392             @Override void write(String value) {
393                 contentsOut.stringDatas.size++;
394                 idsDefsOut.writeInt(stringDataOut.getPosition());
395                 stringDataOut.writeStringData(value);
396             }
397         }.mergeSorted();
398     }
399 
mergeTypeIds()400     private void mergeTypeIds() {
401         new IdMerger<Integer>(idsDefsOut) {
402             @Override TableOfContents.Section getSection(TableOfContents tableOfContents) {
403                 return tableOfContents.typeIds;
404             }
405 
406             @Override Integer read(Dex.Section in, IndexMap indexMap, int index) {
407                 int stringIndex = in.readInt();
408                 return indexMap.adjustString(stringIndex);
409             }
410 
411             @Override void updateIndex(int offset, IndexMap indexMap, int oldIndex, int newIndex) {
412                 if (newIndex < 0 || newIndex > 0xffff) {
413                     throw new DexIndexOverflowException("type ID not in [0, 0xffff]: " + newIndex);
414                 }
415                 indexMap.typeIds[oldIndex] = (short) newIndex;
416             }
417 
418             @Override void write(Integer value) {
419                 idsDefsOut.writeInt(value);
420             }
421         }.mergeSorted();
422     }
423 
mergeTypeLists()424     private void mergeTypeLists() {
425         new IdMerger<TypeList>(typeListOut) {
426             @Override TableOfContents.Section getSection(TableOfContents tableOfContents) {
427                 return tableOfContents.typeLists;
428             }
429 
430             @Override TypeList read(Dex.Section in, IndexMap indexMap, int index) {
431                 return indexMap.adjustTypeList(in.readTypeList());
432             }
433 
434             @Override void updateIndex(int offset, IndexMap indexMap, int oldIndex, int newIndex) {
435                 indexMap.putTypeListOffset(offset, typeListOut.getPosition());
436             }
437 
438             @Override void write(TypeList value) {
439                 typeListOut.writeTypeList(value);
440             }
441         }.mergeUnsorted();
442     }
443 
mergeProtoIds()444     private void mergeProtoIds() {
445         new IdMerger<ProtoId>(idsDefsOut) {
446             @Override TableOfContents.Section getSection(TableOfContents tableOfContents) {
447                 return tableOfContents.protoIds;
448             }
449 
450             @Override ProtoId read(Dex.Section in, IndexMap indexMap, int index) {
451                 return indexMap.adjust(in.readProtoId());
452             }
453 
454             @Override void updateIndex(int offset, IndexMap indexMap, int oldIndex, int newIndex) {
455                 if (newIndex < 0 || newIndex > 0xffff) {
456                     throw new DexIndexOverflowException("proto ID not in [0, 0xffff]: " + newIndex);
457                 }
458                 indexMap.protoIds[oldIndex] = (short) newIndex;
459             }
460 
461             @Override void write(ProtoId value) {
462                 value.writeTo(idsDefsOut);
463             }
464         }.mergeSorted();
465     }
466 
mergeFieldIds()467     private void mergeFieldIds() {
468         new IdMerger<FieldId>(idsDefsOut) {
469             @Override TableOfContents.Section getSection(TableOfContents tableOfContents) {
470                 return tableOfContents.fieldIds;
471             }
472 
473             @Override FieldId read(Dex.Section in, IndexMap indexMap, int index) {
474                 return indexMap.adjust(in.readFieldId());
475             }
476 
477             @Override void updateIndex(int offset, IndexMap indexMap, int oldIndex, int newIndex) {
478                 if (newIndex < 0 || newIndex > 0xffff) {
479                     throw new DexIndexOverflowException("field ID not in [0, 0xffff]: " + newIndex);
480                 }
481                 indexMap.fieldIds[oldIndex] = (short) newIndex;
482             }
483 
484             @Override void write(FieldId value) {
485                 value.writeTo(idsDefsOut);
486             }
487         }.mergeSorted();
488     }
489 
mergeMethodIds()490     private void mergeMethodIds() {
491         new IdMerger<MethodId>(idsDefsOut) {
492             @Override TableOfContents.Section getSection(TableOfContents tableOfContents) {
493                 return tableOfContents.methodIds;
494             }
495 
496             @Override MethodId read(Dex.Section in, IndexMap indexMap, int index) {
497                 return indexMap.adjust(in.readMethodId());
498             }
499 
500             @Override void updateIndex(int offset, IndexMap indexMap, int oldIndex, int newIndex) {
501                 if (newIndex < 0 || newIndex > 0xffff) {
502                     throw new DexIndexOverflowException(
503                         "method ID not in [0, 0xffff]: " + newIndex);
504                 }
505                 indexMap.methodIds[oldIndex] = (short) newIndex;
506             }
507 
508             @Override void write(MethodId methodId) {
509                 methodId.writeTo(idsDefsOut);
510             }
511         }.mergeSorted();
512     }
513 
mergeAnnotations()514     private void mergeAnnotations() {
515         new IdMerger<Annotation>(annotationOut) {
516             @Override TableOfContents.Section getSection(TableOfContents tableOfContents) {
517                 return tableOfContents.annotations;
518             }
519 
520             @Override Annotation read(Dex.Section in, IndexMap indexMap, int index) {
521                 return indexMap.adjust(in.readAnnotation());
522             }
523 
524             @Override void updateIndex(int offset, IndexMap indexMap, int oldIndex, int newIndex) {
525                 indexMap.putAnnotationOffset(offset, annotationOut.getPosition());
526             }
527 
528             @Override void write(Annotation value) {
529                 value.writeTo(annotationOut);
530             }
531         }.mergeUnsorted();
532     }
533 
mergeClassDefs()534     private void mergeClassDefs() {
535         SortableType[] types = getSortedTypes();
536         contentsOut.classDefs.off = idsDefsOut.getPosition();
537         contentsOut.classDefs.size = types.length;
538 
539         for (SortableType type : types) {
540             Dex in = type.getDex();
541             IndexMap indexMap = (in == dexA) ? aIndexMap : bIndexMap;
542             transformClassDef(in, type.getClassDef(), indexMap);
543         }
544     }
545 
546     /**
547      * Returns the union of classes from both files, sorted in order such that
548      * a class is always preceded by its supertype and implemented interfaces.
549      */
getSortedTypes()550     private SortableType[] getSortedTypes() {
551         // size is pessimistic; doesn't include arrays
552         SortableType[] sortableTypes = new SortableType[contentsOut.typeIds.size];
553         readSortableTypes(sortableTypes, dexA, aIndexMap);
554         readSortableTypes(sortableTypes, dexB, bIndexMap);
555 
556         /*
557          * Populate the depths of each sortable type. This makes D iterations
558          * through all N types, where 'D' is the depth of the deepest type. For
559          * example, the deepest class in libcore is Xalan's KeyIterator, which
560          * is 11 types deep.
561          */
562         while (true) {
563             boolean allDone = true;
564             for (SortableType sortableType : sortableTypes) {
565                 if (sortableType != null && !sortableType.isDepthAssigned()) {
566                     allDone &= sortableType.tryAssignDepth(sortableTypes);
567                 }
568             }
569             if (allDone) {
570                 break;
571             }
572         }
573 
574         // Now that all types have depth information, the result can be sorted
575         Arrays.sort(sortableTypes, SortableType.NULLS_LAST_ORDER);
576 
577         // Strip nulls from the end
578         int firstNull = Arrays.asList(sortableTypes).indexOf(null);
579         return firstNull != -1
580                 ? Arrays.copyOfRange(sortableTypes, 0, firstNull)
581                 : sortableTypes;
582     }
583 
584     /**
585      * Reads just enough data on each class so that we can sort it and then find
586      * it later.
587      */
readSortableTypes(SortableType[] sortableTypes, Dex buffer, IndexMap indexMap)588     private void readSortableTypes(SortableType[] sortableTypes, Dex buffer,
589             IndexMap indexMap) {
590         for (ClassDef classDef : buffer.classDefs()) {
591             SortableType sortableType = indexMap.adjust(new SortableType(buffer, classDef));
592             int t = sortableType.getTypeIndex();
593             if (sortableTypes[t] == null) {
594                 sortableTypes[t] = sortableType;
595             } else if (collisionPolicy != CollisionPolicy.KEEP_FIRST) {
596                 throw new DexException("Multiple dex files define "
597                         + buffer.typeNames().get(classDef.getTypeIndex()));
598             }
599         }
600     }
601 
602     /**
603      * Copy annotation sets from each input to the output.
604      *
605      * TODO: this may write multiple copies of the same annotation set.
606      * We should shrink the output by merging rather than unioning
607      */
unionAnnotationSetsAndDirectories()608     private void unionAnnotationSetsAndDirectories() {
609         transformAnnotationSets(dexA, aIndexMap);
610         transformAnnotationSets(dexB, bIndexMap);
611         transformAnnotationSetRefLists(dexA, aIndexMap);
612         transformAnnotationSetRefLists(dexB, bIndexMap);
613         transformAnnotationDirectories(dexA, aIndexMap);
614         transformAnnotationDirectories(dexB, bIndexMap);
615         transformStaticValues(dexA, aIndexMap);
616         transformStaticValues(dexB, bIndexMap);
617     }
618 
transformAnnotationSets(Dex in, IndexMap indexMap)619     private void transformAnnotationSets(Dex in, IndexMap indexMap) {
620         TableOfContents.Section section = in.getTableOfContents().annotationSets;
621         if (section.exists()) {
622             Dex.Section setIn = in.open(section.off);
623             for (int i = 0; i < section.size; i++) {
624                 transformAnnotationSet(indexMap, setIn);
625             }
626         }
627     }
628 
transformAnnotationSetRefLists(Dex in, IndexMap indexMap)629     private void transformAnnotationSetRefLists(Dex in, IndexMap indexMap) {
630         TableOfContents.Section section = in.getTableOfContents().annotationSetRefLists;
631         if (section.exists()) {
632             Dex.Section setIn = in.open(section.off);
633             for (int i = 0; i < section.size; i++) {
634                 transformAnnotationSetRefList(indexMap, setIn);
635             }
636         }
637     }
638 
transformAnnotationDirectories(Dex in, IndexMap indexMap)639     private void transformAnnotationDirectories(Dex in, IndexMap indexMap) {
640         TableOfContents.Section section = in.getTableOfContents().annotationsDirectories;
641         if (section.exists()) {
642             Dex.Section directoryIn = in.open(section.off);
643             for (int i = 0; i < section.size; i++) {
644                 transformAnnotationDirectory(directoryIn, indexMap);
645             }
646         }
647     }
648 
transformStaticValues(Dex in, IndexMap indexMap)649     private void transformStaticValues(Dex in, IndexMap indexMap) {
650         TableOfContents.Section section = in.getTableOfContents().encodedArrays;
651         if (section.exists()) {
652             Dex.Section staticValuesIn = in.open(section.off);
653             for (int i = 0; i < section.size; i++) {
654                 transformStaticValues(staticValuesIn, indexMap);
655             }
656         }
657     }
658 
659     /**
660      * Reads a class_def_item beginning at {@code in} and writes the index and
661      * data.
662      */
transformClassDef(Dex in, ClassDef classDef, IndexMap indexMap)663     private void transformClassDef(Dex in, ClassDef classDef, IndexMap indexMap) {
664         idsDefsOut.assertFourByteAligned();
665         idsDefsOut.writeInt(classDef.getTypeIndex());
666         idsDefsOut.writeInt(classDef.getAccessFlags());
667         idsDefsOut.writeInt(classDef.getSupertypeIndex());
668         idsDefsOut.writeInt(classDef.getInterfacesOffset());
669 
670         int sourceFileIndex = indexMap.adjustString(classDef.getSourceFileIndex());
671         idsDefsOut.writeInt(sourceFileIndex);
672 
673         int annotationsOff = classDef.getAnnotationsOffset();
674         idsDefsOut.writeInt(indexMap.adjustAnnotationDirectory(annotationsOff));
675 
676         int classDataOff = classDef.getClassDataOffset();
677         if (classDataOff == 0) {
678             idsDefsOut.writeInt(0);
679         } else {
680             idsDefsOut.writeInt(classDataOut.getPosition());
681             ClassData classData = in.readClassData(classDef);
682             transformClassData(in, classData, indexMap);
683         }
684 
685         int staticValuesOff = classDef.getStaticValuesOffset();
686         idsDefsOut.writeInt(indexMap.adjustStaticValues(staticValuesOff));
687     }
688 
689     /**
690      * Transform all annotations on a class.
691      */
transformAnnotationDirectory( Dex.Section directoryIn, IndexMap indexMap)692     private void transformAnnotationDirectory(
693             Dex.Section directoryIn, IndexMap indexMap) {
694         contentsOut.annotationsDirectories.size++;
695         annotationsDirectoryOut.assertFourByteAligned();
696         indexMap.putAnnotationDirectoryOffset(
697                 directoryIn.getPosition(), annotationsDirectoryOut.getPosition());
698 
699         int classAnnotationsOffset = indexMap.adjustAnnotationSet(directoryIn.readInt());
700         annotationsDirectoryOut.writeInt(classAnnotationsOffset);
701 
702         int fieldsSize = directoryIn.readInt();
703         annotationsDirectoryOut.writeInt(fieldsSize);
704 
705         int methodsSize = directoryIn.readInt();
706         annotationsDirectoryOut.writeInt(methodsSize);
707 
708         int parameterListSize = directoryIn.readInt();
709         annotationsDirectoryOut.writeInt(parameterListSize);
710 
711         for (int i = 0; i < fieldsSize; i++) {
712             // field index
713             annotationsDirectoryOut.writeInt(indexMap.adjustField(directoryIn.readInt()));
714 
715             // annotations offset
716             annotationsDirectoryOut.writeInt(indexMap.adjustAnnotationSet(directoryIn.readInt()));
717         }
718 
719         for (int i = 0; i < methodsSize; i++) {
720             // method index
721             annotationsDirectoryOut.writeInt(indexMap.adjustMethod(directoryIn.readInt()));
722 
723             // annotation set offset
724             annotationsDirectoryOut.writeInt(
725                     indexMap.adjustAnnotationSet(directoryIn.readInt()));
726         }
727 
728         for (int i = 0; i < parameterListSize; i++) {
729             // method index
730             annotationsDirectoryOut.writeInt(indexMap.adjustMethod(directoryIn.readInt()));
731 
732             // annotations offset
733             annotationsDirectoryOut.writeInt(
734                     indexMap.adjustAnnotationSetRefList(directoryIn.readInt()));
735         }
736     }
737 
738     /**
739      * Transform all annotations on a single type, member or parameter.
740      */
transformAnnotationSet(IndexMap indexMap, Dex.Section setIn)741     private void transformAnnotationSet(IndexMap indexMap, Dex.Section setIn) {
742         contentsOut.annotationSets.size++;
743         annotationSetOut.assertFourByteAligned();
744         indexMap.putAnnotationSetOffset(setIn.getPosition(), annotationSetOut.getPosition());
745 
746         int size = setIn.readInt();
747         annotationSetOut.writeInt(size);
748 
749         for (int j = 0; j < size; j++) {
750             annotationSetOut.writeInt(indexMap.adjustAnnotation(setIn.readInt()));
751         }
752     }
753 
754     /**
755      * Transform all annotation set ref lists.
756      */
transformAnnotationSetRefList(IndexMap indexMap, Dex.Section refListIn)757     private void transformAnnotationSetRefList(IndexMap indexMap, Dex.Section refListIn) {
758         contentsOut.annotationSetRefLists.size++;
759         annotationSetRefListOut.assertFourByteAligned();
760         indexMap.putAnnotationSetRefListOffset(
761                 refListIn.getPosition(), annotationSetRefListOut.getPosition());
762 
763         int parameterCount = refListIn.readInt();
764         annotationSetRefListOut.writeInt(parameterCount);
765         for (int p = 0; p < parameterCount; p++) {
766             annotationSetRefListOut.writeInt(indexMap.adjustAnnotationSet(refListIn.readInt()));
767         }
768     }
769 
transformClassData(Dex in, ClassData classData, IndexMap indexMap)770     private void transformClassData(Dex in, ClassData classData, IndexMap indexMap) {
771         contentsOut.classDatas.size++;
772 
773         ClassData.Field[] staticFields = classData.getStaticFields();
774         ClassData.Field[] instanceFields = classData.getInstanceFields();
775         ClassData.Method[] directMethods = classData.getDirectMethods();
776         ClassData.Method[] virtualMethods = classData.getVirtualMethods();
777 
778         classDataOut.writeUleb128(staticFields.length);
779         classDataOut.writeUleb128(instanceFields.length);
780         classDataOut.writeUleb128(directMethods.length);
781         classDataOut.writeUleb128(virtualMethods.length);
782 
783         transformFields(indexMap, staticFields);
784         transformFields(indexMap, instanceFields);
785         transformMethods(in, indexMap, directMethods);
786         transformMethods(in, indexMap, virtualMethods);
787     }
788 
transformFields(IndexMap indexMap, ClassData.Field[] fields)789     private void transformFields(IndexMap indexMap, ClassData.Field[] fields) {
790         int lastOutFieldIndex = 0;
791         for (ClassData.Field field : fields) {
792             int outFieldIndex = indexMap.adjustField(field.getFieldIndex());
793             classDataOut.writeUleb128(outFieldIndex - lastOutFieldIndex);
794             lastOutFieldIndex = outFieldIndex;
795             classDataOut.writeUleb128(field.getAccessFlags());
796         }
797     }
798 
transformMethods(Dex in, IndexMap indexMap, ClassData.Method[] methods)799     private void transformMethods(Dex in, IndexMap indexMap, ClassData.Method[] methods) {
800         int lastOutMethodIndex = 0;
801         for (ClassData.Method method : methods) {
802             int outMethodIndex = indexMap.adjustMethod(method.getMethodIndex());
803             classDataOut.writeUleb128(outMethodIndex - lastOutMethodIndex);
804             lastOutMethodIndex = outMethodIndex;
805 
806             classDataOut.writeUleb128(method.getAccessFlags());
807 
808             if (method.getCodeOffset() == 0) {
809                 classDataOut.writeUleb128(0);
810             } else {
811                 codeOut.alignToFourBytesWithZeroFill();
812                 classDataOut.writeUleb128(codeOut.getPosition());
813                 transformCode(in, in.readCode(method), indexMap);
814             }
815         }
816     }
817 
transformCode(Dex in, Code code, IndexMap indexMap)818     private void transformCode(Dex in, Code code, IndexMap indexMap) {
819         contentsOut.codes.size++;
820         codeOut.assertFourByteAligned();
821 
822         codeOut.writeUnsignedShort(code.getRegistersSize());
823         codeOut.writeUnsignedShort(code.getInsSize());
824         codeOut.writeUnsignedShort(code.getOutsSize());
825 
826         Code.Try[] tries = code.getTries();
827         Code.CatchHandler[] catchHandlers = code.getCatchHandlers();
828         codeOut.writeUnsignedShort(tries.length);
829 
830         int debugInfoOffset = code.getDebugInfoOffset();
831         if (debugInfoOffset != 0) {
832             codeOut.writeInt(debugInfoOut.getPosition());
833             transformDebugInfoItem(in.open(debugInfoOffset), indexMap);
834         } else {
835             codeOut.writeInt(0);
836         }
837 
838         short[] instructions = code.getInstructions();
839         InstructionTransformer transformer = (in == dexA)
840                 ? aInstructionTransformer
841                 : bInstructionTransformer;
842         short[] newInstructions = transformer.transform(instructions);
843         codeOut.writeInt(newInstructions.length);
844         codeOut.write(newInstructions);
845 
846         if (tries.length > 0) {
847             if (newInstructions.length % 2 == 1) {
848                 codeOut.writeShort((short) 0); // padding
849             }
850 
851             /*
852              * We can't write the tries until we've written the catch handlers.
853              * Unfortunately they're in the opposite order in the dex file so we
854              * need to transform them out-of-order.
855              */
856             Dex.Section triesSection = dexOut.open(codeOut.getPosition());
857             codeOut.skip(tries.length * SizeOf.TRY_ITEM);
858             int[] offsets = transformCatchHandlers(indexMap, catchHandlers);
859             transformTries(triesSection, tries, offsets);
860         }
861     }
862 
863     /**
864      * Writes the catch handlers to {@code codeOut} and returns their indices.
865      */
transformCatchHandlers(IndexMap indexMap, Code.CatchHandler[] catchHandlers)866     private int[] transformCatchHandlers(IndexMap indexMap, Code.CatchHandler[] catchHandlers) {
867         int baseOffset = codeOut.getPosition();
868         codeOut.writeUleb128(catchHandlers.length);
869         int[] offsets = new int[catchHandlers.length];
870         for (int i = 0; i < catchHandlers.length; i++) {
871             offsets[i] = codeOut.getPosition() - baseOffset;
872             transformEncodedCatchHandler(catchHandlers[i], indexMap);
873         }
874         return offsets;
875     }
876 
transformTries(Dex.Section out, Code.Try[] tries, int[] catchHandlerOffsets)877     private void transformTries(Dex.Section out, Code.Try[] tries,
878             int[] catchHandlerOffsets) {
879         for (Code.Try tryItem : tries) {
880             out.writeInt(tryItem.getStartAddress());
881             out.writeUnsignedShort(tryItem.getInstructionCount());
882             out.writeUnsignedShort(catchHandlerOffsets[tryItem.getCatchHandlerIndex()]);
883         }
884     }
885 
886     private static final byte DBG_END_SEQUENCE = 0x00;
887     private static final byte DBG_ADVANCE_PC = 0x01;
888     private static final byte DBG_ADVANCE_LINE = 0x02;
889     private static final byte DBG_START_LOCAL = 0x03;
890     private static final byte DBG_START_LOCAL_EXTENDED = 0x04;
891     private static final byte DBG_END_LOCAL = 0x05;
892     private static final byte DBG_RESTART_LOCAL = 0x06;
893     private static final byte DBG_SET_PROLOGUE_END = 0x07;
894     private static final byte DBG_SET_EPILOGUE_BEGIN = 0x08;
895     private static final byte DBG_SET_FILE = 0x09;
896 
transformDebugInfoItem(Dex.Section in, IndexMap indexMap)897     private void transformDebugInfoItem(Dex.Section in, IndexMap indexMap) {
898         contentsOut.debugInfos.size++;
899         int lineStart = in.readUleb128();
900         debugInfoOut.writeUleb128(lineStart);
901 
902         int parametersSize = in.readUleb128();
903         debugInfoOut.writeUleb128(parametersSize);
904 
905         for (int p = 0; p < parametersSize; p++) {
906             int parameterName = in.readUleb128p1();
907             debugInfoOut.writeUleb128p1(indexMap.adjustString(parameterName));
908         }
909 
910         int addrDiff;    // uleb128   address delta.
911         int lineDiff;    // sleb128   line delta.
912         int registerNum; // uleb128   register number.
913         int nameIndex;   // uleb128p1 string index.    Needs indexMap adjustment.
914         int typeIndex;   // uleb128p1 type index.      Needs indexMap adjustment.
915         int sigIndex;    // uleb128p1 string index.    Needs indexMap adjustment.
916 
917         while (true) {
918             int opcode = in.readByte();
919             debugInfoOut.writeByte(opcode);
920 
921             switch (opcode) {
922             case DBG_END_SEQUENCE:
923                 return;
924 
925             case DBG_ADVANCE_PC:
926                 addrDiff = in.readUleb128();
927                 debugInfoOut.writeUleb128(addrDiff);
928                 break;
929 
930             case DBG_ADVANCE_LINE:
931                 lineDiff = in.readSleb128();
932                 debugInfoOut.writeSleb128(lineDiff);
933                 break;
934 
935             case DBG_START_LOCAL:
936             case DBG_START_LOCAL_EXTENDED:
937                 registerNum = in.readUleb128();
938                 debugInfoOut.writeUleb128(registerNum);
939                 nameIndex = in.readUleb128p1();
940                 debugInfoOut.writeUleb128p1(indexMap.adjustString(nameIndex));
941                 typeIndex = in.readUleb128p1();
942                 debugInfoOut.writeUleb128p1(indexMap.adjustType(typeIndex));
943                 if (opcode == DBG_START_LOCAL_EXTENDED) {
944                     sigIndex = in.readUleb128p1();
945                     debugInfoOut.writeUleb128p1(indexMap.adjustString(sigIndex));
946                 }
947                 break;
948 
949             case DBG_END_LOCAL:
950             case DBG_RESTART_LOCAL:
951                 registerNum = in.readUleb128();
952                 debugInfoOut.writeUleb128(registerNum);
953                 break;
954 
955             case DBG_SET_FILE:
956                 nameIndex = in.readUleb128p1();
957                 debugInfoOut.writeUleb128p1(indexMap.adjustString(nameIndex));
958                 break;
959 
960             case DBG_SET_PROLOGUE_END:
961             case DBG_SET_EPILOGUE_BEGIN:
962             default:
963                 break;
964             }
965         }
966     }
967 
transformEncodedCatchHandler(Code.CatchHandler catchHandler, IndexMap indexMap)968     private void transformEncodedCatchHandler(Code.CatchHandler catchHandler, IndexMap indexMap) {
969         int catchAllAddress = catchHandler.getCatchAllAddress();
970         int[] typeIndexes = catchHandler.getTypeIndexes();
971         int[] addresses = catchHandler.getAddresses();
972 
973         if (catchAllAddress != -1) {
974             codeOut.writeSleb128(-typeIndexes.length);
975         } else {
976             codeOut.writeSleb128(typeIndexes.length);
977         }
978 
979         for (int i = 0; i < typeIndexes.length; i++) {
980             codeOut.writeUleb128(indexMap.adjustType(typeIndexes[i]));
981             codeOut.writeUleb128(addresses[i]);
982         }
983 
984         if (catchAllAddress != -1) {
985             codeOut.writeUleb128(catchAllAddress);
986         }
987     }
988 
transformStaticValues(Dex.Section in, IndexMap indexMap)989     private void transformStaticValues(Dex.Section in, IndexMap indexMap) {
990         contentsOut.encodedArrays.size++;
991         indexMap.putStaticValuesOffset(in.getPosition(), encodedArrayOut.getPosition());
992         indexMap.adjustEncodedArray(in.readEncodedArray()).writeTo(encodedArrayOut);
993     }
994 
995     /**
996      * Byte counts for the sections written when creating a dex. Target sizes
997      * are defined in one of two ways:
998      * <ul>
999      * <li>By pessimistically guessing how large the union of dex files will be.
1000      *     We're pessimistic because we can't predict the amount of duplication
1001      *     between dex files, nor can we predict the length of ULEB-encoded
1002      *     offsets or indices.
1003      * <li>By exactly measuring an existing dex.
1004      * </ul>
1005      */
1006     private static class WriterSizes {
1007         private int header = SizeOf.HEADER_ITEM;
1008         private int idsDefs;
1009         private int mapList;
1010         private int typeList;
1011         private int classData;
1012         private int code;
1013         private int stringData;
1014         private int debugInfo;
1015         private int encodedArray;
1016         private int annotationsDirectory;
1017         private int annotationsSet;
1018         private int annotationsSetRefList;
1019         private int annotation;
1020 
1021         /**
1022          * Compute sizes for merging a and b.
1023          */
WriterSizes(Dex a, Dex b)1024         public WriterSizes(Dex a, Dex b) {
1025             plus(a.getTableOfContents(), false);
1026             plus(b.getTableOfContents(), false);
1027             fourByteAlign();
1028         }
1029 
WriterSizes(DexMerger dexMerger)1030         public WriterSizes(DexMerger dexMerger) {
1031             header = dexMerger.headerOut.used();
1032             idsDefs = dexMerger.idsDefsOut.used();
1033             mapList = dexMerger.mapListOut.used();
1034             typeList = dexMerger.typeListOut.used();
1035             classData = dexMerger.classDataOut.used();
1036             code = dexMerger.codeOut.used();
1037             stringData = dexMerger.stringDataOut.used();
1038             debugInfo = dexMerger.debugInfoOut.used();
1039             encodedArray = dexMerger.encodedArrayOut.used();
1040             annotationsDirectory = dexMerger.annotationsDirectoryOut.used();
1041             annotationsSet = dexMerger.annotationSetOut.used();
1042             annotationsSetRefList = dexMerger.annotationSetRefListOut.used();
1043             annotation = dexMerger.annotationOut.used();
1044             fourByteAlign();
1045         }
1046 
plus(TableOfContents contents, boolean exact)1047         private void plus(TableOfContents contents, boolean exact) {
1048             idsDefs += contents.stringIds.size * SizeOf.STRING_ID_ITEM
1049                     + contents.typeIds.size * SizeOf.TYPE_ID_ITEM
1050                     + contents.protoIds.size * SizeOf.PROTO_ID_ITEM
1051                     + contents.fieldIds.size * SizeOf.MEMBER_ID_ITEM
1052                     + contents.methodIds.size * SizeOf.MEMBER_ID_ITEM
1053                     + contents.classDefs.size * SizeOf.CLASS_DEF_ITEM;
1054             mapList = SizeOf.UINT + (contents.sections.length * SizeOf.MAP_ITEM);
1055             typeList += fourByteAlign(contents.typeLists.byteCount); // We count each dex's
1056             // typelists section as realigned on 4 bytes, because each typelist of each dex's
1057             // typelists section is aligned on 4 bytes. If we didn't, there is a case where each
1058             // size of both dex's typelists section is a multiple of 2 but not a multiple of 4,
1059             // and the sum of both sizes is a multiple of 4 but would not be sufficient to write
1060             // each typelist aligned on 4 bytes.
1061             stringData += contents.stringDatas.byteCount;
1062             annotationsDirectory += contents.annotationsDirectories.byteCount;
1063             annotationsSet += contents.annotationSets.byteCount;
1064             annotationsSetRefList += contents.annotationSetRefLists.byteCount;
1065 
1066             if (exact) {
1067                 code += contents.codes.byteCount;
1068                 classData += contents.classDatas.byteCount;
1069                 encodedArray += contents.encodedArrays.byteCount;
1070                 annotation += contents.annotations.byteCount;
1071                 debugInfo += contents.debugInfos.byteCount;
1072             } else {
1073                 // at most 1/4 of the bytes in a code section are uleb/sleb
1074                 code += (int) Math.ceil(contents.codes.byteCount * 1.25);
1075                 // at most 1/3 of the bytes in a class data section are uleb/sleb
1076                 classData += (int) Math.ceil(contents.classDatas.byteCount * 1.34);
1077                 // all of the bytes in an encoding arrays section may be uleb/sleb
1078                 encodedArray += contents.encodedArrays.byteCount * 2;
1079                 // all of the bytes in an annotations section may be uleb/sleb
1080                 annotation += (int) Math.ceil(contents.annotations.byteCount * 2);
1081                 // all of the bytes in a debug info section may be uleb/sleb
1082                 debugInfo += contents.debugInfos.byteCount * 2;
1083             }
1084         }
1085 
fourByteAlign()1086         private void fourByteAlign() {
1087             header = fourByteAlign(header);
1088             idsDefs = fourByteAlign(idsDefs);
1089             mapList = fourByteAlign(mapList);
1090             typeList = fourByteAlign(typeList);
1091             classData = fourByteAlign(classData);
1092             code = fourByteAlign(code);
1093             stringData = fourByteAlign(stringData);
1094             debugInfo = fourByteAlign(debugInfo);
1095             encodedArray = fourByteAlign(encodedArray);
1096             annotationsDirectory = fourByteAlign(annotationsDirectory);
1097             annotationsSet = fourByteAlign(annotationsSet);
1098             annotationsSetRefList = fourByteAlign(annotationsSetRefList);
1099             annotation = fourByteAlign(annotation);
1100         }
1101 
fourByteAlign(int position)1102         private static int fourByteAlign(int position) {
1103             return (position + 3) & ~3;
1104         }
1105 
size()1106         public int size() {
1107             return header + idsDefs + mapList + typeList + classData + code + stringData + debugInfo
1108                     + encodedArray + annotationsDirectory + annotationsSet + annotationsSetRefList
1109                     + annotation;
1110         }
1111     }
1112 
main(String[] args)1113     public static void main(String[] args) throws IOException {
1114         if (args.length < 2) {
1115             printUsage();
1116             return;
1117         }
1118 
1119         Dex merged = new Dex(new File(args[1]));
1120         for (int i = 2; i < args.length; i++) {
1121             Dex toMerge = new Dex(new File(args[i]));
1122             merged = new DexMerger(merged, toMerge, CollisionPolicy.KEEP_FIRST).merge();
1123         }
1124         merged.writeTo(new File(args[0]));
1125     }
1126 
printUsage()1127     private static void printUsage() {
1128         System.out.println("Usage: DexMerger <out.dex> <a.dex> <b.dex> ...");
1129         System.out.println();
1130         System.out.println(
1131             "If a class is defined in several dex, the class found in the first dex will be used.");
1132     }
1133 }
1134