Change filePerClass semanthic to filePerInputClass

Meaning that synthesized classes are emitted in the individual dex files
of the input classes they were generated from. Shared synthetic classes
may then by emitted in several individual dex files, those duplicates
will then be removed by merging.

Bug: 65725257

Change-Id: Ib1032c81dcaa08ab9c490caa45d8243fce3fbdf6
diff --git a/src/main/java/com/android/tools/r8/BaseCompilerCommand.java b/src/main/java/com/android/tools/r8/BaseCompilerCommand.java
index 5e8f975..7645242 100644
--- a/src/main/java/com/android/tools/r8/BaseCompilerCommand.java
+++ b/src/main/java/com/android/tools/r8/BaseCompilerCommand.java
@@ -136,7 +136,7 @@
 
     protected void validate() throws CompilationException {
       super.validate();
-      if (getAppBuilder().hasMainDexList() && outputMode == OutputMode.FilePerClass) {
+      if (getAppBuilder().hasMainDexList() && outputMode == OutputMode.FilePerInputClass) {
         throw new CompilationException(
             "Option --main-dex-list cannot be used with --file-per-class");
       }
diff --git a/src/main/java/com/android/tools/r8/D8Command.java b/src/main/java/com/android/tools/r8/D8Command.java
index 76779bb..f2aebc7 100644
--- a/src/main/java/com/android/tools/r8/D8Command.java
+++ b/src/main/java/com/android/tools/r8/D8Command.java
@@ -158,7 +158,7 @@
           builder.setMode(CompilationMode.RELEASE);
           modeSet = CompilationMode.RELEASE;
         } else if (arg.equals("--file-per-class")) {
-          builder.setOutputMode(OutputMode.FilePerClass);
+          builder.setOutputMode(OutputMode.FilePerInputClass);
         } else if (arg.equals("--output")) {
           String output = args[++i];
           if (outputPath != null) {
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 46ee85c..58c3a5b 100644
--- a/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java
+++ b/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java
@@ -138,10 +138,10 @@
 
       // Distribute classes into dex files.
       VirtualFile.Distributor distributor = null;
-      if (options.outputMode == OutputMode.FilePerClass) {
+      if (options.outputMode == OutputMode.FilePerInputClass) {
         assert packageDistribution == null :
             "Cannot combine package distribution definition with file-per-class option.";
-        distributor = new VirtualFile.FilePerClassDistributor(this);
+        distributor = new VirtualFile.FilePerInputClassDistributor(this);
       } else if (!options.canUseMultidex()
           && options.mainDexKeepRules.isEmpty()
           && application.mainDexList.isEmpty()) {
@@ -178,7 +178,15 @@
       AndroidApp.Builder builder = AndroidApp.builder();
       try {
         for (Map.Entry<VirtualFile, Future<byte[]>> entry : dexDataFutures.entrySet()) {
-          builder.addDexProgramData(entry.getValue().get(), entry.getKey().getClassDescriptors());
+          VirtualFile virtualFile = entry.getKey();
+          if (virtualFile.getPrimaryClassDescriptor() != null) {
+            builder.addDexProgramData(
+                entry.getValue().get(),
+                virtualFile.getClassDescriptors(),
+                virtualFile.getPrimaryClassDescriptor());
+          } else {
+            builder.addDexProgramData(entry.getValue().get(), virtualFile.getClassDescriptors());
+          }
         }
       } catch (InterruptedException e) {
         throw new RuntimeException("Interrupted while waiting for future.", e);
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 d820747..1e70b9b 100644
--- a/src/main/java/com/android/tools/r8/dex/FileWriter.java
+++ b/src/main/java/com/android/tools/r8/dex/FileWriter.java
@@ -617,7 +617,7 @@
     }
   }
 
-  private void writeEncodedMethods(DexEncodedMethod[] methods) {
+  private void writeEncodedMethods(DexEncodedMethod[] methods, boolean clearBodies) {
     assert isSorted(methods);
     int currentOffset = 0;
     for (DexEncodedMethod method : methods) {
@@ -633,7 +633,9 @@
         dest.putUleb128(mixedSectionOffsets.getOffsetFor(method.getCode().asDexCode()));
         // Writing the methods starts to take up memory so we are going to flush the
         // code objects since they are no longer necessary after this.
-        method.removeCode();
+        if (clearBodies) {
+          method.removeCode();
+        }
       }
     }
   }
@@ -647,8 +649,10 @@
     dest.putUleb128(clazz.virtualMethods().length);
     writeEncodedFields(clazz.staticFields());
     writeEncodedFields(clazz.instanceFields());
-    writeEncodedMethods(clazz.directMethods());
-    writeEncodedMethods(clazz.virtualMethods());
+
+    boolean isSharedSynthetic = clazz.getSynthesizedFrom().size() > 1;
+    writeEncodedMethods(clazz.directMethods(), !isSharedSynthetic);
+    writeEncodedMethods(clazz.virtualMethods(), !isSharedSynthetic);
   }
 
   private void addStaticFieldValues(DexProgramClass clazz) {
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 c40d5a5..32fb175 100644
--- a/src/main/java/com/android/tools/r8/dex/VirtualFile.java
+++ b/src/main/java/com/android/tools/r8/dex/VirtualFile.java
@@ -4,8 +4,8 @@
 package com.android.tools.r8.dex;
 
 import com.android.tools.r8.errors.CompilationError;
-import com.android.tools.r8.errors.InternalCompilerError;
 import com.android.tools.r8.errors.DexOverflowException;
+import com.android.tools.r8.errors.InternalCompilerError;
 import com.android.tools.r8.graph.DexApplication;
 import com.android.tools.r8.graph.DexCallSite;
 import com.android.tools.r8.graph.DexClass;
@@ -76,10 +76,17 @@
   private final VirtualFileIndexedItemCollection indexedItems;
   private final IndexedItemTransaction transaction;
 
+  private final DexProgramClass primaryClass;
+
   private VirtualFile(int id, NamingLens namingLens) {
+    this(id, namingLens, null);
+  }
+
+  private VirtualFile(int id, NamingLens namingLens, DexProgramClass primaryClass) {
     this.id = id;
     this.indexedItems = new VirtualFileIndexedItemCollection(id);
     this.transaction = new IndexedItemTransaction(indexedItems, namingLens);
+    this.primaryClass = primaryClass;
   }
 
   public int getId() {
@@ -95,6 +102,10 @@
     return classDescriptors;
   }
 
+  public String getPrimaryClassDescriptor() {
+    return primaryClass == null ? null : primaryClass.type.descriptor.toString();
+  }
+
   public static String deriveCommonPrefixAndSanityCheck(List<String> fileNames) {
     Iterator<String> nameIterator = fileNames.iterator();
     String first = nameIterator.next();
@@ -221,21 +232,40 @@
         throws ExecutionException, IOException, DexOverflowException;
   }
 
-  public static class FilePerClassDistributor extends Distributor {
+  /**
+   * Distribute each type to its individual virtual except for types synthesized during this
+   * compilation. Synthesized classes are emitted in the individual virtual files
+   * of the input classes they were generated from. Shared synthetic classes
+   * may then be distributed in several individual virtual files.
+   */
+  public static class FilePerInputClassDistributor extends Distributor {
 
-    FilePerClassDistributor(ApplicationWriter writer) {
+    FilePerInputClassDistributor(ApplicationWriter writer) {
       super(writer);
     }
 
     @Override
     public Map<Integer, VirtualFile> run() {
+      HashMap<DexProgramClass, VirtualFile> files = new HashMap<>();
+      Collection<DexProgramClass> synthetics = new ArrayList<>();
       // Assign dedicated virtual files for all program classes.
       for (DexProgramClass clazz : application.classes()) {
-        VirtualFile file = new VirtualFile(nameToFileMap.size(), writer.namingLens);
-        nameToFileMap.put(nameToFileMap.size(), file);
-        file.addClass(clazz);
-        file.commitTransaction();
+        if (clazz.getSynthesizedFrom().isEmpty()) {
+          VirtualFile file = new VirtualFile(nameToFileMap.size(), writer.namingLens, clazz);
+          nameToFileMap.put(nameToFileMap.size(), file);
+          file.addClass(clazz);
+          files.put(clazz, file);
+        } else {
+          synthetics.add(clazz);
+        }
       }
+      for (DexProgramClass synthetic : synthetics) {
+        for (DexProgramClass inputType : synthetic.getSynthesizedFrom()) {
+          VirtualFile file = files.get(inputType);
+          file.addClass(synthetic);
+        }
+      }
+      files.values().forEach(file -> file.commitTransaction());
       return nameToFileMap;
     }
   }
diff --git a/src/main/java/com/android/tools/r8/graph/DexProgramClass.java b/src/main/java/com/android/tools/r8/graph/DexProgramClass.java
index 1bdc56c..d27a3be 100644
--- a/src/main/java/com/android/tools/r8/graph/DexProgramClass.java
+++ b/src/main/java/com/android/tools/r8/graph/DexProgramClass.java
@@ -7,11 +7,16 @@
 import com.android.tools.r8.dex.IndexedItemCollection;
 import com.android.tools.r8.dex.MixedSectionCollection;
 import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Set;
 import java.util.function.Supplier;
 
 public class DexProgramClass extends DexClass implements Supplier<DexProgramClass> {
 
   private DexEncodedArray staticValues;
+  private final Collection<DexProgramClass> synthesizedFrom;
 
   public DexProgramClass(DexType type,
       Resource.Kind origin,
@@ -24,9 +29,36 @@
       DexEncodedField[] instanceFields,
       DexEncodedMethod[] directMethods,
       DexEncodedMethod[] virtualMethods) {
+    this(type,
+        origin,
+        accessFlags,
+        superType,
+        interfaces,
+        sourceFile,
+        classAnnotations,
+        staticFields,
+        instanceFields,
+        directMethods,
+        virtualMethods,
+        Collections.emptyList());
+  }
+
+  public DexProgramClass(DexType type,
+      Resource.Kind origin,
+      DexAccessFlags accessFlags,
+      DexType superType,
+      DexTypeList interfaces,
+      DexString sourceFile,
+      DexAnnotationSet classAnnotations,
+      DexEncodedField[] staticFields,
+      DexEncodedField[] instanceFields,
+      DexEncodedMethod[] directMethods,
+      DexEncodedMethod[] virtualMethods,
+      Collection<DexProgramClass> synthesizedDirectlyFrom) {
     super(sourceFile, interfaces, accessFlags, superType, type, staticFields,
         instanceFields, directMethods, virtualMethods, classAnnotations, origin);
     assert classAnnotations != null;
+    this.synthesizedFrom = accumulateSynthesizedFrom(new HashSet<>(), synthesizedDirectlyFrom);
   }
 
   @Override
@@ -54,6 +86,10 @@
     }
   }
 
+  public Collection<DexProgramClass> getSynthesizedFrom() {
+    return synthesizedFrom;
+  }
+
   @Override
   void collectMixedSectionItems(MixedSectionCollection mixedItems) {
     if (hasAnnotations()) {
@@ -129,6 +165,19 @@
     return methods != null && Arrays.stream(methods).anyMatch(DexEncodedMethod::hasAnnotation);
   }
 
+  private static Collection<DexProgramClass> accumulateSynthesizedFrom(
+      Set<DexProgramClass> accumulated,
+      Collection<DexProgramClass> toAccumulate) {
+    for (DexProgramClass dexProgramClass : toAccumulate) {
+      if (dexProgramClass.synthesizedFrom.isEmpty()) {
+        accumulated.add(dexProgramClass);
+      } else {
+        accumulateSynthesizedFrom(accumulated, dexProgramClass.synthesizedFrom);
+      }
+    }
+    return accumulated;
+  }
+
   public void setStaticValues(DexEncodedArray staticValues) {
     this.staticValues = staticValues;
   }
diff --git a/src/main/java/com/android/tools/r8/ir/desugar/InterfaceMethodRewriter.java b/src/main/java/com/android/tools/r8/ir/desugar/InterfaceMethodRewriter.java
index aa93aa8..3d7d254 100644
--- a/src/main/java/com/android/tools/r8/ir/desugar/InterfaceMethodRewriter.java
+++ b/src/main/java/com/android/tools/r8/ir/desugar/InterfaceMethodRewriter.java
@@ -60,7 +60,8 @@
 //       forward the call to an appropriate method in interface companion class.
 //
 public final class InterfaceMethodRewriter {
-  private static final String COMPANION_CLASS_NAME_SUFFIX = "-CC";
+  // public for testing
+  public static final String COMPANION_CLASS_NAME_SUFFIX = "-CC";
   private static final String DEFAULT_METHOD_PREFIX = "$default$";
 
   private final IRConverter converter;
diff --git a/src/main/java/com/android/tools/r8/ir/desugar/InterfaceProcessor.java b/src/main/java/com/android/tools/r8/ir/desugar/InterfaceProcessor.java
index 2d06abe..02d6dc9 100644
--- a/src/main/java/com/android/tools/r8/ir/desugar/InterfaceProcessor.java
+++ b/src/main/java/com/android/tools/r8/ir/desugar/InterfaceProcessor.java
@@ -17,6 +17,7 @@
 import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.graph.DexTypeList;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Map;
@@ -135,7 +136,8 @@
         DexEncodedField.EMPTY_ARRAY,
         DexEncodedField.EMPTY_ARRAY,
         companionMethods.toArray(new DexEncodedMethod[companionMethods.size()]),
-        DexEncodedMethod.EMPTY_ARRAY
+        DexEncodedMethod.EMPTY_ARRAY,
+        Collections.singletonList(iface)
     );
     companionClasses.put(iface, companionClass);
   }
diff --git a/src/main/java/com/android/tools/r8/ir/desugar/LambdaClass.java b/src/main/java/com/android/tools/r8/ir/desugar/LambdaClass.java
index f03a619..0fe0f99 100644
--- a/src/main/java/com/android/tools/r8/ir/desugar/LambdaClass.java
+++ b/src/main/java/com/android/tools/r8/ir/desugar/LambdaClass.java
@@ -27,6 +27,8 @@
 import com.android.tools.r8.graph.DexValue.DexValueNull;
 import com.android.tools.r8.ir.code.Invoke;
 import com.android.tools.r8.ir.synthetic.SynthesizedCode;
+import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicBoolean;
 
@@ -58,6 +60,7 @@
   final DexField instanceField;
   final Target target;
   final AtomicBoolean addToMainDexList = new AtomicBoolean(false);
+  private Collection<DexProgramClass> synthesizedFrom = new ArrayList<DexProgramClass>(1);
 
   LambdaClass(LambdaRewriter rewriter, DexType accessedFrom,
       DexType lambdaClassType, LambdaDescriptor descriptor) {
@@ -125,7 +128,8 @@
         synthesizeStaticFields(),
         synthesizeInstanceFields(),
         synthesizeDirectMethods(),
-        synthesizeVirtualMethods()
+        synthesizeVirtualMethods(),
+        synthesizedFrom
     );
   }
 
@@ -138,6 +142,11 @@
     return descriptor.isStateless();
   }
 
+  synchronized void addSynthesizedFrom(DexProgramClass synthesizedFrom) {
+    assert synthesizedFrom != null;
+    this.synthesizedFrom.add(synthesizedFrom);
+  }
+
   // Synthesize virtual methods.
   private DexEncodedMethod[] synthesizeVirtualMethods() {
     DexEncodedMethod[] methods = new DexEncodedMethod[1 + descriptor.bridges.size()];
diff --git a/src/main/java/com/android/tools/r8/ir/desugar/LambdaRewriter.java b/src/main/java/com/android/tools/r8/ir/desugar/LambdaRewriter.java
index 86ad391..c31675f 100644
--- a/src/main/java/com/android/tools/r8/ir/desugar/LambdaRewriter.java
+++ b/src/main/java/com/android/tools/r8/ir/desugar/LambdaRewriter.java
@@ -54,7 +54,8 @@
   private static final String METAFACTORY_ALT_METHOD_NAME = "altMetafactory";
   private static final String DESERIALIZE_LAMBDA_METHOD_NAME = "$deserializeLambda$";
 
-  static final String LAMBDA_CLASS_NAME_PREFIX = "-$$Lambda$";
+  // public for testing
+  public static final String LAMBDA_CLASS_NAME_PREFIX = "-$$Lambda$";
   static final String EXPECTED_LAMBDA_METHOD_PREFIX = "lambda$";
   static final String LAMBDA_INSTANCE_FIELD_NAME = "INSTANCE";
 
@@ -238,6 +239,7 @@
       lambdaClass = putIfAbsent(knownLambdaClasses, lambdaClassType,
           new LambdaClass(this, accessedFrom, lambdaClassType, descriptor));
     }
+    lambdaClass.addSynthesizedFrom(appInfo.definitionFor(accessedFrom).asProgramClass());
     if (isInMainDexList(accessedFrom)) {
       lambdaClass.addToMainDexList.set(true);
     }
diff --git a/src/main/java/com/android/tools/r8/utils/AndroidApp.java b/src/main/java/com/android/tools/r8/utils/AndroidApp.java
index 54d10a0..9d31763 100644
--- a/src/main/java/com/android/tools/r8/utils/AndroidApp.java
+++ b/src/main/java/com/android/tools/r8/utils/AndroidApp.java
@@ -14,8 +14,10 @@
 import com.android.tools.r8.dex.VDexFile;
 import com.android.tools.r8.dex.VDexFileReader;
 import com.android.tools.r8.errors.CompilationError;
+import com.android.tools.r8.errors.Unreachable;
 import com.android.tools.r8.shaking.FilteredClassPath;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.io.ByteStreams;
 import com.google.common.io.Closer;
 import java.io.ByteArrayOutputStream;
@@ -34,7 +36,9 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -51,6 +55,7 @@
   public static final String DEFAULT_PROGUARD_MAP_FILE = "proguard.map";
 
   private final ImmutableList<Resource> programResources;
+  private final ImmutableMap<Resource, String> programResourcesMainDescriptor;
   private final ImmutableList<ClassFileResourceProvider> classpathResourceProviders;
   private final ImmutableList<ClassFileResourceProvider> libraryResourceProviders;
 
@@ -66,6 +71,7 @@
   // See factory methods and AndroidApp.Builder below.
   private AndroidApp(
       ImmutableList<Resource> programResources,
+      ImmutableMap<Resource, String> programResourcesMainDescriptor,
       ImmutableList<ProgramFileArchiveReader> programFileArchiveReaders,
       ImmutableList<ClassFileResourceProvider> classpathResourceProviders,
       ImmutableList<ClassFileResourceProvider> libraryResourceProviders,
@@ -77,6 +83,7 @@
       List<String> mainDexClasses,
       Resource mainDexListOutput) {
     this.programResources = programResources;
+    this.programResourcesMainDescriptor = programResourcesMainDescriptor;
     this.programFileArchiveReaders = programFileArchiveReaders;
     this.classpathResourceProviders = classpathResourceProviders;
     this.libraryResourceProviders = libraryResourceProviders;
@@ -331,7 +338,7 @@
     try (Closer closer = Closer.create()) {
       List<Resource> dexProgramSources = getDexProgramResources();
       for (int i = 0; i < dexProgramSources.size(); i++) {
-        Path filePath = directory.resolve(outputMode.getOutputPath(dexProgramSources.get(i), i));
+        Path filePath = directory.resolve(getOutputPath(outputMode, dexProgramSources.get(i), i));
         if (!Files.exists(filePath.getParent())) {
           Files.createDirectories(filePath.getParent());
         }
@@ -340,6 +347,19 @@
     }
   }
 
+  private String getOutputPath(OutputMode outputMode, Resource resource, int index) {
+    switch (outputMode) {
+      case Indexed:
+        return index == 0 ? "classes.dex" : ("classes" + (index + 1) + ".dex");
+      case FilePerInputClass:
+        String classDescriptor = programResourcesMainDescriptor.get(resource);
+        assert classDescriptor!= null && DescriptorUtils.isClassDescriptor(classDescriptor);
+        return classDescriptor.substring(1, classDescriptor.length() - 1) + ".dex";
+      default:
+        throw new Unreachable("Unknown output mode: " + outputMode);
+    }
+  }
+
   private static boolean isClassesDexFile(Path file) {
     String name = file.getFileName().toString().toLowerCase();
     if (!name.startsWith("classes") || !name.endsWith(".dex")) {
@@ -389,7 +409,7 @@
       try (ZipOutputStream out = new ZipOutputStream(Files.newOutputStream(archive, options))) {
         List<Resource> dexProgramSources = getDexProgramResources();
         for (int i = 0; i < dexProgramSources.size(); i++) {
-          ZipEntry zipEntry = new ZipEntry(outputMode.getOutputPath(dexProgramSources.get(i), i));
+          ZipEntry zipEntry = new ZipEntry(getOutputPath(outputMode, dexProgramSources.get(i), i));
           byte[] bytes =
               ByteStreams.toByteArray(closer.register(dexProgramSources.get(i).getStream()));
           zipEntry.setSize(bytes.length);
@@ -426,12 +446,17 @@
     out.write(ByteStreams.toByteArray(input));
   }
 
+  String getPrimaryClassDescriptor(Resource resource) {
+    return programResourcesMainDescriptor.get(resource);
+  }
+
   /**
    * Builder interface for constructing an AndroidApp.
    */
   public static class Builder {
 
     private final List<Resource> programResources = new ArrayList<>();
+    private final Map<Resource, String> programResourcesMainDescriptor = new HashMap<>();
     private final List<ProgramFileArchiveReader> programFileArchiveReaders = new ArrayList<>();
     private final List<ClassFileResourceProvider> classpathResourceProviders = new ArrayList<>();
     private final List<ClassFileResourceProvider> libraryResourceProviders = new ArrayList<>();
@@ -564,6 +589,19 @@
     }
 
     /**
+     * Add dex program-data with class descriptor and primary class.
+     */
+    public Builder addDexProgramData(
+        byte[] data,
+        Set<String> classDescriptors,
+        String primaryClassDescriptor) {
+      Resource resource = Resource.fromBytes(Resource.Kind.DEX, data, classDescriptors);
+      programResources.add(resource);
+      programResourcesMainDescriptor.put(resource, primaryClassDescriptor);
+      return this;
+    }
+
+    /**
      * Add dex program-data.
      */
     public Builder addDexProgramData(byte[]... data) {
@@ -727,6 +765,7 @@
     public AndroidApp build() {
       return new AndroidApp(
           ImmutableList.copyOf(programResources),
+          ImmutableMap.copyOf(programResourcesMainDescriptor),
           ImmutableList.copyOf(programFileArchiveReaders),
           ImmutableList.copyOf(classpathResourceProviders),
           ImmutableList.copyOf(libraryResourceProviders),
diff --git a/src/main/java/com/android/tools/r8/utils/OutputMode.java b/src/main/java/com/android/tools/r8/utils/OutputMode.java
index 1bdf501..4520cd6 100644
--- a/src/main/java/com/android/tools/r8/utils/OutputMode.java
+++ b/src/main/java/com/android/tools/r8/utils/OutputMode.java
@@ -3,28 +3,8 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.utils;
 
-import com.android.tools.r8.Resource;
-import java.util.Set;
-
 /** Defines way the output is formed. */
 public enum OutputMode {
-  Indexed {
-    @Override
-    String getOutputPath(Resource resource, int index) {
-      return index == 0 ? "classes.dex" : ("classes" + (index + 1) + ".dex");
-    }
-  },
-  FilePerClass {
-    @Override
-    String getOutputPath(Resource resource, int index) {
-      Set<String> classDescriptors = resource.getClassDescriptors();
-      assert classDescriptors != null;
-      assert classDescriptors.size() == 1;
-      String classDescriptor = classDescriptors.iterator().next();
-      assert DescriptorUtils.isClassDescriptor(classDescriptor);
-      return classDescriptor.substring(1, classDescriptor.length() - 1) + ".dex";
-    }
-  };
-
-  abstract String getOutputPath(Resource resource, int index);
+  Indexed,
+  FilePerInputClass;
 }
diff --git a/src/test/java/com/android/tools/r8/D8IncrementalRunExamplesAndroidOTest.java b/src/test/java/com/android/tools/r8/D8IncrementalRunExamplesAndroidOTest.java
index 6461a57..711453b 100644
--- a/src/test/java/com/android/tools/r8/D8IncrementalRunExamplesAndroidOTest.java
+++ b/src/test/java/com/android/tools/r8/D8IncrementalRunExamplesAndroidOTest.java
@@ -11,9 +11,13 @@
 import com.android.tools.r8.errors.CompilationError;
 import com.android.tools.r8.errors.InternalCompilerError;
 import com.android.tools.r8.errors.Unimplemented;
+import com.android.tools.r8.ir.desugar.InterfaceMethodRewriter;
+import com.android.tools.r8.ir.desugar.LambdaRewriter;
 import com.android.tools.r8.utils.AndroidApp;
+import com.android.tools.r8.utils.DescriptorUtils;
 import com.android.tools.r8.utils.OffOrAuto;
 import com.android.tools.r8.utils.OutputMode;
+import com.android.tools.r8.utils.UtilsHelper;
 import com.beust.jcommander.internal.Lists;
 import com.google.common.io.Closer;
 import java.io.ByteArrayOutputStream;
@@ -77,13 +81,19 @@
       TreeMap<String, Resource> fileToResource = new TreeMap<>();
       List<String> classFiles = collectClassFiles(testJarFile);
       AndroidApp app = compileClassFiles(
-          testJarFile, classFiles, output, OutputMode.FilePerClass);
+          testJarFile, classFiles, output, OutputMode.FilePerInputClass);
       for (Resource resource : app.getDexProgramResources()) {
         Set<String> descriptors = resource.getClassDescriptors();
-        Assert.assertNotNull(descriptors);
-        Assert.assertEquals(1, descriptors.size());
-        String classDescriptor = descriptors.iterator().next();
-        classDescriptor = classDescriptor.substring(1, classDescriptor.length() - 1);
+        String mainClassDescriptor = UtilsHelper.getMainClassDescriptor(app, resource);
+        for (String descriptor : descriptors) {
+          // classes are either lambda classes used by the main class, companion classes of the main
+          // interface or the main class/interface
+          Assert.assertTrue(descriptor.contains(LambdaRewriter.LAMBDA_CLASS_NAME_PREFIX)
+              || descriptor.endsWith(InterfaceMethodRewriter.COMPANION_CLASS_NAME_SUFFIX + ";")
+              || descriptor.equals(mainClassDescriptor));
+        }
+        String classDescriptor =
+            DescriptorUtils.getClassBinaryNameFromDescriptor(mainClassDescriptor);
         String classFilePath = classDescriptor + ".class";
         if (File.separatorChar != '/') {
           classFilePath = classFilePath.replace('/', File.separatorChar);
diff --git a/src/test/java/com/android/tools/r8/maindexlist/MainDexTracingTest.java b/src/test/java/com/android/tools/r8/maindexlist/MainDexTracingTest.java
index ba5092f..cc56417 100644
--- a/src/test/java/com/android/tools/r8/maindexlist/MainDexTracingTest.java
+++ b/src/test/java/com/android/tools/r8/maindexlist/MainDexTracingTest.java
@@ -15,6 +15,7 @@
 import com.android.tools.r8.ToolHelper;
 import com.android.tools.r8.graph.DexClass;
 import com.android.tools.r8.graph.DexType;
+import com.android.tools.r8.ir.desugar.LambdaRewriter;
 import com.android.tools.r8.utils.DescriptorUtils;
 import com.android.tools.r8.utils.InternalOptions;
 import java.nio.charset.StandardCharsets;
@@ -179,7 +180,7 @@
   }
 
   private boolean isLambda(String mainDexEntry) {
-    return mainDexEntry.contains("-$$Lambda$");
+    return mainDexEntry.contains(LambdaRewriter.LAMBDA_CLASS_NAME_PREFIX);
   }
 
   private String mainDexStringToDescriptor(String mainDexString) {
diff --git a/src/test/java/com/android/tools/r8/utils/OutputModeTest.java b/src/test/java/com/android/tools/r8/utils/OutputModeTest.java
deleted file mode 100644
index 5929b60..0000000
--- a/src/test/java/com/android/tools/r8/utils/OutputModeTest.java
+++ /dev/null
@@ -1,34 +0,0 @@
-// Copyright (c) 2017, the R8 project authors. Please see the AUTHORS file
-// for details. All rights reserved. Use of this source code is governed by a
-// BSD-style license that can be found in the LICENSE file.
-package com.android.tools.r8.utils;
-
-import static org.junit.Assert.assertEquals;
-
-import com.android.tools.r8.Resource;
-import java.util.Collections;
-import org.junit.Test;
-
-public class OutputModeTest {
-  @Test
-  public void testIndexedFileName() {
-    assertEquals("classes.dex", OutputMode.Indexed.getOutputPath(null, 0));
-    assertEquals("classes2.dex", OutputMode.Indexed.getOutputPath(null, 1));
-  }
-
-  @Test
-  public void testFilePerClass() {
-    Resource test =
-        Resource.fromBytes(Resource.Kind.CLASSFILE, new byte[]{}, Collections.singleton("LTest;"));
-    assertEquals("Test.dex", OutputMode.FilePerClass.getOutputPath(test, 0));
-    Resource comTest =
-        Resource.fromBytes(
-            Resource.Kind.CLASSFILE, new byte[]{}, Collections.singleton("Lcom/Test;"));
-    assertEquals("com/Test.dex", OutputMode.FilePerClass.getOutputPath(comTest, 0));
-    Resource comExampleTest =
-        Resource.fromBytes(
-            Resource.Kind.CLASSFILE, new byte[]{}, Collections.singleton("Lcom/example/Test;"));
-    assertEquals("com/example/Test.dex", OutputMode.FilePerClass.getOutputPath(comExampleTest, 0));
-    assertEquals("com/example/Test.dex", OutputMode.FilePerClass.getOutputPath(comExampleTest, 1));
-  }
-}