Use dedicated per-file maps instead of field on IndexedDexItem for
file membership.

This reduces the memory usage and number of allocations, especially
when many files are generated.

Bug: 67286151
Change-Id: I5324b7eb2f07136fc0c1b4fd7d3793e4dbfbb570
diff --git a/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java b/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java
index 752c66b..c8e3aa9 100644
--- a/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java
+++ b/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java
@@ -32,8 +32,8 @@
 import java.io.PrintStream;
 import java.io.PrintWriter;
 import java.io.Writer;
+import java.util.Collection;
 import java.util.LinkedHashMap;
-import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
@@ -229,7 +229,7 @@
    * be used.
    */
   private static void rewriteCodeWithJumboStrings(ObjectToOffsetMapping mapping,
-      List<DexProgramClass> classes, DexApplication application) {
+      Collection<DexProgramClass> classes, DexApplication application) {
     // If there are no strings with jumbo indices at all this is a no-op.
     if (!mapping.hasJumboStrings()) {
       return;
diff --git a/src/main/java/com/android/tools/r8/dex/Constants.java b/src/main/java/com/android/tools/r8/dex/Constants.java
index 3a06f1d..50af413 100644
--- a/src/main/java/com/android/tools/r8/dex/Constants.java
+++ b/src/main/java/com/android/tools/r8/dex/Constants.java
@@ -135,7 +135,6 @@
   public static final String CLASS_INITIALIZER_NAME = "<clinit>";
 
   public static final int MAX_NON_JUMBO_INDEX = U16BIT_MAX;
-  public static final int FIRST_JUMBO_INDEX = MAX_NON_JUMBO_INDEX + 1;
 
   public static final int KILOBYTE = 1 << 10;
 }
diff --git a/src/main/java/com/android/tools/r8/dex/FileWriter.java b/src/main/java/com/android/tools/r8/dex/FileWriter.java
index 70fd655..020580f 100644
--- a/src/main/java/com/android/tools/r8/dex/FileWriter.java
+++ b/src/main/java/com/android/tools/r8/dex/FileWriter.java
@@ -38,6 +38,7 @@
 import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.graph.DexTypeList;
 import com.android.tools.r8.graph.DexValue;
+import com.android.tools.r8.graph.IndexedDexItem;
 import com.android.tools.r8.graph.KeyedDexItem;
 import com.android.tools.r8.graph.ObjectToOffsetMapping;
 import com.android.tools.r8.graph.PresortedComparable;
@@ -316,7 +317,7 @@
     }
   }
 
