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