1 files changed, 10 insertions, 6 deletions
diff --git a/java/dagger/internal/codegen/writing/InaccessibleMapKeyProxyGenerator.java b/java/dagger/internal/codegen/writing/InaccessibleMapKeyProxyGenerator.java
index 889f8985c..d52735963 100644
--- a/java/dagger/internal/codegen/writing/InaccessibleMapKeyProxyGenerator.java
+++ b/java/dagger/internal/codegen/writing/InaccessibleMapKeyProxyGenerator.java
@@ -22,13 +22,14 @@ import static javax.lang.model.element.Modifier.FINAL;
 import static javax.lang.model.element.Modifier.PRIVATE;
 import static javax.lang.model.element.Modifier.PUBLIC;
 
-import com.google.common.collect.ImmutableList;
+import com.squareup.javapoet.ClassName;
 import com.squareup.javapoet.TypeSpec;
 import dagger.internal.codegen.base.SourceFileGenerator;
 import dagger.internal.codegen.binding.ContributionBinding;
 import dagger.internal.codegen.binding.MapKeys;
 import dagger.internal.codegen.langmodel.DaggerElements;
 import dagger.internal.codegen.langmodel.DaggerTypes;
+import java.util.Optional;
 import javax.annotation.processing.Filer;
 import javax.inject.Inject;
 import javax.lang.model.SourceVersion;
@@ -52,21 +53,24 @@ public final class InaccessibleMapKeyProxyGenerator
   }
 
   @Override
+  public ClassName nameGeneratedType(ContributionBinding binding) {
+    return MapKeys.mapKeyProxyClassName(binding);
+  }
+
+  @Override
   public Element originatingElement(ContributionBinding binding) {
     // a map key is only ever present on bindings that have a binding element
     return binding.bindingElement().get();
   }
 
   @Override
-  public ImmutableList<TypeSpec.Builder> topLevelTypes(ContributionBinding binding) {
+  public Optional<TypeSpec.Builder> write(ContributionBinding binding) {
     return MapKeys.mapKeyFactoryMethod(binding, types, elements)
         .map(
             method ->
-                classBuilder(MapKeys.mapKeyProxyClassName(binding))
+                classBuilder(nameGeneratedType(binding))
                     .addModifiers(PUBLIC, FINAL)
                     .addMethod(constructorBuilder().addModifiers(PRIVATE).build())
-                    .addMethod(method))
-        .map(ImmutableList::of)
-        .orElse(ImmutableList.of());
+                    .addMethod(method));
   }
 }