-  private <T extends DexItem> void writeFixedSectionItems(T[] items, int offset,
+  private <T extends IndexedDexItem> void writeFixedSectionItems(Collection<T> items, int offset,
       ThrowingConsumer<T, ApiLevelException> writer) throws ApiLevelException {
     assert dest.position() == offset;
     for (T item : items) {
@@ -324,6 +325,14 @@
     }
   }
 
+  private void writeFixedSectionItems(DexProgramClass[] items, int offset,
+      ThrowingConsumer<DexProgramClass, ApiLevelException> writer) throws ApiLevelException {
+    assert dest.position() == offset;
+    for (DexProgramClass item : items) {
+      writer.accept(item);
+    }
+  }
+
   private <T extends DexItem> void writeItems(Collection<T> items, Consumer<Integer> offsetSetter,
       Consumer<T> writer) {
     writeItems(items, offsetSetter, writer, 1);
@@ -691,21 +700,21 @@
     int size = 0;
     size += writeMapItem(Constants.TYPE_HEADER_ITEM, 0, 1);
     size += writeMapItem(Constants.TYPE_STRING_ID_ITEM, layout.stringIdsOffset,
-        mapping.getStrings().length);
+        mapping.getStrings().size());
     size += writeMapItem(Constants.TYPE_TYPE_ID_ITEM, layout.typeIdsOffset,
-        mapping.getTypes().length);
+        mapping.getTypes().size());
     size += writeMapItem(Constants.TYPE_PROTO_ID_ITEM, layout.protoIdsOffset,
-        mapping.getProtos().length);
+        mapping.getProtos().size());
     size += writeMapItem(Constants.TYPE_FIELD_ID_ITEM, layout.fieldIdsOffset,
-        mapping.getFields().length);
+        mapping.getFields().size());
     size += writeMapItem(Constants.TYPE_METHOD_ID_ITEM, layout.methodIdsOffset,
-        mapping.getMethods().length);
+        mapping.getMethods().size());
     size += writeMapItem(Constants.TYPE_CLASS_DEF_ITEM, layout.classDefsOffset,
         mapping.getClasses().length);
     size += writeMapItem(Constants.TYPE_CALL_SITE_ID_ITEM, layout.callSiteIdsOffset,
-        mapping.getCallSites().length);
+        mapping.getCallSites().size());
     size += writeMapItem(Constants.TYPE_METHOD_HANDLE_ITEM, layout.methodHandleIdsOffset,
-        mapping.getMethodHandles().length);
+        mapping.getMethodHandles().size());
     size += writeMapItem(Constants.TYPE_CODE_ITEM, layout.getCodesOffset(),
         mixedSectionOffsets.getCodes().size());
     size += writeMapItem(Constants.TYPE_DEBUG_INFO_ITEM, layout.getDebugInfosOffset(),
@@ -750,19 +759,19 @@
     dest.putInt(0);
     dest.putInt(0);
     dest.putInt(layout.getMapOffset());
-    int numberOfStrings = mapping.getStrings().length;
+    int numberOfStrings = mapping.getStrings().size();
     dest.putInt(numberOfStrings);
     dest.putInt(numberOfStrings == 0 ? 0 : layout.stringIdsOffset);
-    int numberOfTypes = mapping.getTypes().length;
+    int numberOfTypes = mapping.getTypes().size();
     dest.putInt(numberOfTypes);
     dest.putInt(numberOfTypes == 0 ? 0 : layout.typeIdsOffset);
-    int numberOfProtos = mapping.getProtos().length;
+    int numberOfProtos = mapping.getProtos().size();
     dest.putInt(numberOfProtos);
     dest.putInt(numberOfProtos == 0 ? 0 : layout.protoIdsOffset);
-    int numberOfFields = mapping.getFields().length;
+    int numberOfFields = mapping.getFields().size();
     dest.putInt(numberOfFields);
     dest.putInt(numberOfFields == 0 ? 0 : layout.fieldIdsOffset);
-    int numberOfMethods = mapping.getMethods().length;
+    int numberOfMethods = mapping.getMethods().size();
     dest.putInt(numberOfMethods);
     dest.putInt(numberOfMethods == 0 ? 0 : layout.methodIdsOffset);
     int numberOfClasses = mapping.getClasses().length;
@@ -853,14 +862,14 @@
       int offset = 0;
       return new Layout(
           offset = Constants.TYPE_HEADER_ITEM_SIZE,
-          offset += mapping.getStrings().length * Constants.TYPE_STRING_ID_ITEM_SIZE,
-          offset += mapping.getTypes().length * Constants.TYPE_TYPE_ID_ITEM_SIZE,
-          offset += mapping.getProtos().length * Constants.TYPE_PROTO_ID_ITEM_SIZE,
-          offset += mapping.getFields().length * Constants.TYPE_FIELD_ID_ITEM_SIZE,
-          offset += mapping.getMethods().length * Constants.TYPE_METHOD_ID_ITEM_SIZE,
+          offset += mapping.getStrings().size() * Constants.TYPE_STRING_ID_ITEM_SIZE,
+          offset += mapping.getTypes().size() * Constants.TYPE_TYPE_ID_ITEM_SIZE,
+          offset += mapping.getProtos().size() * Constants.TYPE_PROTO_ID_ITEM_SIZE,
+          offset += mapping.getFields().size() * Constants.TYPE_FIELD_ID_ITEM_SIZE,
+          offset += mapping.getMethods().size() * Constants.TYPE_METHOD_ID_ITEM_SIZE,
           offset += mapping.getClasses().length * Constants.TYPE_CLASS_DEF_ITEM_SIZE,
-          offset += mapping.getCallSites().length * Constants.TYPE_CALL_SITE_ID_ITEM_SIZE,
-          offset += mapping.getMethodHandles().length * Constants.TYPE_METHOD_HANDLE_ITEM_SIZE);
+          offset += mapping.getCallSites().size() * Constants.TYPE_CALL_SITE_ID_ITEM_SIZE,
+          offset += mapping.getMethodHandles().size() * Constants.TYPE_METHOD_HANDLE_ITEM_SIZE);
     }
 
     int getDataSectionSize() {
diff --git a/src/main/java/com/android/tools/r8/dex/VirtualFile.java b/src/main/java/com/android/tools/r8/dex/VirtualFile.java
index eeb5023..7eb399d 100644
--- a/src/main/java/com/android/tools/r8/dex/VirtualFile.java
+++ b/src/main/java/com/android/tools/r8/dex/VirtualFile.java
@@ -17,7 +17,6 @@
 import com.android.tools.r8.graph.DexProto;
 import com.android.tools.r8.graph.DexString;
 import com.android.tools.r8.graph.DexType;
-import com.android.tools.r8.graph.IndexedDexItem;
 import com.android.tools.r8.graph.ObjectToOffsetMapping;
 import com.android.tools.r8.naming.ClassNameMapper;
 import com.android.tools.r8.naming.NamingLens;
@@ -74,7 +73,7 @@
 
   private VirtualFile(int id, NamingLens namingLens, DexProgramClass primaryClass) {
     this.id = id;
-    this.indexedItems = new VirtualFileIndexedItemCollection(id);
+    this.indexedItems = new VirtualFileIndexedItemCollection(namingLens);
     this.transaction = new IndexedItemTransaction(indexedItems, namingLens);
     this.primaryClass = primaryClass;
   }
@@ -150,16 +149,15 @@
   public ObjectToOffsetMapping computeMapping(DexApplication application) {
     assert transaction.isEmpty();
     return new ObjectToOffsetMapping(
-        id,
         application,
-        indexedItems.classes.toArray(new DexProgramClass[indexedItems.classes.size()]),
-        indexedItems.protos.toArray(new DexProto[indexedItems.protos.size()]),
-        indexedItems.types.toArray(new DexType[indexedItems.types.size()]),
-        indexedItems.methods.toArray(new DexMethod[indexedItems.methods.size()]),
-        indexedItems.fields.toArray(new DexField[indexedItems.fields.size()]),
-        indexedItems.strings.toArray(new DexString[indexedItems.strings.size()]),
-        indexedItems.callSites.toArray(new DexCallSite[indexedItems.callSites.size()]),
-        indexedItems.methodHandles.toArray(new DexMethodHandle[indexedItems.methodHandles.size()]));
+        indexedItems.classes,
+        indexedItems.protos,
+        indexedItems.types,
+        indexedItems.methods,
+        indexedItems.fields,
+        indexedItems.strings,
+        indexedItems.callSites,
+        indexedItems.methodHandles);
   }
 
   private void addClass(DexProgramClass clazz) {
@@ -204,7 +202,7 @@
     return indexedItems.classes.isEmpty();
   }
 
