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