dx/src/com/android/dx/merge/DexMerger.java

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