1 files changed, 20 insertions, 21 deletions
diff --git a/java/dagger/internal/codegen/base/SourceFileGenerator.java b/java/dagger/internal/codegen/base/SourceFileGenerator.java
index ada67d30d..02348a4f4 100644
--- a/java/dagger/internal/codegen/base/SourceFileGenerator.java
+++ b/java/dagger/internal/codegen/base/SourceFileGenerator.java
@@ -16,20 +16,17 @@
 
 package dagger.internal.codegen.base;
 
-
 import static com.google.auto.common.GeneratedAnnotations.generatedAnnotation;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static dagger.internal.codegen.javapoet.AnnotationSpecs.Suppression.RAWTYPES;
 import static dagger.internal.codegen.javapoet.AnnotationSpecs.Suppression.UNCHECKED;
 
 import com.google.common.base.Throwables;
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.squareup.javapoet.AnnotationSpec;
 import com.squareup.javapoet.ClassName;
 import com.squareup.javapoet.JavaFile;
 import com.squareup.javapoet.TypeSpec;
-import dagger.internal.DaggerGenerated;
 import dagger.internal.codegen.javapoet.AnnotationSpecs;
 import dagger.internal.codegen.javapoet.AnnotationSpecs.Suppression;
 import dagger.internal.codegen.langmodel.DaggerElements;
@@ -77,21 +74,22 @@ public abstract class SourceFileGenerator<T> {
 
   /** Generates a source file to be compiled for {@code T}. */
   public void generate(T input) throws SourceFileGenerationException {
-    for (TypeSpec.Builder type : topLevelTypes(input)) {
-      try {
-        buildJavaFile(input, type).writeTo(filer);
-      } catch (Exception e) {
-        // if the code above threw a SFGE, use that
-        Throwables.propagateIfPossible(e, SourceFileGenerationException.class);
-        // otherwise, throw a new one
-        throw new SourceFileGenerationException(Optional.empty(), e, originatingElement(input));
-      }
+    Optional<TypeSpec.Builder> type = write(input);
+    if (!type.isPresent()) {
+      return;
+    }
+    try {
+      buildJavaFile(input, type.get()).writeTo(filer);
+    } catch (Exception e) {
+      // if the code above threw a SFGE, use that
+      Throwables.propagateIfPossible(e, SourceFileGenerationException.class);
+      // otherwise, throw a new one
+      throw new SourceFileGenerationException(Optional.empty(), e, originatingElement(input));
     }
   }
 
   private JavaFile buildJavaFile(T input, TypeSpec.Builder typeSpecBuilder) {
     typeSpecBuilder.addOriginatingElement(originatingElement(input));
-    typeSpecBuilder.addAnnotation(DaggerGenerated.class);
     Optional<AnnotationSpec> generatedAnnotation =
         generatedAnnotation(elements, sourceVersion)
             .map(
@@ -111,9 +109,7 @@ public abstract class SourceFileGenerator<T> {
                 .build()));
 
     JavaFile.Builder javaFileBuilder =
-        JavaFile.builder(
-                elements.getPackageOf(originatingElement(input)).getQualifiedName().toString(),
-                typeSpecBuilder.build())
+        JavaFile.builder(nameGeneratedType(input).packageName(), typeSpecBuilder.build())
             .skipJavaLangImports(true);
     if (!generatedAnnotation.isPresent()) {
       javaFileBuilder.addFileComment("Generated by Dagger ($L).", GENERATED_COMMENTS);
@@ -121,16 +117,19 @@ public abstract class SourceFileGenerator<T> {
     return javaFileBuilder.build();
   }
 
+  /** Implementations should return the {@link ClassName} for the top-level type to be generated. */
+  public abstract ClassName nameGeneratedType(T input);
+
   /** Returns the originating element of the generating type. */
   public abstract Element originatingElement(T input);
 
   /**
-   * Returns {@link TypeSpec.Builder types} be generated for {@code T}, or an empty list if no types
-   * should be generated.
-   *
-   * <p>Every type will be generated in its own file.
+   * Returns a {@link TypeSpec.Builder type} to be generated for {@code T}, or {@link
+   * Optional#empty()} if no file should be generated.
    */
-  public abstract ImmutableList<TypeSpec.Builder> topLevelTypes(T input);
+  // TODO(ronshapiro): write() makes more sense in JavaWriter where all writers are mutable.
+  // consider renaming to something like typeBuilder() which conveys the mutability of the result
+  public abstract Optional<TypeSpec.Builder> write(T input);
 
   /** Returns {@link Suppression}s that are applied to files generated by this generator. */
   // TODO(b/134590785): When suppressions are removed locally, remove this and inline the usages