Add `@LazyClassKey` for injecting class keyed map and doesn't load class when unused.

RELNOTES=n/a
PiperOrigin-RevId: 596634724

5 files changed, 51 insertions, 18 deletions

diff --git a/java/dagger/internal/codegen/javapoet/CodeBlocks.java b/java/dagger/internal/codegen/javapoet/CodeBlocks.java
index 29457e267..877a0a540 100644
--- a/java/dagger/internal/codegen/javapoet/CodeBlocks.java
+++ b/java/dagger/internal/codegen/javapoet/CodeBlocks.java
@@ -19,6 +19,7 @@ package dagger.internal.codegen.javapoet;
 import static com.squareup.javapoet.MethodSpec.methodBuilder;
 import static com.squareup.javapoet.TypeSpec.anonymousClassBuilder;
 import static dagger.internal.codegen.javapoet.TypeNames.daggerProviderOf;
+import static dagger.internal.codegen.javapoet.TypeNames.lazyOf;
 import static java.util.stream.StreamSupport.stream;
 import static javax.lang.model.element.Modifier.PUBLIC;
 
@@ -98,6 +99,21 @@ public final class CodeBlocks {
             .build());
   }
 
+  public static CodeBlock anonymousLazy(TypeName providedType, CodeBlock body) {
+    return CodeBlock.of(
+        "$L",
+        anonymousClassBuilder("")
+            .superclass(lazyOf(providedType))
+            .addMethod(
+                methodBuilder("get")
+                    .addAnnotation(Override.class)
+                    .addModifiers(PUBLIC)
+                    .returns(providedType)
+                    .addCode(body)
+                    .build())
+            .build());
+  }
+
   /** Returns {@code expression} cast to a type. */
   public static CodeBlock cast(CodeBlock expression, ClassName castTo) {
     return CodeBlock.of("($T) $L", castTo, expression);
diff --git a/java/dagger/internal/codegen/javapoet/TypeNames.java b/java/dagger/internal/codegen/javapoet/TypeNames.java
index 251ddbe3f..7d77bf2d1 100644
--- a/java/dagger/internal/codegen/javapoet/TypeNames.java
+++ b/java/dagger/internal/codegen/javapoet/TypeNames.java
@@ -53,6 +53,7 @@ public final class TypeNames {
   public static final ClassName SUBCOMPONENT_FACTORY = SUBCOMPONENT.nestedClass("Factory");
 
   // Dagger Internal classnames
+
   public static final ClassName DELEGATE_FACTORY =
       ClassName.get("dagger.internal", "DelegateFactory");
   public static final ClassName DOUBLE_CHECK = ClassName.get("dagger.internal", "DoubleCheck");
@@ -62,6 +63,7 @@ public final class TypeNames {
       ClassName.get("dagger.internal", "InjectedFieldSignature");
   public static final ClassName INSTANCE_FACTORY =
       ClassName.get("dagger.internal", "InstanceFactory");
+  public static final ClassName MAP_BUILDER = ClassName.get("dagger.internal", "MapBuilder");
   public static final ClassName MAP_FACTORY = ClassName.get("dagger.internal", "MapFactory");
   public static final ClassName MAP_PROVIDER_FACTORY =
       ClassName.get("dagger.internal", "MapProviderFactory");
@@ -142,6 +144,8 @@ public final class TypeNames {
   public static final ClassName ERROR = ClassName.get("java.lang", "Error");
   public static final ClassName EXCEPTION = ClassName.get("java.lang", "Exception");
   public static final ClassName RUNTIME_EXCEPTION = ClassName.get("java.lang", "RuntimeException");
+  public static final ClassName STRING = ClassName.get("java.lang", "String");
+
   public static final ClassName MAP = ClassName.get("java.util", "Map");
   public static final ClassName KOTLIN_METADATA = ClassName.get("kotlin", "Metadata");
   public static final ClassName IMMUTABLE_MAP =
diff --git a/java/dagger/internal/codegen/writing/ComponentImplementation.java b/java/dagger/internal/codegen/writing/ComponentImplementation.java
index 585181825..abc39f3d1 100644
--- a/java/dagger/internal/codegen/writing/ComponentImplementation.java
+++ b/java/dagger/internal/codegen/writing/ComponentImplementation.java
@@ -472,7 +472,6 @@ public final class ComponentImplementation {
     private ShardImplementation(ClassName name) {
       this.name = name;
       this.switchingProviders = new SwitchingProviders(this, processingEnv);
-
       if (graph.componentDescriptor().isProduction()) {
         claimMethodName(CANCELLATION_LISTENER_METHOD_NAME);
       }
diff --git a/java/dagger/internal/codegen/writing/MapFactoryCreationExpression.java b/java/dagger/internal/codegen/writing/MapFactoryCreationExpression.java
index 1f857dbf1..3512353e7 100644
--- a/java/dagger/internal/codegen/writing/MapFactoryCreationExpression.java
+++ b/java/dagger/internal/codegen/writing/MapFactoryCreationExpression.java
@@ -22,8 +22,9 @@ import static dagger.internal.codegen.binding.SourceFiles.mapFactoryClassName;
 import static dagger.internal.codegen.extension.DaggerCollectors.toOptional;
 
 import androidx.room.compiler.processing.XProcessingEnv;
-import androidx.room.compiler.processing.XType;
+import com.squareup.javapoet.ClassName;
 import com.squareup.javapoet.CodeBlock;
+import com.squareup.javapoet.TypeName;
 import dagger.assisted.Assisted;
 import dagger.assisted.AssistedFactory;
 import dagger.assisted.AssistedInject;
@@ -58,18 +59,24 @@ final class MapFactoryCreationExpression extends MultibindingFactoryCreationExpr
 
   @Override
   public CodeBlock creationExpression() {
-    CodeBlock.Builder builder = CodeBlock.builder().add("$T.", mapFactoryClassName(binding));
+    ClassName mapFactoryClassName = mapFactoryClassName(binding);
+    CodeBlock.Builder builder = CodeBlock.builder().add("$T.", mapFactoryClassName);
+    TypeName valueTypeName = TypeName.OBJECT;
     if (!useRawType()) {
       MapType mapType = MapType.from(binding.key());
       // TODO(ronshapiro): either inline this into mapFactoryClassName, or add a
       // mapType.unwrappedValueType() method that doesn't require a framework type
-      XType valueType =
+      valueTypeName =
           Stream.of(TypeNames.PROVIDER, TypeNames.PRODUCER, TypeNames.PRODUCED)
               .filter(mapType::valuesAreTypeOf)
               .map(mapType::unwrappedValueType)
               .collect(toOptional())
-              .orElseGet(mapType::valueType);
-      builder.add("<$T, $T>", mapType.keyType().getTypeName(), valueType.getTypeName());
+              .orElseGet(mapType::valueType)
+              .getTypeName();
+      builder.add(
+          "<$T, $T>",
+          mapType.keyType().getTypeName(),
+          valueTypeName);
     }
 
     builder.add("builder($L)", binding.dependencies().size());
@@ -78,12 +85,12 @@ final class MapFactoryCreationExpression extends MultibindingFactoryCreationExpr
       ContributionBinding contributionBinding = graph.contributionBinding(dependency.key());
       builder.add(
           ".put($L, $L)",
-          getMapKeyExpression(contributionBinding, componentImplementation.name(), processingEnv),
+          getMapKeyExpression(
+          contributionBinding, componentImplementation.name(), processingEnv),
           multibindingDependencyExpression(dependency));
     }
-    builder.add(".build()");
 
-    return builder.build();
+    return builder.add(".build()").build();
   }
 
   @AssistedFactory
diff --git a/java/dagger/internal/codegen/writing/MapRequestRepresentation.java b/java/dagger/internal/codegen/writing/MapRequestRepresentation.java
index a04ba1ffe..d42b6bfc7 100644
--- a/java/dagger/internal/codegen/writing/MapRequestRepresentation.java
+++ b/java/dagger/internal/codegen/writing/MapRequestRepresentation.java
@@ -73,6 +73,12 @@ final class MapRequestRepresentation extends RequestRepresentation {
 
   @Override
   Expression getDependencyExpression(ClassName requestingClass) {
+    MapType mapType = MapType.from(binding.key());
+    Expression dependencyExpression = getUnderlyingMapExpression(requestingClass);
+    return dependencyExpression;
+  }
+
+  private Expression getUnderlyingMapExpression(ClassName requestingClass) {
     // TODO(ronshapiro): We should also make an ImmutableMap version of MapFactory
     boolean isImmutableMapAvailable = isImmutableMapAvailable();
     // TODO(ronshapiro, gak): Use Maps.immutableEnumMap() if it's available?
@@ -84,9 +90,7 @@ final class MapRequestRepresentation extends RequestRepresentation {
               .add(maybeTypeParameters(requestingClass))
               .add(
                   "of($L)",
-                  dependencies
-                      .keySet()
-                      .stream()
+                  dependencies.keySet().stream()
                       .map(dependency -> keyAndValueExpression(dependency, requestingClass))
                       .collect(toParametersCodeBlock()))
               .build());
@@ -101,10 +105,10 @@ final class MapRequestRepresentation extends RequestRepresentation {
                 "singletonMap($L)",
                 keyAndValueExpression(getOnlyElement(dependencies.keySet()), requestingClass)));
       default:
-        CodeBlock.Builder instantiation = CodeBlock.builder();
-        instantiation
-            .add("$T.", isImmutableMapAvailable ? ImmutableMap.class : MapBuilder.class)
-            .add(maybeTypeParameters(requestingClass));
+        CodeBlock.Builder instantiation =
+            CodeBlock.builder()
+                .add("$T.", isImmutableMapAvailable ? ImmutableMap.class : MapBuilder.class)
+                .add(maybeTypeParameters(requestingClass));
         if (isImmutableMapBuilderWithExpectedSizeAvailable()) {
           instantiation.add("builderWithExpectedSize($L)", dependencies.size());
         } else if (isImmutableMapAvailable) {
@@ -132,7 +136,8 @@ final class MapRequestRepresentation extends RequestRepresentation {
   private CodeBlock keyAndValueExpression(DependencyRequest dependency, ClassName requestingClass) {
     return CodeBlock.of(
         "$L, $L",
-        getMapKeyExpression(dependencies.get(dependency), requestingClass, processingEnv),
+        getMapKeyExpression(
+        dependencies.get(dependency), requestingClass, processingEnv),
         componentRequestRepresentations
             .getDependencyExpression(bindingRequest(dependency), requestingClass)
             .codeBlock());
@@ -154,7 +159,9 @@ final class MapRequestRepresentation extends RequestRepresentation {
     MapType mapType = MapType.from(binding.key());
     return isTypeAccessibleFrom(bindingKeyType, requestingClass.packageName())
         ? CodeBlock.of(
-            "<$T, $T>", mapType.keyType().getTypeName(), mapType.valueType().getTypeName())
+            "<$T, $T>",
+            mapType.keyType().getTypeName(),
+            mapType.valueType().getTypeName())
         : CodeBlock.of("");
   }