Fixes unnecessary multi-merge steps.

Changes merging from quadratic in the number of classes to linear.
This has a tremendous speed up while merging many dexes at the same time.
A sample test (iosched app) with 29 dexes goes from 6 seconds to 1.

Change-Id: Iff02a0dc44d098b0878e88d18f5f4083804a6495

8 files changed, 163 insertions, 157 deletions

diff --git a/dx/src/com/android/dx/command/dexer/Main.java b/dx/src/com/android/dx/command/dexer/Main.java
index 099264ca7..e93c6a32c 100644
--- a/dx/src/com/android/dx/command/dexer/Main.java
+++ b/dx/src/com/android/dx/command/dexer/Main.java
@@ -478,7 +478,7 @@ public class Main {
         } else if (dexB == null) {
             result = dexA;
         } else {
-            result = new DexMerger(dexA, dexB, CollisionPolicy.KEEP_FIRST).merge();
+            result = new DexMerger(new Dex[] {dexA, dexB}, CollisionPolicy.KEEP_FIRST).merge();
         }
 
         ByteArrayOutputStream bytesOut = new ByteArrayOutputStream();
@@ -491,19 +491,18 @@ public class Main {
      * same type, this fails with an exception.
      */
     private static byte[] mergeLibraryDexBuffers(byte[] outArray) throws IOException {
+        ArrayList<Dex> dexes = new ArrayList<Dex>();
+        if (outArray != null) {
+            dexes.add(new Dex(outArray));
+        }
         for (byte[] libraryDex : libraryDexBuffers) {
-            if (outArray == null) {
-                outArray = libraryDex;
-                continue;
-            }
-
-            Dex a = new Dex(outArray);
-            Dex b = new Dex(libraryDex);
-            Dex ab = new DexMerger(a, b, CollisionPolicy.FAIL).merge();
-            outArray = ab.getBytes();
+            dexes.add(new Dex(libraryDex));
         }
-
-        return outArray;
+        if (dexes.isEmpty()) {
+            return null;
+        }
+        Dex merged = new DexMerger(dexes.toArray(new Dex[dexes.size()]), CollisionPolicy.FAIL).merge();
+        return merged.getBytes();
     }
 
     /**
diff --git a/dx/src/com/android/dx/merge/DexMerger.java b/dx/src/com/android/dx/merge/DexMerger.java
index 6b167d704..041ebd2ce 100644
--- a/dx/src/com/android/dx/merge/DexMerger.java
+++ b/dx/src/com/android/dx/merge/DexMerger.java
@@ -32,17 +32,15 @@ import com.android.dex.TypeList;
 
 import java.io.File;
 import java.io.IOException;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.List;
+import java.util.*;
 
 /**
  * Combine two dex files into one.
  */
 public final class DexMerger {
-    private final Dex dexA;
-    private final Dex dexB;
+    private final Dex[] dexes;
+    private final IndexMap[] indexMaps;
+
     private final CollisionPolicy collisionPolicy;
     private final WriterSizes writerSizes;
 
@@ -81,34 +79,29 @@ public final class DexMerger {
 
     private final TableOfContents contentsOut;
 
-    private final IndexMap aIndexMap;
-    private final IndexMap bIndexMap;
-    private final InstructionTransformer aInstructionTransformer;
-    private final InstructionTransformer bInstructionTransformer;
+    private final InstructionTransformer instructionTransformer;
 
     /** minimum number of wasted bytes before it's worthwhile to compact the result */
     private int compactWasteThreshold = 1024 * 1024; // 1MiB
 
-    public DexMerger(Dex dexA, Dex dexB, CollisionPolicy collisionPolicy)
+    public DexMerger(Dex[] dexes, CollisionPolicy collisionPolicy)
             throws IOException {
-        this(dexA, dexB, collisionPolicy, new WriterSizes(dexA, dexB));
+        this(dexes, collisionPolicy, new WriterSizes(dexes));
     }
 
-    private DexMerger(Dex dexA, Dex dexB, CollisionPolicy collisionPolicy,
+    private DexMerger(Dex[] dexes, CollisionPolicy collisionPolicy,
             WriterSizes writerSizes) throws IOException {
-        this.dexA = dexA;
-        this.dexB = dexB;
+        this.dexes = dexes;
         this.collisionPolicy = collisionPolicy;
         this.writerSizes = writerSizes;
 
         dexOut = new Dex(writerSizes.size());
 
-        TableOfContents aContents = dexA.getTableOfContents();
-        TableOfContents bContents = dexB.getTableOfContents();
-        aIndexMap = new IndexMap(dexOut, aContents);
-        bIndexMap = new IndexMap(dexOut, bContents);
-        aInstructionTransformer = new InstructionTransformer(aIndexMap);
-        bInstructionTransformer = new InstructionTransformer(bIndexMap);
+        indexMaps = new IndexMap[dexes.length];
+        for (int i = 0; i < dexes.length; i++) {
+            indexMaps[i] = new IndexMap(dexOut, dexes[i].getTableOfContents());
+        }
+        instructionTransformer = new InstructionTransformer();
 
         headerOut = dexOut.appendSection(writerSizes.header, "header");
         idsDefsOut = dexOut.appendSection(writerSizes.idsDefs, "ids defs");
@@ -185,6 +178,12 @@ public final class DexMerger {
     }
 
     public Dex merge() throws IOException {
+        if (dexes.length == 1) {
+            return dexes[0];
+        } else if (dexes.length == 0) {
+            return null;
+        }
+
         long start = System.nanoTime();
         Dex result = mergeDexes();
 
@@ -197,7 +196,7 @@ public final class DexMerger {
         int wastedByteCount = writerSizes.size() - compactedSizes.size();
         if (wastedByteCount >  + compactWasteThreshold) {
             DexMerger compacter = new DexMerger(
-                    dexOut, new Dex(0), CollisionPolicy.FAIL, compactedSizes);
+                    new Dex[] {dexOut, new Dex(0)}, CollisionPolicy.FAIL, compactedSizes);
             result = compacter.mergeDexes();
             System.out.printf("Result compacted from %.1fKiB to %.1fKiB to save %.1fKiB%n",
                     dexOut.getLength() / 1024f,
@@ -206,12 +205,13 @@ public final class DexMerger {
         }
 
         long elapsed = System.nanoTime() - start;
-        System.out.printf("Merged dex A (%d defs/%.1fKiB) with dex B "
-                + "(%d defs/%.1fKiB). Result is %d defs/%.1fKiB. Took %.1fs%n",
-                dexA.getTableOfContents().classDefs.size,
-                dexA.getLength() / 1024f,
-                dexB.getTableOfContents().classDefs.size,
-                dexB.getLength() / 1024f,
+        for (int i = 0; i < dexes.length; i++) {
+            System.out.printf("Merged dex #%d (%d defs/%.1fKiB)%n",
+                i + 1,
+                dexes[i].getTableOfContents().classDefs.size,
+                dexes[i].getLength() / 1024f);
+        }
+        System.out.printf("Result is %d defs/%.1fKiB. Took %.1fs%n",
                 result.getTableOfContents().classDefs.size,
                 result.getLength() / 1024f,
                 elapsed / 1000000000f);
@@ -231,69 +231,60 @@ public final class DexMerger {
         }
 
         /**
-         * Merges already-sorted sections, reading only two values into memory
+         * Merges already-sorted sections, reading one value from each dex into memory
          * at a time.
          */
         public final void mergeSorted() {
-            TableOfContents.Section aSection = getSection(dexA.getTableOfContents());
-            TableOfContents.Section bSection = getSection(dexB.getTableOfContents());
+            TableOfContents.Section[] sections = new TableOfContents.Section[dexes.length];
+            Dex.Section[] dexSections = new Dex.Section[dexes.length];
+            int[] offsets = new int[dexes.length];
+            int[] indexes = new int[dexes.length];
+
+            // values contains one value from each dex, sorted for fast retrieval of
+            // the smallest value. The list associated with a value has the indexes
+            // of the dexes that had that value.
+            TreeMap<T, List<Integer>> values = new TreeMap<T, List<Integer>>();
+
+            for (int i = 0; i < dexes.length; i++) {
+                sections[i] = getSection(dexes[i].getTableOfContents());
+                dexSections[i] = sections[i].exists() ? dexes[i].open(sections[i].off) : null;
+                // Fill in values with the first value of each dex.
+                offsets[i] = readIntoMap(
+                        dexSections[i], sections[i], indexMaps[i], indexes[i], values, i);
+            }
             getSection(contentsOut).off = out.getPosition();
 
-            Dex.Section inA = aSection.exists() ? dexA.open(aSection.off) : null;
-            Dex.Section inB = bSection.exists() ? dexB.open(bSection.off) : null;
-            int aOffset = -1;
-            int bOffset = -1;
-            int aIndex = 0;
-            int bIndex = 0;
             int outCount = 0;
-            T a = null;
-            T b = null;
-
-            while (true) {
-                if (a == null && aIndex < aSection.size) {
-                    aOffset = inA.getPosition();
-                    a = read(inA, aIndexMap, aIndex);
-                }
-                if (b == null && bIndex < bSection.size) {
-                    bOffset = inB.getPosition();
-                    b = read(inB, bIndexMap, bIndex);
+            while (!values.isEmpty()) {
+                Map.Entry<T, List<Integer>> first = values.pollFirstEntry();
+                for (Integer dex : first.getValue()) {
+                    updateIndex(offsets[dex], indexMaps[dex], indexes[dex]++, outCount);
+                    // Fetch the next value of the dexes we just polled out
+                    offsets[dex] = readIntoMap(dexSections[dex], sections[dex],
+                            indexMaps[dex], indexes[dex], values, dex);
                 }
-
-                // Write the smaller of a and b. If they're equal, write only once
-                boolean advanceA;
-                boolean advanceB;
-                if (a != null && b != null) {
-                    int compare = a.compareTo(b);
-                    advanceA = compare <= 0;
-                    advanceB = compare >= 0;
-                } else {
-                    advanceA = (a != null);
-                    advanceB = (b != null);
-                }
-
-                T toWrite = null;
-                if (advanceA) {
-                    toWrite = a;
-                    updateIndex(aOffset, aIndexMap, aIndex++, outCount);
-                    a = null;
-                    aOffset = -1;
-                }
-                if (advanceB) {
-                    toWrite = b;
-                    updateIndex(bOffset, bIndexMap, bIndex++, outCount);
-                    b = null;
-                    bOffset = -1;
-                }
-                if (toWrite == null) {
-                    break; // advanceA == false && advanceB == false
-                }
-                write(toWrite);
+                write(first.getKey());
                 outCount++;
             }
 
             getSection(contentsOut).size = outCount;
         }
 
+        private int readIntoMap(Dex.Section in, TableOfContents.Section section, IndexMap indexMap,
+                                int index, TreeMap<T, List<Integer>> values, int dex) {
+            int offset = in != null ? in.getPosition() : -1;
+            if (index < section.size) {
+                T v = read(in, indexMap, index);
+                List<Integer> l = values.get(v);
+                if (l == null) {
+                    l = new ArrayList<Integer>();
+                    values.put(v, l);
+                }
+                l.add(new Integer(dex));
+            }
+            return offset;
+        }
+
         /**
          * Merges unsorted sections by reading them completely into memory and
          * sorting in memory.
@@ -302,18 +293,19 @@ public final class DexMerger {
             getSection(contentsOut).off = out.getPosition();
 
             List<UnsortedValue> all = new ArrayList<UnsortedValue>();
-            all.addAll(readUnsortedValues(dexA, aIndexMap));
-            all.addAll(readUnsortedValues(dexB, bIndexMap));
+            for (int i = 0; i < dexes.length; i++) {
+                all.addAll(readUnsortedValues(dexes[i], indexMaps[i]));
+            }
             Collections.sort(all);
 
             int outCount = 0;
             for (int i = 0; i < all.size(); ) {
                 UnsortedValue e1 = all.get(i++);
-                updateIndex(e1.offset, getIndexMap(e1.source), e1.index, outCount - 1);
+                updateIndex(e1.offset, e1.indexMap, e1.index, outCount - 1);
 
                 while (i < all.size() && e1.compareTo(all.get(i)) == 0) {
                     UnsortedValue e2 = all.get(i++);
-                    updateIndex(e2.offset, getIndexMap(e2.source), e2.index, outCount - 1);
+                    updateIndex(e2.offset, e2.indexMap, e2.index, outCount - 1);
                 }
 
                 write(e1.value);
@@ -365,16 +357,6 @@ public final class DexMerger {
         }
     }
 
-    private IndexMap getIndexMap(Dex dex) {
-        if (dex == dexA) {
-            return aIndexMap;
-        } else if (dex == dexB) {
-            return bIndexMap;
-        } else {
-            throw new IllegalArgumentException();
-        }
-    }
-
     private void mergeStringIds() {
         new IdMerger<String>(idsDefsOut) {
             @Override TableOfContents.Section getSection(TableOfContents tableOfContents) {
@@ -538,8 +520,7 @@ public final class DexMerger {
 
         for (SortableType type : types) {
             Dex in = type.getDex();
-            IndexMap indexMap = (in == dexA) ? aIndexMap : bIndexMap;
-            transformClassDef(in, type.getClassDef(), indexMap);
+            transformClassDef(in, type.getClassDef(), type.getIndexMap());
         }
     }
 
@@ -550,8 +531,9 @@ public final class DexMerger {
     private SortableType[] getSortedTypes() {
         // size is pessimistic; doesn't include arrays
         SortableType[] sortableTypes = new SortableType[contentsOut.typeIds.size];
-        readSortableTypes(sortableTypes, dexA, aIndexMap);
-        readSortableTypes(sortableTypes, dexB, bIndexMap);
+        for (int i = 0; i < dexes.length; i++) {
+            readSortableTypes(sortableTypes, dexes[i], indexMaps[i]);
+        }
 
         /*
          * Populate the depths of each sortable type. This makes D iterations
@@ -588,7 +570,8 @@ public final class DexMerger {
     private void readSortableTypes(SortableType[] sortableTypes, Dex buffer,
             IndexMap indexMap) {
         for (ClassDef classDef : buffer.classDefs()) {
-            SortableType sortableType = indexMap.adjust(new SortableType(buffer, classDef));
+            SortableType sortableType = indexMap.adjust(
+                    new SortableType(buffer, indexMap, classDef));
             int t = sortableType.getTypeIndex();
             if (sortableTypes[t] == null) {
                 sortableTypes[t] = sortableType;
@@ -606,14 +589,18 @@ public final class DexMerger {
      * We should shrink the output by merging rather than unioning
      */
     private void unionAnnotationSetsAndDirectories() {
-        transformAnnotationSets(dexA, aIndexMap);
-        transformAnnotationSets(dexB, bIndexMap);
-        transformAnnotationSetRefLists(dexA, aIndexMap);
-        transformAnnotationSetRefLists(dexB, bIndexMap);
-        transformAnnotationDirectories(dexA, aIndexMap);
-        transformAnnotationDirectories(dexB, bIndexMap);
-        transformStaticValues(dexA, aIndexMap);
-        transformStaticValues(dexB, bIndexMap);
+        for (int i = 0; i < dexes.length; i++) {
+            transformAnnotationSets(dexes[i], indexMaps[i]);
+        }
+        for (int i = 0; i < dexes.length; i++) {
+            transformAnnotationSetRefLists(dexes[i], indexMaps[i]);
+        }
+        for (int i = 0; i < dexes.length; i++) {
+            transformAnnotationDirectories(dexes[i], indexMaps[i]);
+        }
+        for (int i = 0; i < dexes.length; i++) {
+            transformStaticValues(dexes[i], indexMaps[i]);
+        }
     }
 
     private void transformAnnotationSets(Dex in, IndexMap indexMap) {
@@ -836,10 +823,7 @@ public final class DexMerger {
         }
 
         short[] instructions = code.getInstructions();
-        InstructionTransformer transformer = (in == dexA)
-                ? aInstructionTransformer
-                : bInstructionTransformer;
-        short[] newInstructions = transformer.transform(instructions);
+        short[] newInstructions = instructionTransformer.transform(indexMap, instructions);
         codeOut.writeInt(newInstructions.length);
         codeOut.write(newInstructions);
 
@@ -1019,11 +1003,12 @@ public final class DexMerger {
         private int annotation;
 
         /**
-         * Compute sizes for merging a and b.
+         * Compute sizes for merging several dexes.
          */
-        public WriterSizes(Dex a, Dex b) {
-            plus(a.getTableOfContents(), false);
-            plus(b.getTableOfContents(), false);
+        public WriterSizes(Dex[] dexes) {
+            for (int i = 0; i < dexes.length; i++) {
+                plus(dexes[i].getTableOfContents(), false);
+            }
             fourByteAlign();
         }
 
@@ -1116,11 +1101,11 @@ public final class DexMerger {
             return;
         }
 
-        Dex merged = new Dex(new File(args[1]));
-        for (int i = 2; i < args.length; i++) {
-            Dex toMerge = new Dex(new File(args[i]));
-            merged = new DexMerger(merged, toMerge, CollisionPolicy.KEEP_FIRST).merge();
+        Dex[] dexes = new Dex[args.length - 1];
+        for (int i = 1; i < args.length; i++) {
+            dexes[i - 1] = new Dex(new File(args[i]));
         }
+        Dex merged = new DexMerger(dexes, CollisionPolicy.KEEP_FIRST).merge();
         merged.writeTo(new File(args[0]));
     }
 
diff --git a/dx/src/com/android/dx/merge/IndexMap.java b/dx/src/com/android/dx/merge/IndexMap.java
index 739238a0a..f47c03ccb 100644
--- a/dx/src/com/android/dx/merge/IndexMap.java
+++ b/dx/src/com/android/dx/merge/IndexMap.java
@@ -220,7 +220,8 @@ public final class IndexMap {
     }
 
     public SortableType adjust(SortableType sortableType) {
-        return new SortableType(sortableType.getDex(), adjust(sortableType.getClassDef()));
+        return new SortableType(sortableType.getDex(),
+                sortableType.getIndexMap(), adjust(sortableType.getClassDef()));
     }
 
     public EncodedValue adjustEncodedValue(EncodedValue encodedValue) {
diff --git a/dx/src/com/android/dx/merge/InstructionTransformer.java b/dx/src/com/android/dx/merge/InstructionTransformer.java
index f8c75e558..461af0f35 100644
--- a/dx/src/com/android/dx/merge/InstructionTransformer.java
+++ b/dx/src/com/android/dx/merge/InstructionTransformer.java
@@ -24,13 +24,13 @@ import com.android.dx.io.instructions.DecodedInstruction;
 import com.android.dx.io.instructions.ShortArrayCodeOutput;
 
 final class InstructionTransformer {
-    private final IndexMap indexMap;
     private final CodeReader reader;
+
     private DecodedInstruction[] mappedInstructions;
     private int mappedAt;
+    private IndexMap indexMap;
 
-    public InstructionTransformer(IndexMap indexMap) {
-        this.indexMap = indexMap;
+    public InstructionTransformer() {
         this.reader = new CodeReader();
         this.reader.setAllVisitors(new GenericVisitor());
         this.reader.setStringVisitor(new StringVisitor());
@@ -39,11 +39,12 @@ final class InstructionTransformer {
         this.reader.setMethodVisitor(new MethodVisitor());
     }
 
-    public short[] transform(short[] encodedInstructions) throws DexException {
+    public short[] transform(IndexMap indexMap, short[] encodedInstructions) throws DexException {
         DecodedInstruction[] decodedInstructions =
             DecodedInstruction.decodeAll(encodedInstructions);
         int size = decodedInstructions.length;
 
+        this.indexMap = indexMap;
         mappedInstructions = new DecodedInstruction[size];
         mappedAt = 0;
         reader.visitAll(decodedInstructions);
@@ -55,6 +56,7 @@ final class InstructionTransformer {
             }
         }
 
+        this.indexMap = null;
         return out.getArray();
     }
 
diff --git a/dx/src/com/android/dx/merge/SortableType.java b/dx/src/com/android/dx/merge/SortableType.java
index 2ae68fb23..3389b0e0f 100644
--- a/dx/src/com/android/dx/merge/SortableType.java
+++ b/dx/src/com/android/dx/merge/SortableType.java
@@ -44,11 +44,13 @@ final class SortableType {
     };
 
     private final Dex dex;
+    private final IndexMap indexMap;
     private ClassDef classDef;
     private int depth = -1;
 
-    public SortableType(Dex dex, ClassDef classDef) {
+    public SortableType(Dex dex, IndexMap indexMap, ClassDef classDef) {
         this.dex = dex;
+        this.indexMap = indexMap;
         this.classDef = classDef;
     }
 
@@ -56,6 +58,10 @@ final class SortableType {
         return dex;
     }
 
+    public IndexMap getIndexMap() {
+        return indexMap;
+    }
+
     public ClassDef getClassDef() {
         return classDef;
     }
diff --git a/dx/tests/115-merge/com/android/dx/merge/DexMergeTest.java b/dx/tests/115-merge/com/android/dx/merge/DexMergeTest.java
index f51e2516a..e68b6cf70 100644
--- a/dx/tests/115-merge/com/android/dx/merge/DexMergeTest.java
+++ b/dx/tests/115-merge/com/android/dx/merge/DexMergeTest.java
@@ -50,10 +50,10 @@ public final class DexMergeTest extends TestCase {
                 new byte[] { 0, 1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89, -112, -23, 121 },
                 (byte[]) fillArrayData.getMethod("newByteArray").invoke(null)));
         assertTrue(Arrays.equals(
-                new char[] { 0xFFFF, 0x4321, 0xABCD, 0, 'a', 'b', 'c' },
+                new char[]{0xFFFF, 0x4321, 0xABCD, 0, 'a', 'b', 'c'},
                 (char[]) fillArrayData.getMethod("newCharArray").invoke(null)));
         assertTrue(Arrays.equals(
-                new long[] { 4660046610375530309L, 7540113804746346429L, -6246583658587674878L },
+                new long[]{4660046610375530309L, 7540113804746346429L, -6246583658587674878L},
                 (long[]) fillArrayData.getMethod("newLongArray").invoke(null)));
     }
 
@@ -120,7 +120,7 @@ public final class DexMergeTest extends TestCase {
                 + "c=@testdata.Annotated$Nested(e=, f=0, g=0), d=[])",
                 field.getAnnotation(marker).toString());
         assertEquals("@testdata.Annotated$Marker(a=on parameter, b=[], "
-                + "c=@testdata.Annotated$Nested(e=, f=0, g=0), d=[])",
+                        + "c=@testdata.Annotated$Nested(e=, f=0, g=0), d=[])",
                 method.getParameterAnnotations()[1][0].toString());
     }
 
@@ -141,11 +141,11 @@ public final class DexMergeTest extends TestCase {
 
         Dex dexA = resourceToDexBuffer("/testdata/Basic.dex");
         Dex dexB = resourceToDexBuffer("/testdata/TryCatchFinally.dex");
-        Dex merged = new DexMerger(dexA, dexB, CollisionPolicy.KEEP_FIRST).merge();
+        Dex merged = new DexMerger(new Dex[]{dexA, dexB}, CollisionPolicy.KEEP_FIRST).merge();
 
         int maxLength = 0;
         for (int i = 0; i < steps; i++) {
-            DexMerger dexMerger = new DexMerger(dexA, merged, CollisionPolicy.KEEP_FIRST);
+            DexMerger dexMerger = new DexMerger(new Dex[]{dexA, merged}, CollisionPolicy.KEEP_FIRST);
             dexMerger.setCompactWasteThreshold(compactWasteThreshold);
             merged = dexMerger.merge();
             maxLength = Math.max(maxLength, merged.getLength());
@@ -158,7 +158,7 @@ public final class DexMergeTest extends TestCase {
     public ClassLoader mergeAndLoad(String dexAResource, String dexBResource) throws Exception {
         Dex dexA = resourceToDexBuffer(dexAResource);
         Dex dexB = resourceToDexBuffer(dexBResource);
-        Dex merged = new DexMerger(dexA, dexB, CollisionPolicy.KEEP_FIRST).merge();
+        Dex merged = new DexMerger(new Dex[]{dexA, dexB}, CollisionPolicy.KEEP_FIRST).merge();
         File mergedDex = File.createTempFile("DexMergeTest", ".classes.dex");
         merged.writeTo(mergedDex);
         File mergedJar = dexToJar(mergedDex);
diff --git a/dx/tests/119-merge-conflict/com/android/dx/merge/MergeConflictTest.java b/dx/tests/119-merge-conflict/com/android/dx/merge/MergeConflictTest.java
index 0248e9677..c0713e915 100644
--- a/dx/tests/119-merge-conflict/com/android/dx/merge/MergeConflictTest.java
+++ b/dx/tests/119-merge-conflict/com/android/dx/merge/MergeConflictTest.java
@@ -28,10 +28,10 @@ public final class MergeConflictTest extends TestCase {
         Dex b = resourceToDexBuffer("/testdata/B.dex");
 
         // a and b don't overlap; this should succeed
-        Dex ab = new DexMerger(a, b, CollisionPolicy.FAIL).merge();
+        Dex ab = new DexMerger(new Dex[]{a, b}, CollisionPolicy.FAIL).merge();
 
         // a and ab overlap; this should fail
-        DexMerger dexMerger = new DexMerger(a, ab, CollisionPolicy.FAIL);
+        DexMerger dexMerger = new DexMerger(new Dex[]{a, ab}, CollisionPolicy.FAIL);
         try {
             dexMerger.merge();
             fail();
diff --git a/dx/tests/127-merge-stress/com/android/dx/merge/MergeTest.java b/dx/tests/127-merge-stress/com/android/dx/merge/MergeTest.java
index 1abdfbba0..951b2ebd5 100644
--- a/dx/tests/127-merge-stress/com/android/dx/merge/MergeTest.java
+++ b/dx/tests/127-merge-stress/com/android/dx/merge/MergeTest.java
@@ -4,9 +4,12 @@ import com.android.dex.Dex;
 import com.android.dex.DexIndexOverflowException;
 
 import java.io.File;
+import java.util.Arrays;
+import java.util.Random;
 
 /**
- * This test tries to merge given dex files at random, 2 by 2.
+ * This test tries to merge given dex files at random, a first pass at 2 by 2, followed by
+ * a second pass doing multi-way merges.
  */
 public class MergeTest {
 
@@ -14,19 +17,29 @@ public class MergeTest {
 
   public static void main(String[] args) throws Throwable {
 
-    for (int i = 0; i < NUMBER_OF_TRIES; i++) {
-      String fileName1 = args[(int) (Math.random() * args.length)];
-      String fileName2 = args[(int) (Math.random() * args.length)];
-      try {
-        Dex toMerge = new Dex(new File(fileName1));
-        Dex toMerge2 = new Dex(new File(fileName2));
-        new DexMerger(toMerge, toMerge2, CollisionPolicy.KEEP_FIRST).merge();
-      } catch (DexIndexOverflowException e) {
-        // ignore index overflow
-      } catch (Throwable t) {
-        System.err.println(
-            "Problem merging those 2 dexes: \"" + fileName1 + "\" and \"" + fileName2 + "\"");
-        throw t;
+    Random random = new Random();
+    for (int pass = 0; pass < 2; pass++) {
+      for (int i = 0; i < NUMBER_OF_TRIES; i++) {
+        // On the first pass only do 2-way merges, then do from 3 to 10 way merges
+        // but not more to avoid dex index overflow.
+        int numDex = pass == 0 ? 2 : random.nextInt(8) + 3;
+
+        String[] fileNames = new String[numDex]; // only for the error message
+        try {
+          Dex[] dexesToMerge = new Dex[numDex];
+          for (int j = 0; j < numDex; j++) {
+            String fileName = args[random.nextInt(args.length)];
+            fileNames[j] = fileName;
+            dexesToMerge[j] = new Dex(new File(fileName));
+          }
+          new DexMerger(dexesToMerge, CollisionPolicy.KEEP_FIRST).merge();
+        } catch (DexIndexOverflowException e) {
+          // ignore index overflow
+        } catch (Throwable t) {
+          System.err.println(
+                  "Problem merging those dexes: " + Arrays.toString(fileNames));
+          throw t;
+        }
       }
     }
   }