-  public List<DexProgramClass> classes() {
+  public Collection<DexProgramClass> classes() {
     return indexedItems.classes;
   }
 
@@ -397,87 +395,90 @@
 
   private static class VirtualFileIndexedItemCollection implements IndexedItemCollection {
 
-    final int id;
+    private final NamingLens namingLens;
 
-    private final List<DexProgramClass> classes = new ArrayList<>();
-    private final List<DexProto> protos = new ArrayList<>();
-    private final List<DexType> types = new ArrayList<>();
-    private final List<DexMethod> methods = new ArrayList<>();
-    private final List<DexField> fields = new ArrayList<>();
-    private final List<DexString> strings = new ArrayList<>();
-    private final List<DexCallSite> callSites = new ArrayList<>();
-    private final List<DexMethodHandle> methodHandles = new ArrayList<>();
+    private final Set<DexProgramClass> classes = Sets.newIdentityHashSet();
+    private final Set<DexProto> protos = Sets.newIdentityHashSet();
+    private final Set<DexType> types = Sets.newIdentityHashSet();
+    private final Set<DexMethod> methods = Sets.newIdentityHashSet();
+    private final Set<DexField> fields = Sets.newIdentityHashSet();
+    private final Set<DexString> strings = Sets.newIdentityHashSet();
+    private final Set<DexCallSite> callSites = Sets.newIdentityHashSet();
+    private final Set<DexMethodHandle> methodHandles = Sets.newIdentityHashSet();
 
-    private final Set<DexClass> seenClasses = Sets.newIdentityHashSet();
+    public VirtualFileIndexedItemCollection(
+        NamingLens namingLens) {
+      this.namingLens = namingLens;
 
-    private VirtualFileIndexedItemCollection(int id) {
-      this.id = id;
-    }
-
-    private <T extends IndexedDexItem> boolean addItem(T item, List<T> itemList) {
-      assert item != null;
-      if (item.assignToVirtualFile(id)) {
-        itemList.add(item);
-        return true;
-      }
-      return false;
     }
 
     @Override
     public boolean addClass(DexProgramClass clazz) {
-      if (seenClasses.add(clazz)) {
-        classes.add(clazz);
-        return true;
-      }
-      return false;
+      return classes.add(clazz);
     }
 
     @Override
     public boolean addField(DexField field) {
-      return addItem(field, fields);
+      return fields.add(field);
     }
 
     @Override
     public boolean addMethod(DexMethod method) {
-      return addItem(method, methods);
+      return methods.add(method);
     }
 
     @Override
     public boolean addString(DexString string) {
-      return addItem(string, strings);
+      return strings.add(string);
     }
 
     @Override
     public boolean addProto(DexProto proto) {
-      return addItem(proto, protos);
-    }
-
-    @Override
-    public boolean addCallSite(DexCallSite callSite) {
-      return addItem(callSite, callSites);
-    }
-
-    @Override
-    public boolean addMethodHandle(DexMethodHandle methodHandle) {
-      return addItem(methodHandle, methodHandles);
+      return protos.add(proto);
     }
 
     @Override
     public boolean addType(DexType type) {
-      return addItem(type, types);
+      return types.add(type);
     }
 
-    public int getNumberOfMethods() {
+    @Override
+    public boolean addCallSite(DexCallSite callSite) {
+      return callSites.add(callSite);
+    }
+
+    @Override
+    public boolean addMethodHandle(DexMethodHandle methodHandle) {
+      return methodHandles.add(methodHandle);
+    }
+
+    int getNumberOfMethods() {
       return methods.size();
     }
 
-    public int getNumberOfFields() {
+    int getNumberOfFields() {
       return fields.size();
     }
 
-    public int getNumberOfStrings() {
+    int getNumberOfStrings() {
       return strings.size();
     }
+
+    @Override
+    public DexString getRenamedDescriptor(DexType type) {
+      return namingLens.lookupDescriptor(type);
+    }
+
+    @Override
+    public DexString getRenamedName(DexMethod method) {
+      assert namingLens.checkTargetCanBeTranslated(method);
+      return namingLens.lookupName(method);
+    }
+
+    @Override
+    public DexString getRenamedName(DexField field) {
+      return namingLens.lookupName(field);
+    }
   }
 
   private static class IndexedItemTransaction implements IndexedItemCollection {
@@ -500,8 +501,8 @@
       this.namingLens = namingLens;
     }
 
-    private <T extends IndexedDexItem> boolean maybeInsert(T item, Set<T> set) {
-      if (item.hasVirtualFileData(base.id) || set.contains(item)) {
+    private <T extends DexItem> boolean maybeInsert(T item, Set<T> set, Set<T> baseSet) {
+      if (baseSet.contains(item) || set.contains(item)) {
         return false;
       }
       set.add(item);
@@ -514,46 +515,42 @@
 
     @Override
     public boolean addClass(DexProgramClass dexProgramClass) {
-      if (base.seenClasses.contains(dexProgramClass) || classes.contains(dexProgramClass)) {
-        return false;
-      }
-      classes.add(dexProgramClass);
-      return true;
+      return maybeInsert(dexProgramClass, classes, base.classes);
     }
 
     @Override
     public boolean addField(DexField field) {
-      return maybeInsert(field, fields);
+      return maybeInsert(field, fields, base.fields);
     }
 
     @Override
     public boolean addMethod(DexMethod method) {
-      return maybeInsert(method, methods);
+      return maybeInsert(method, methods, base.methods);
     }
 
     @Override
     public boolean addString(DexString string) {
-      return maybeInsert(string, strings);
+      return maybeInsert(string, strings, base.strings);
     }
 
     @Override
     public boolean addProto(DexProto proto) {
-      return maybeInsert(proto, protos);
+      return maybeInsert(proto, protos, base.protos);
     }
 
     @Override
     public boolean addType(DexType type) {
-      return maybeInsert(type, types);
+      return maybeInsert(type, types, base.types);
     }
 
     @Override
     public boolean addCallSite(DexCallSite callSite) {
-      return maybeInsert(callSite, callSites);
+      return maybeInsert(callSite, callSites, base.callSites);
     }
 
     @Override
     public boolean addMethodHandle(DexMethodHandle methodHandle) {
-      return maybeInsert(methodHandle, methodHandles);
+      return maybeInsert(methodHandle, methodHandles, base.methodHandles);
     }
 
     @Override
diff --git a/src/main/java/com/android/tools/r8/graph/DexItem.java b/src/main/java/com/android/tools/r8/graph/DexItem.java
index c0df752..88665e9 100644
--- a/src/main/java/com/android/tools/r8/graph/DexItem.java
+++ b/src/main/java/com/android/tools/r8/graph/DexItem.java
@@ -5,6 +5,7 @@
 
 import com.android.tools.r8.dex.IndexedItemCollection;
 import com.android.tools.r8.dex.MixedSectionCollection;
+import java.util.Collection;
 import java.util.function.Consumer;
 
 public abstract class DexItem {
@@ -17,6 +18,11 @@
     consumeArray(items, (T item) -> item.collectMixedSectionItems(mixedItems));
   }
 
+  public static <T extends DexItem> void collectAll(MixedSectionCollection mixedItems,
+      Collection<T> items) {
+    items.forEach((T item) -> item.collectMixedSectionItems(mixedItems));
+  }
+
   /**
    * Helper method to iterate over elements in an array.
    * Handles the case where the array is null.
diff --git a/src/main/java/com/android/tools/r8/graph/IndexedDexItem.java b/src/main/java/com/android/tools/r8/graph/IndexedDexItem.java
index 5c8f439..ec36b60 100644
--- a/src/main/java/com/android/tools/r8/graph/IndexedDexItem.java
+++ b/src/main/java/com/android/tools/r8/graph/IndexedDexItem.java
@@ -5,7 +5,6 @@
 
 import com.android.tools.r8.dex.IndexedItemCollection;
 import com.android.tools.r8.dex.MixedSectionCollection;
-import java.util.Arrays;
 
 /**
  * Subset of dex items that are referenced by some table index.
@@ -14,27 +13,6 @@
 
   private static final int SORTED_INDEX_UNKNOWN = -1;
   private int sortedIndex = SORTED_INDEX_UNKNOWN; // assigned globally after reading.
-  /**
-   * Contains the indexes assigned to this item for the various virtual output files.
-   *
-   * <p>One DexItem might be assigned to multiple virtual files.
-   *
-   * <p>For a certain virtual file this DexItem has the value:
-   * <ul>
-   * <li>{@link #UNASSOCIATED_VALUE}, when not associated with the virtual file.
-   * <li>{@link #ASSOCIATED_VALUE}, when associated with the virtual file but no index allocated.
-   * <li>A zero or greater value when this item has been associated by the virtual file
-   * and the value denotes the assigned index.
-   * </ul>
-   * <p> Note that, in case of multiple files, for a specific IndexedDexItem, we may not have
-   * as many entries in the index as there are files (we only expand when we need to). If we lookup
-   * the value of an entry that is out of bounds it is equivalent to {@link #UNASSOCIATED_VALUE}
-   *
-   * <p>This field is initialized on first write in {@link #updateVirtualFileData(int)}}. It
-   * is assumed that multiple files are processed concurrently and thus the allocation of the
-   * array is synchronized. However, for any a given file id, sequential access is assumed.
-   */
-  private int[] virtualFileIndexes;
 
   public abstract void collectIndexedItems(IndexedItemCollection indexedItems);
 
@@ -46,90 +24,6 @@
 
   public abstract int getOffset(ObjectToOffsetMapping mapping);
 
-  /**
-   * Constants used inside virtualFileIndexes.
-   */
-  public static final int UNASSOCIATED_VALUE = -2;
-  public static final int ASSOCIATED_VALUE = -1;
-  public static final int MIN_VALID_VALUE = 0;
-
-  /**
-   * Returns whether this item is assigned to the given file id.
-   */
-  public boolean hasVirtualFileData(int virtualFileId) {
-    return getVirtualFileIndex(virtualFileId) != UNASSOCIATED_VALUE;
-  }
-
-  /**
-   * Assigns this item to the given file id if it has not been assigned previously.
-   *
-   * <p>This method returns 'true' if the item was newly assigned, i.e., it was not previously
-   * assigned to the file id.
-   */
-  public boolean assignToVirtualFile(int virtualFileId) {
-    // Fast lock-free check whether already assigned.
-    if (hasVirtualFileData(virtualFileId)) {
-      return false;
-    }
-    return updateVirtualFileData(virtualFileId);
-  }
-
-  /**
-   * Assigns this item to the given file id.
-   *
-   * <p>As a side effect, the {@link #virtualFileIndexes} field might be initialized or expanded.
-   * Hence this method is synchronized. Note that the field is queried without synchronization.
-   * Therefor it has to remain in a valid state at all times and must transition atomically from
-   * null to an initialized allocated value.
-   */
-  private synchronized boolean updateVirtualFileData(int virtualFileId) {
-    if (virtualFileIndexes == null) {
-      int[] fileIndices = new int[virtualFileId + 1];
-      Arrays.fill(fileIndices, UNASSOCIATED_VALUE);
-      // This has to be an atomic transition from null to an initialized array.
-      virtualFileIndexes = fileIndices;
-    }
-    // We increased the number of files, increase the index size.
-    if (virtualFileId >= virtualFileIndexes.length) {
-      int oldLength = virtualFileIndexes.length;
-      int[] fileIndices = Arrays.copyOf(virtualFileIndexes, virtualFileId + 1);
-      Arrays.fill(fileIndices, oldLength, virtualFileId + 1, UNASSOCIATED_VALUE);
-      virtualFileIndexes = fileIndices;
-    }
-    assert virtualFileId < virtualFileIndexes.length;
-    boolean wasAdded = virtualFileIndexes[virtualFileId] == UNASSOCIATED_VALUE;
-    virtualFileIndexes[virtualFileId] = ASSOCIATED_VALUE;
-    return wasAdded;
-  }
-
-  /**
-   * Assigns an actual index for this item in the given file.
-   *
-   * <p>May only be used after this item has been assigned to the file using {@link
-   * #assignToVirtualFile(int)}.
-   */
-  public void assignVirtualFileIndex(int virtualFileId, int index) {
-    assert virtualFileIndexes != null;
-    assert virtualFileIndexes[virtualFileId] < MIN_VALID_VALUE;
-    virtualFileIndexes[virtualFileId] = index;
-  }
-
-  /**
-   * Returns the index associated with this item for the given file id or {@link
-   * #UNASSOCIATED_VALUE} if the item is not associated to the given file id.
-   */
-  public int getVirtualFileIndex(int virtualFileId) {
-    if (virtualFileIndexes == null) {
-      return UNASSOCIATED_VALUE;
-    }
-    // If more files were added, but this entry not associated with it, we would not have extended
-    // the size of the array. So if the {@link virtualFileId} is out of bounds, it means
-    // {@link #UNASSOCIATED_VALUE}
-    return virtualFileIndexes.length > virtualFileId
-        ? virtualFileIndexes[virtualFileId]
-        : UNASSOCIATED_VALUE;
-  }
-
   // Partial implementation of PresortedComparable.
 
   final public void setSortedIndex(int sortedIndex) {
diff --git a/src/main/java/com/android/tools/r8/graph/ObjectToOffsetMapping.java b/src/main/java/com/android/tools/r8/graph/ObjectToOffsetMapping.java
index c8157d6..98af471 100644
--- a/src/main/java/com/android/tools/r8/graph/ObjectToOffsetMapping.java
+++ b/src/main/java/com/android/tools/r8/graph/ObjectToOffsetMapping.java
@@ -4,36 +4,41 @@
 package com.android.tools.r8.graph;
 
 import com.android.tools.r8.dex.Constants;
+import com.android.tools.r8.errors.CompilationError;
 import com.google.common.collect.Sets;
-import java.util.Arrays;
-import java.util.Collections;
+import it.unimi.dsi.fastutil.objects.Reference2IntLinkedOpenHashMap;
+import it.unimi.dsi.fastutil.objects.Reference2IntMap;
+import java.util.Collection;
+import java.util.Comparator;
 import java.util.Set;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
 
 public class ObjectToOffsetMapping {
 
-  private final int virtualFileId;
+  private final static int NOT_FOUND = -1;
+  private final static int NOT_SET = -2;
 
   private final DexProgramClass[] classes;
-  private final DexProto[] protos;
-  private final DexType[] types;
-  private final DexMethod[] methods;
-  private final DexField[] fields;
-  private final DexString[] strings;
-  private final DexCallSite[] callSites;
-  private final DexMethodHandle[] methodHandles;
+  private final Reference2IntMap<DexProto> protos;
+  private final Reference2IntMap<DexType> types;
+  private final Reference2IntMap<DexMethod> methods;
+  private final Reference2IntMap<DexField> fields;
+  private final Reference2IntMap<DexString> strings;
+  private final Reference2IntMap<DexCallSite> callSites;
+  private final Reference2IntMap<DexMethodHandle> methodHandles;
   private DexString firstJumboString;
 
   public ObjectToOffsetMapping(
-      int virtualFileId,
       DexApplication application,
-      DexProgramClass[] classes,
-      DexProto[] protos,
-      DexType[] types,
-      DexMethod[] methods,
-      DexField[] fields,
-      DexString[] strings,
-      DexCallSite[] callSites,
-      DexMethodHandle[] methodHandles) {
+      Collection<DexProgramClass> classes,
+      Collection<DexProto> protos,
+      Collection<DexType> types,
+      Collection<DexMethod> methods,
+      Collection<DexField> fields,
+      Collection<DexString> strings,
+      Collection<DexCallSite> callSites,
+      Collection<DexMethodHandle> methodHandles) {
     assert application != null;
     assert classes != null;
     assert protos != null;
@@ -44,91 +49,83 @@
     assert callSites != null;
     assert methodHandles != null;
 
-    this.virtualFileId = virtualFileId;
     this.classes = sortClasses(application, classes);
-    this.protos = protos;
-    this.types = types;
-    this.methods = methods;
-    this.fields = fields;
-    this.strings = strings;
-    this.callSites = callSites;
-    this.methodHandles = methodHandles;
-
-    Arrays.sort(protos);
-    setIndexes(protos);
-
-    Arrays.sort(types);
-    setIndexes(types);
-
-    Arrays.sort(methods);
-    setIndexes(methods);
-
-    Arrays.sort(fields);
-    setIndexes(fields);
-
-    Arrays.sort(strings);
-    setIndexes(strings);
-
+    this.protos = createMap(protos, true, this::failOnOverflow);
+    this.types = createMap(types, true, this::failOnOverflow);
+    this.methods = createMap(methods, true, this::failOnOverflow);
+    this.fields = createMap(fields, true, this::failOnOverflow);
+    this.strings = createMap(strings, true, this::setFirstJumboString);
     // No need to sort CallSite, they will be written in data section in the callSites order,
     // consequently offset of call site used into the call site section will be in ascending order.
-    setIndexes(callSites);
-
+    this.callSites = createMap(callSites, false, this::failOnOverflow);
     // No need to sort method handle
-    setIndexes(methodHandles);
+    this.methodHandles = createMap(methodHandles, false, this::failOnOverflow);
   }
 
-  private static DexProgramClass[] sortClasses(
-      DexApplication application, DexProgramClass[] classes) {
-    Arrays.sort(classes, (o1, o2) -> o1.type.descriptor.slowCompareTo(o2.type.descriptor));
+  private void setFirstJumboString(DexString string) {
+    assert firstJumboString == null;
+    firstJumboString = string;
+  }
+
+  private void failOnOverflow(DexItem item) {
+    throw new CompilationError("Index overflow for " + item.getClass());
+  }
+
+  private <T extends IndexedDexItem> Reference2IntMap<T> createMap(Collection<T> items,
+      boolean sort, Consumer<T> onUInt16Overflow) {
+    Reference2IntMap<T> map = new Reference2IntLinkedOpenHashMap<>();
+    map.defaultReturnValue(NOT_FOUND);
+    Collection<T> sorted = sort ? items.stream().sorted().collect(Collectors.toList()) : items;
+    int index = 0;
+    for (T item : sorted) {
+      if (index == Constants.U16BIT_MAX + 1) {
+        onUInt16Overflow.accept(item);
+      }
+      map.put(item, index++);
+    }
+    return map;
+  }
+
+  private static DexProgramClass[] sortClasses(DexApplication application,
+      Collection<DexProgramClass> classes) {
     SortingProgramClassVisitor classVisitor = new SortingProgramClassVisitor(application, classes);
-    classVisitor.run(classes);
+    // Collect classes in subtyping order, based on a sorted list of classes to start with.
+    classVisitor.run(
+        classes.stream().sorted(Comparator.comparing(DexClass::getType))
+            .collect(Collectors.toList()));
     return classVisitor.getSortedClasses();
   }
 
-  private void setIndexes(IndexedDexItem[] items) {
-    int index = 0;
-    for (IndexedDexItem item : items) {
-      item.assignVirtualFileIndex(virtualFileId, index);
-      // For strings collect the first jumbo string (if any).
-      if ((index > Constants.MAX_NON_JUMBO_INDEX) && (item instanceof DexString)) {
-        if (index == Constants.FIRST_JUMBO_INDEX) {
-          firstJumboString = (DexString) item;
-        }
-      }
-      index++;
-    }
-  }
-
-  public DexMethod[] getMethods() {
-    return methods;
+  public Collection<DexMethod> getMethods() {
+    return methods.keySet();
   }
 
   public DexProgramClass[] getClasses() {
     return classes;
   }
 
-  public DexType[] getTypes() {
-    return types;
+  public Collection<DexType> getTypes() {
+    return types.keySet();
   }
 
-  public DexProto[] getProtos() {
-    return protos;
+  public Collection<DexProto> getProtos() {
+    return protos.keySet();
   }
 
-  public DexField[] getFields() {
-    return fields;
+  public Collection<DexField> getFields() {
+    return fields.keySet();
   }
 
-  public DexString[] getStrings() {
-    return strings;
+  public Collection<DexString> getStrings() {
+    return strings.keySet();
   }
 
-  public DexCallSite[] getCallSites() {
-    return callSites;
+  public Collection<DexCallSite> getCallSites() {
+    return callSites.keySet();
   }
 
-  public DexMethodHandle[] getMethodHandles() {
-    return methodHandles;
+  public Collection<DexMethodHandle> getMethodHandles() {
+    return methodHandles.keySet();
   }
 
   public boolean hasJumboStrings() {
@@ -139,43 +136,39 @@
     return firstJumboString;
   }
 
-  private boolean isContainedInMapping(IndexedDexItem item) {
-    return item.getVirtualFileIndex(virtualFileId) != IndexedDexItem.UNASSOCIATED_VALUE;
+  private <T extends IndexedDexItem> int getOffsetFor(T item, Reference2IntMap<T> map) {
+    int index = map.getInt(item);
+    assert index != NOT_SET : "Index was not set: " + item;
+    assert index != NOT_FOUND : "Missing dependency: " + item;
+    return index;
   }
 
   public int getOffsetFor(DexProto proto) {
-    assert isContainedInMapping(proto) : "Missing dependency: " + proto;
-    return proto.getVirtualFileIndex(virtualFileId);
+    return getOffsetFor(proto, protos);
   }
 
   public int getOffsetFor(DexField field) {
-    assert isContainedInMapping(field) : "Missing dependency: " + field;
-    return field.getVirtualFileIndex(virtualFileId);
+    return getOffsetFor(field, fields);
   }
 
   public int getOffsetFor(DexMethod method) {
-    assert isContainedInMapping(method) : "Missing dependency: " + method;
-    return method.getVirtualFileIndex(virtualFileId);
+    return getOffsetFor(method, methods);
   }
 
   public int getOffsetFor(DexString string) {
-    assert isContainedInMapping(string) : "Missing dependency: " + string;
-    return string.getVirtualFileIndex(virtualFileId);
+    return getOffsetFor(string, strings);
   }
 
   public int getOffsetFor(DexType type) {
-    assert isContainedInMapping(type) : "Missing dependency: " + type;
-    return type.getVirtualFileIndex(virtualFileId);
+    return getOffsetFor(type, types);
   }
 
   public int getOffsetFor(DexCallSite callSite) {
-    assert isContainedInMapping(callSite) : "Missing dependency: " + callSite;
-    return callSite.getVirtualFileIndex(virtualFileId);
+    return getOffsetFor(callSite, callSites);
   }
 
   public int getOffsetFor(DexMethodHandle methodHandle) {
-    assert isContainedInMapping(methodHandle) : "Missing dependency: " + methodHandle;
-    return methodHandle.getVirtualFileIndex(virtualFileId);
+    return getOffsetFor(methodHandle, methodHandles);
   }
 
   private static class SortingProgramClassVisitor extends ProgramClassVisitor {
@@ -184,10 +177,11 @@
 
     private int index = 0;
 
-    public SortingProgramClassVisitor(DexApplication application, DexProgramClass[] classes) {
+    public SortingProgramClassVisitor(DexApplication application,
+        Collection<DexProgramClass> classes) {
       super(application);
-      this.sortedClasses = new DexProgramClass[classes.length];
-      Collections.addAll(classSet, classes);
+      this.sortedClasses = new DexProgramClass[classes.size()];
+      classSet.addAll(classes);
     }
 
     @Override
diff --git a/src/main/java/com/android/tools/r8/graph/ProgramClassVisitor.java b/src/main/java/com/android/tools/r8/graph/ProgramClassVisitor.java
index 826e260..3414037 100644
--- a/src/main/java/com/android/tools/r8/graph/ProgramClassVisitor.java
+++ b/src/main/java/com/android/tools/r8/graph/ProgramClassVisitor.java
@@ -68,12 +68,16 @@
     }
   }
 
-  public void run() {
-    for (DexProgramClass clazz : application.classes()) {
+  public void run(Iterable<DexProgramClass> classes) {
+    for (DexProgramClass clazz : classes) {
       accept(clazz);
     }
   }
 
+  public void run() {
+    run(application.classes());
+  }
+
   /**
    * Called for each library class used in the class hierarchy. A library class is a class that is
    * not present in the application.
diff --git a/src/test/java/com/android/tools/r8/dex/DebugByteCodeWriterTest.java b/src/test/java/com/android/tools/r8/dex/DebugByteCodeWriterTest.java
index a08f274..9060e55 100644
--- a/src/test/java/com/android/tools/r8/dex/DebugByteCodeWriterTest.java
+++ b/src/test/java/com/android/tools/r8/dex/DebugByteCodeWriterTest.java
@@ -4,41 +4,33 @@
 package com.android.tools.r8.dex;
 
 import com.android.tools.r8.graph.DexApplication;
-import com.android.tools.r8.graph.DexApplication.Builder;
-import com.android.tools.r8.graph.DexCallSite;
 import com.android.tools.r8.graph.DexDebugEvent;
 import com.android.tools.r8.graph.DexDebugInfo;
-import com.android.tools.r8.graph.DexField;
 import com.android.tools.r8.graph.DexItemFactory;
-import com.android.tools.r8.graph.DexMethod;
-import com.android.tools.r8.graph.DexMethodHandle;
-import com.android.tools.r8.graph.DexProgramClass;
-import com.android.tools.r8.graph.DexProto;
 import com.android.tools.r8.graph.DexString;
-import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.graph.ObjectToOffsetMapping;
+import java.util.Collections;
 import org.junit.Assert;
 import org.junit.Test;
 
 public class DebugByteCodeWriterTest {
 
-  ObjectToOffsetMapping emptyObjectTObjectMapping() {
+  private ObjectToOffsetMapping emptyObjectTObjectMapping() {
     return new ObjectToOffsetMapping(
-        0,
         DexApplication.builder(new DexItemFactory(), null).build(),
-        new DexProgramClass[] {},
-        new DexProto[] {},
-        new DexType[] {},
-        new DexMethod[] {},
-        new DexField[] {},
-        new DexString[] {},
-        new DexCallSite[] {},
-        new DexMethodHandle[] {});
+        Collections.emptyList(),
+        Collections.emptyList(),
+        Collections.emptyList(),
+        Collections.emptyList(),
+        Collections.emptyList(),
+        Collections.emptyList(),
+        Collections.emptyList(),
+        Collections.emptyList());
   }
 
   @Test
   public void testEmptyDebugInfo() {
-    DexDebugInfo debugInfo = new DexDebugInfo(1, new DexString[]{}, new DexDebugEvent[]{});
+    DexDebugInfo debugInfo = new DexDebugInfo(1, DexString.EMPTY_ARRAY, new DexDebugEvent[]{});
     DebugBytecodeWriter writer = new DebugBytecodeWriter(debugInfo, emptyObjectTObjectMapping());
     Assert.assertEquals(3, writer.generate().length);
   }
diff --git a/src/test/java/com/android/tools/r8/jasmin/JumboStringTests.java b/src/test/java/com/android/tools/r8/jasmin/JumboStringTests.java
index 61b624e..27d8d8b 100644
--- a/src/test/java/com/android/tools/r8/jasmin/JumboStringTests.java
+++ b/src/test/java/com/android/tools/r8/jasmin/JumboStringTests.java
@@ -23,7 +23,7 @@
   // String constants are split into several class files to ensure both the constant-pool and
   // instruction count are below the class-file limits.
   private static int CLASSES_COUNT = 10;
-  private static int MIN_STRING_COUNT = Constants.FIRST_JUMBO_INDEX + 1;
+  private static int MIN_STRING_COUNT = Constants.MAX_NON_JUMBO_INDEX + 2;
   private static int EXTRA_STRINGS_PER_CLASSES_COUNT = MIN_STRING_COUNT % CLASSES_COUNT;
   private static int STRINGS_PER_CLASSES_COUNT =
       EXTRA_STRINGS_PER_CLASSES_COUNT + MIN_STRING_COUNT / CLASSES_COUNT;