Support Map<K, Lazy<V>> as a multibinding request type.

This change adds `dagger.Lazy` to the set of recognized framework types and allows `RequestKind.LAZY` for map multibinding values. A new test case is added to verify the generated code for injecting `Map<Long, Lazy<String>>`. Debug print statements were added to `MapType.unwrappedFrameworkValueType`.

RELNOTES=Supported `Map<K, Lazy<V>>` as a multibinding request type.
PiperOrigin-RevId: 843328737

Author: trexem

dagger-compiler/main/java/dagger/internal/codegen/base/FrameworkTypes.java

@@ -52,6 +52,8 @@ public final class FrameworkTypes {
 
   public static final ImmutableSet<XClassName> MAP_VALUE_FRAMEWORK_TYPES =
       ImmutableSet.<XClassName>builder()
+          .add(XTypeNames.PROVIDER_OF_LAZY)
+          .add(XTypeNames.LAZY)
           .addAll(providerTypeNames())
           .add(XTypeNames.PRODUCED)
           .add(XTypeNames.PRODUCER)

dagger-compiler/main/java/dagger/internal/codegen/base/MapType.java

@@ -53,12 +53,16 @@ public static MapType from(XType type) {
     return new MapType(type);
   }
 
-  // TODO(b/28555349): support PROVIDER_OF_LAZY here too
   // TODO(b/376124787): We could consolidate this with a similar list in FrameworkTypes
   // if we had a better way to go from RequestKind to framework class name or vice versa
   /** The valid framework request kinds allowed on a multibinding map value. */
   private static final ImmutableSet<RequestKind> VALID_FRAMEWORK_REQUEST_KINDS =
-      ImmutableSet.of(RequestKind.PROVIDER, RequestKind.PRODUCER, RequestKind.PRODUCED);
+      ImmutableSet.of(
+          RequestKind.PROVIDER,
+          RequestKind.PRODUCER,
+          RequestKind.PRODUCED,
+          RequestKind.LAZY,
+          RequestKind.PROVIDER_OF_LAZY);
 
   private final XType type;
 
@@ -124,6 +128,16 @@ public boolean valuesAreProvider() {
     return providerTypeNames().stream().anyMatch(this::valuesAreTypeOf);
   }
 
+  /** Returns {@code true} if the raw type of {@link #valueType()} is a lazy type. */
+  public boolean valuesAreLazy() {
+    return valueRequestKind().equals(RequestKind.LAZY);
+  }
+
+  /** Returns {@code true} if the raw type of {@link #valueType()} is a provider of lazy type. */
+  public boolean valuesAreProviderOfLazy() {
+    return valueRequestKind().equals(RequestKind.PROVIDER_OF_LAZY);
+  }
+
   /**
    * Returns the map's {@link #valueType()} without any wrapping framework type, if one exists.
    *
@@ -145,6 +159,9 @@ public XType unwrappedFrameworkValueType() {
   }
 
   public XTypeName unwrappedFrameworkValueTypeName() {
+    if (valueRequestKind().equals(RequestKind.PROVIDER_OF_LAZY)) {
+      return unwrap(unwrap(valueTypeName()));
+    }
     return valuesAreFrameworkType() ? unwrap(valueTypeName()) : valueTypeName();
   }
 
@@ -158,11 +175,6 @@ public RequestKind valueRequestKind() {
     RequestKind requestKind = RequestKinds.getRequestKind(valueType());
     if (VALID_FRAMEWORK_REQUEST_KINDS.contains(requestKind)) {
       return requestKind;
-    } else if (requestKind == RequestKind.PROVIDER_OF_LAZY) {
-      // This is kind of a weird case. We don't support Map<K, Lazy<V>>, so we also don't support
-      // Map<K, Provider<Lazy<V>>> directly. However, if the user bound that themselves, we don't
-      // want that to get confused as a normal instance request, so return PROVIDER here.
-      return RequestKind.PROVIDER;
     } else {
       // Not all RequestKinds are supported, so if there's a map value that matches an unsupported
       // RequestKind, just treat it like it is a normal instance request.

dagger-compiler/main/java/dagger/internal/codegen/binding/ComponentDeclarations.java

@@ -21,6 +21,8 @@
 import static com.google.common.collect.Iterables.getOnlyElement;
 import static dagger.internal.codegen.binding.SourceFiles.generatedMonitoringModuleName;
 import static dagger.internal.codegen.extension.DaggerStreams.toImmutableSet;
+import static dagger.internal.codegen.xprocessing.XTypeNames.LAZY;
+import static dagger.internal.codegen.xprocessing.XTypeNames.providerTypeNames;
 
 import androidx.room3.compiler.codegen.XClassName;
 import androidx.room3.compiler.codegen.XTypeName;
@@ -353,15 +355,32 @@ private static boolean isValidSetMultibindingTypeName(TypeName typeName) {
         && !(getOnlyElement(parameterizedTypeName.typeArguments) instanceof WildcardTypeName);
   }
 
+  /**
+   * Returns the given {@code typeName} with any framework type wrappers removed.
+   *
+   * <p>For example, if {@code frameworkTypeNames} contains {@code Produced}, then {@code
+   * Produced<Foo>} is unwrapped to {@code Foo}. And if {@code Provider<Lazy<Foo>>} is encountered,
+   * we unwrap the {@code Provider} and {@code Lazy} layers and return {@code Foo}.
+   */
   private static TypeName unwrapFrameworkTypeName(
       TypeName typeName, ImmutableSet<XClassName> frameworkTypeNames) {
-    if (typeName instanceof ParameterizedTypeName) {
-      ParameterizedTypeName parameterizedTypeName = (ParameterizedTypeName) typeName;
-      if (frameworkTypeNames.contains(toXPoet(parameterizedTypeName.rawType))) {
-        typeName = getOnlyElement(parameterizedTypeName.typeArguments);
-      }
+    if (!(typeName instanceof ParameterizedTypeName)) {
+      return typeName;
     }
-    return typeName;
+    ParameterizedTypeName parameterizedTypeName = (ParameterizedTypeName) typeName;
+    XClassName rawType = toXPoet(parameterizedTypeName.rawType);
+    if (!frameworkTypeNames.contains(rawType)) {
+      return typeName;
+    }
+    // If we have a framework type, unwrap it.
+    TypeName unwrappedTypeName = getOnlyElement(parameterizedTypeName.typeArguments);
+
+    // Account for Provider<Lazy<T>>:
+    // If we just unwrapped Provider<T> and T is Lazy<U>, we need to unwrap Lazy<U> as well.
+    if (providerTypeNames().contains(rawType)) {
+      return unwrapFrameworkTypeName(unwrappedTypeName, ImmutableSet.of(LAZY));
+    }
+    return unwrappedTypeName;
   }
 
   /**

dagger-compiler/main/java/dagger/internal/codegen/binding/SourceFiles.java

@@ -276,8 +276,16 @@ public static XClassName mapFactoryClassName(MultiboundMapBinding binding) {
     MapType mapType = MapType.from(binding.key());
     switch (binding.bindingType()) {
       case PROVISION:
-        return mapType.valuesAreProvider()
-            ? XTypeNames.MAP_PROVIDER_FACTORY : XTypeNames.MAP_FACTORY;
+        if (mapType.valuesAreProviderOfLazy()) {
+          return XTypeNames.MAP_PROVIDER_LAZY_FACTORY;
+        }
+        if (mapType.valuesAreLazy()) {
+          return XTypeNames.MAP_LAZY_FACTORY;
+        }
+        if (mapType.valuesAreProvider()) {
+          return XTypeNames.MAP_PROVIDER_FACTORY;
+        }
+        return XTypeNames.MAP_FACTORY;
       case PRODUCTION:
         return mapType.valuesAreFrameworkType()
             ? mapType.valuesAreTypeOf(XTypeNames.PRODUCER)

dagger-compiler/main/java/dagger/internal/codegen/writing/MapFactoryCreationExpression.java

@@ -79,6 +79,9 @@ final class MapFactoryCreationExpression extends MultibindingFactoryCreationExpr
 
   @Override
   public XCodeBlock creationExpression() {
+    // TODO(b/467104887): Revisit how DependencyRequests are generated for Map<K, Lazy<V>> and
+    // Provider<Map<K, Lazy<V>>> to ensure consistent handling of Provider wrapping. See bug for
+    // more details on the observed differences between direct and Provider-based map injections.
     return XCodeBlock.of(
         "%N(%L)", methodName(), parameterNames(shardImplementation.constructorParameters()));
   }

dagger-compiler/main/java/dagger/internal/codegen/writing/MapRequestRepresentation.java

@@ -90,6 +90,9 @@ final class MapRequestRepresentation extends RequestRepresentation {
 
   @Override
   XExpression getDependencyExpression(XClassName requestingClass) {
+    // TODO(b/467104887): Revisit how DependencyRequests are generated for Map<K, Lazy<V>> and
+    // Provider<Map<K, Lazy<V>>> to ensure consistent handling of Provider wrapping. See bug for more
+    // details on the observed differences between direct and Provider-based map injections.
     XExpression dependencyExpression = getUnderlyingMapExpression(requestingClass);
     // LazyClassKey is backed with a string map, therefore needs to be wrapped.
     if (useLazyClassKey) {

dagger-compiler/main/java/dagger/internal/codegen/xprocessing/XTypeNames.kt

@@ -101,6 +101,9 @@ object XTypeNames {
   @JvmField val MAP_BUILDER = XClassName.get("dagger.internal", "MapBuilder")
   @JvmField val MAP_FACTORY = XClassName.get("dagger.internal", "MapFactory")
   @JvmField val MAP_PROVIDER_FACTORY = XClassName.get("dagger.internal", "MapProviderFactory")
+  @JvmField val MAP_LAZY_FACTORY = XClassName.get("dagger.internal", "MapLazyFactory")
+  @JvmField
+  val MAP_PROVIDER_LAZY_FACTORY = XClassName.get("dagger.internal", "MapProviderLazyFactory")
   @JvmField val MEMBERS_INJECTOR = XClassName.get("dagger", "MembersInjector")
   @JvmField val MEMBERS_INJECTORS = XClassName.get("dagger.internal", "MembersInjectors")
   @JvmField val DAGGER_PROVIDER = XClassName.get("dagger.internal", "Provider")

javatests/dagger/functional/jakarta/JakartaProviderTest.java

@@ -172,14 +172,6 @@ static Bar provideBar(
     @LazyClassKey(Bar.class)
     abstract Bar bindBarIntoClassMap(Bar bar);
 
-    // TODO(b/65118638): Use @Binds @IntoMap Lazy<T> once that works properly.
-    @Provides
-    @IntoMap
-    @StringKey("bar")
-    static Lazy<Bar> provideLazyIntoMap(Lazy<Bar> bar) {
-      return bar;
-    }
-
     // Manually provide two Provider maps to make sure they don't conflict.
     @Provides
     static Map<Long, Provider<Long>> manuallyProvidedJakartaMap() {

javatests/dagger/functional/kotlinsrc/multibindings/LazyMaps.kt

@@ -16,6 +16,7 @@
 
 package dagger.functional.kotlinsrc.multibindings
 
+import dagger.Binds
 import dagger.Component
 import dagger.Lazy
 import dagger.Module
@@ -34,31 +35,34 @@ import javax.inject.Singleton
  */
 class LazyMaps {
   @Module
-  internal object TestModule {
-    @Provides @Singleton fun provideAtomicInteger(): AtomicInteger = AtomicInteger()
+  abstract class TestModule {
 
-    @Provides
-    fun provideString(atomicInteger: AtomicInteger): String =
-      "value-${atomicInteger.incrementAndGet()}"
+    @Module
+    companion object {
+      @Provides @Singleton fun provideAtomicInteger(): AtomicInteger = AtomicInteger()
+
+      @Provides
+      fun provideString(atomicInteger: AtomicInteger): String =
+        "value-${atomicInteger.incrementAndGet()}"
+
+      @Provides @IntoMap @StringKey("key") fun mapContribution(string: String): String = string
+    }
 
-    /* TODO(b/65118638) Replace once @Binds @IntoMap Lazy<T> methods work properly.
     @Binds
     @IntoMap
     @StringKey("binds-key")
-    abstract Lazy<String> mapContributionAsBinds(Lazy<String> lazy);
-    */
-    @Provides
-    @IntoMap
-    @StringKey("key")
-    fun mapContribution(lazy: Lazy<String>): Lazy<String> = lazy
+    abstract fun mapContributionAsBinds(string: String): String
   }
 
   @Singleton
   @Component(modules = [TestModule::class])
   interface TestComponent {
     fun mapOfLazy(): Map<String, Lazy<String>>
+
     fun mapOfProviderOfLazy(): Map<String, Provider<Lazy<String>>>
+
     fun providerForMapOfLazy(): Provider<Map<String, Lazy<String>>>
+
     fun providerForMapOfProviderOfLazy(): Provider<Map<String, Provider<Lazy<String>>>>
   }
 }

javatests/dagger/functional/kotlinsrc/multibindings/LazyMapsTest.kt

@@ -17,9 +17,6 @@
 package dagger.functional.kotlinsrc.multibindings
 
 import com.google.common.truth.Truth.assertThat
-import dagger.Lazy
-import dagger.functional.kotlinsrc.multibindings.LazyMaps.TestComponent
-import javax.inject.Provider
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.JUnit4
@@ -35,6 +32,11 @@ class LazyMapsTest {
     assertThat(firstGet).isEqualTo("value-1")
     assertThat(firstGet).isSameInstanceAs(laziesMap["key"]!!.get())
     assertThat(component.mapOfLazy()["key"]!!.get()).isEqualTo("value-2")
+
+    val bindsGet = laziesMap["binds-key"]!!.get()
+    assertThat(bindsGet).isEqualTo("value-3")
+    assertThat(bindsGet).isSameInstanceAs(laziesMap["binds-key"]!!.get())
+    assertThat(component.mapOfLazy()["binds-key"]!!.get()).isEqualTo("value-4")
   }
 
   @Test