Support @LazyClassKey usage with dagger producers.

RELNOTES=n/a
PiperOrigin-RevId: 627152997

java/dagger/internal/LazyClassKeyMap.java

@@ -111,24 +111,39 @@ public void putAll(Map<? extends Class<?>, ? extends V> map) {
     throw new UnsupportedOperationException("Dagger map bindings are immutable");
   }
 
-  // TODO(wanyingd) Support @LazyClassKey on producer
-  /** A factory for {@code LazyClassKeyMap}. */
-  @SuppressWarnings("unchecked")
-  public static class Factory<V> implements Provider<Map<Class<?>, V>> {
-    Provider<?> delegate;
-
-    // MapProviderFactory or MapFactory
-    public static <V> Factory<V> of(Provider<?> delegate) {
-      return new Factory<V>(delegate);
+  /** Wrapper around {@link MapFactory}. */
+  public static class MapFactory<V> implements Factory<Map<Class<?>, V>> {
+    Factory<Map<String, V>> delegate;
+
+    public static <V> MapFactory<V> of(Factory<Map<String, V>> delegate) {
+      return new MapFactory<V>(delegate);
     }
 
-    private Factory(Provider<?> delegate) {
+    private MapFactory(Factory<Map<String, V>> delegate) {
       this.delegate = delegate;
     }
 
     @Override
     public Map<Class<?>, V> get() {
-      return LazyClassKeyMap.of((Map<String, V>) delegate.get());
+      return LazyClassKeyMap.of(delegate.get());
+    }
+  }
+
+  /** Wrapper around for {@link MapProviderFactory}. */
+  public static class MapProviderFactory<V> implements Factory<Map<Class<?>, Provider<V>>> {
+    Factory<Map<String, Provider<V>>> delegate;
+
+    public static <V> MapProviderFactory<V> of(Factory<Map<String, Provider<V>>> delegate) {
+      return new MapProviderFactory<V>(delegate);
+    }
+
+    private MapProviderFactory(Factory<Map<String, Provider<V>>> delegate) {
+      this.delegate = delegate;
+    }
+
+    @Override
+    public Map<Class<?>, Provider<V>> get() {
+      return LazyClassKeyMap.of(delegate.get());
     }
   }
 }

java/dagger/internal/codegen/javapoet/TypeNames.java

@@ -61,7 +61,15 @@ public final class TypeNames {
   public static final ClassName LAZY_CLASS_KEY_MAP =
       ClassName.get("dagger.internal", "LazyClassKeyMap");
   public static final ClassName LAZY_CLASS_KEY_MAP_FACTORY =
-      ClassName.get("dagger.internal", "LazyClassKeyMap", "Factory");
+      ClassName.get("dagger.internal", "LazyClassKeyMap", "MapFactory");
+  public static final ClassName LAZY_CLASS_KEY_MAP_PROVIDER_FACTORY =
+      ClassName.get("dagger.internal", "LazyClassKeyMap", "MapProviderFactory");
+  public static final ClassName LAZY_MAP_OF_PRODUCED_PRODUCER =
+      ClassName.get("dagger.producers.internal", "LazyMapOfProducedProducer");
+  public static final ClassName LAZY_MAP_OF_PRODUCER_PRODUCER =
+      ClassName.get("dagger.producers.internal", "LazyMapOfProducerProducer");
+  public static final ClassName LAZY_MAP_PRODUCER =
+      ClassName.get("dagger.producers.internal", "LazyMapProducer");
 
   public static final ClassName DELEGATE_FACTORY =
       ClassName.get("dagger.internal", "DelegateFactory");

java/dagger/internal/codegen/writing/MapFactoryCreationExpression.java

@@ -101,12 +101,30 @@ public CodeBlock creationExpression() {
     return useLazyClassKey
         ? CodeBlock.of(
             "$T.<$T>of($L)",
-            TypeNames.LAZY_CLASS_KEY_MAP_FACTORY,
+            lazyMapFactoryClassName(binding),
             valueTypeName,
             builder.add(".build()").build())
         : builder.add(".build()").build();
   }
 
+  private static ClassName lazyMapFactoryClassName(ContributionBinding binding) {
+    MapType mapType = MapType.from(binding.key());
+    switch (binding.bindingType()) {
+      case PROVISION:
+        return mapType.valuesAreTypeOf(TypeNames.PROVIDER)
+            ? TypeNames.LAZY_CLASS_KEY_MAP_PROVIDER_FACTORY
+            : TypeNames.LAZY_CLASS_KEY_MAP_FACTORY;
+      case PRODUCTION:
+        return mapType.valuesAreFrameworkType()
+            ? mapType.valuesAreTypeOf(TypeNames.PRODUCER)
+                ? TypeNames.LAZY_MAP_OF_PRODUCER_PRODUCER
+                : TypeNames.LAZY_MAP_OF_PRODUCED_PRODUCER
+            : TypeNames.LAZY_MAP_PRODUCER;
+      default:
+        throw new IllegalArgumentException(binding.bindingType().toString());
+    }
+  }
+
   @AssistedFactory
   static interface Factory {
     MapFactoryCreationExpression create(ContributionBinding binding);

java/dagger/producers/internal/LazyMapOfProducedProducer.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2024 The Dagger Authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package dagger.producers.internal;
+
+import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
+
+import com.google.common.base.Function;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import dagger.internal.LazyClassKeyMap;
+import dagger.producers.Produced;
+import java.util.Map;
+
+/**
+ * Wrapper around {@link MapOfProducedProducer} to be compatible with @LazyClassKey annotated map.
+ */
+public final class LazyMapOfProducedProducer<V>
+    extends AbstractProducer<Map<Class<?>, Produced<V>>> {
+  AbstractProducer<Map<String, Produced<V>>> delegate;
+
+  public static <V> LazyMapOfProducedProducer<V> of(
+      AbstractProducer<Map<String, Produced<V>>> delegate) {
+    return new LazyMapOfProducedProducer<V>(delegate);
+  }
+
+  private LazyMapOfProducedProducer(AbstractProducer<Map<String, Produced<V>>> delegate) {
+    this.delegate = delegate;
+  }
+
+  @Override
+  public ListenableFuture<Map<Class<?>, Produced<V>>> compute() {
+    return Futures.transform(
+        delegate.compute(),
+        new Function<Map<String, Produced<V>>, Map<Class<?>, Produced<V>>>() {
+          @Override
+          public Map<Class<?>, Produced<V>> apply(Map<String, Produced<V>> classMap) {
+            return LazyClassKeyMap.of(classMap);
+          }
+        },
+        directExecutor());
+  }
+}

java/dagger/producers/internal/LazyMapOfProducerProducer.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2024 The Dagger Authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package dagger.producers.internal;
+
+import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
+
+import com.google.common.base.Function;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import dagger.internal.LazyClassKeyMap;
+import dagger.producers.Producer;
+import java.util.Map;
+
+/**
+ * Wrapper around {@link MapOfProducerProducer} to be compatible with @LazyClassKey annotated map.
+ */
+public final class LazyMapOfProducerProducer<V>
+    extends AbstractProducer<Map<Class<?>, Producer<V>>> {
+  AbstractProducer<Map<String, Producer<V>>> delegate;
+
+  public static <V> LazyMapOfProducerProducer<V> of(
+      AbstractProducer<Map<String, Producer<V>>> delegate) {
+    return new LazyMapOfProducerProducer<V>(delegate);
+  }
+
+  private LazyMapOfProducerProducer(AbstractProducer<Map<String, Producer<V>>> delegate) {
+    this.delegate = delegate;
+  }
+
+  @Override
+  public ListenableFuture<Map<Class<?>, Producer<V>>> compute() {
+    return Futures.transform(
+        delegate.compute(),
+        new Function<Map<String, Producer<V>>, Map<Class<?>, Producer<V>>>() {
+          @Override
+          public Map<Class<?>, Producer<V>> apply(Map<String, Producer<V>> classMap) {
+            return LazyClassKeyMap.of((Map<String, Producer<V>>) classMap);
+          }
+        },
+        directExecutor());
+  }
+}

java/dagger/producers/internal/LazyMapProducer.java

@@ -0,0 +1,51 @@
+/*
+ * Copyright (C) 2024 The Dagger Authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package dagger.producers.internal;
+
+import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
+
+import com.google.common.base.Function;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import dagger.internal.LazyClassKeyMap;
+import java.util.Map;
+
+/** Wrapper around {@link MapProducer} to be compatible with @LazyClassKey annotated map. */
+public final class LazyMapProducer<V> extends AbstractProducer<Map<Class<?>, V>> {
+  AbstractProducer<Map<String, V>> delegate;
+
+  public static <V> LazyMapProducer<V> of(AbstractProducer<Map<String, V>> delegate) {
+    return new LazyMapProducer<V>(delegate);
+  }
+
+  private LazyMapProducer(AbstractProducer<Map<String, V>> delegate) {
+    this.delegate = delegate;
+  }
+
+  @Override
+  public ListenableFuture<Map<Class<?>, V>> compute() {
+    return Futures.transform(
+        delegate.compute(),
+        new Function<Map<String, V>, Map<Class<?>, V>>() {
+          @Override
+          public Map<Class<?>, V> apply(Map<String, V> classMap) {
+            return LazyClassKeyMap.of((Map<String, V>) classMap);
+          }
+        },
+        directExecutor());
+  }
+}

javatests/dagger/functional/producers/multibindings/MultibindingComponent.java

@@ -29,7 +29,6 @@
 import java.util.Map;
 import java.util.Set;
 
-
 @ProductionComponent(
   modules = {ExecutorModule.class, MultibindingProducerModule.class, MultibindingModule.class}
 )
@@ -48,6 +47,12 @@ interface MultibindingComponent {
 
   ListenableFuture<Map<Integer, Produced<String>>> mapOfProduced();
 
+  ListenableFuture<Map<Class<?>, Produced<String>>> lazyClassKeyMapOfProduced();
+
+  ListenableFuture<Map<Class<?>, String>> lazyClassKeyMapOfString();
+
+  ListenableFuture<Map<Class<?>, Producer<String>>> lazyClassKeyMapOfProducer();
+
   @PossiblyThrowingMap
   ListenableFuture<Map<Integer, String>> possiblyThrowingMap();
 

javatests/dagger/functional/producers/multibindings/MultibindingProducerModule.java

@@ -27,6 +27,7 @@
 import dagger.multibindings.IntKey;
 import dagger.multibindings.IntoMap;
 import dagger.multibindings.IntoSet;
+import dagger.multibindings.LazyClassKey;
 import dagger.multibindings.Multibinds;
 import dagger.producers.Produced;
 import dagger.producers.ProducerModule;
@@ -36,6 +37,8 @@
 
 @ProducerModule
 abstract class MultibindingProducerModule {
+  static class Foo {}
+
   @Produces
   @IntoSet
   static ListenableFuture<String> futureStr() {
@@ -122,6 +125,13 @@ static String throwingValueFor15() {
     throw new RuntimeException("monkey");
   }
 
+  @Produces
+  @IntoMap
+  @LazyClassKey(Foo.class)
+  static String produceName() {
+    return "foo";
+  }
+
   @Multibinds
   abstract Set<Object> objs();
 

javatests/dagger/internal/codegen/goldens/LazyClassKeyMapBindingComponentProcessorTest_mapBindingsWithInaccessibleKeys_DEFAULT_MODE_test.DaggerTestComponent

@@ -66,7 +66,7 @@ final class DaggerTestComponent {
 
     @SuppressWarnings("unchecked")
     private void initialize() {
-      this.mapOfClassOfAndIntegerProvider = LazyClassKeyMap.Factory.<Integer>of(MapFactory.<String, Integer>builder(1).put(LazyClassKeyProvider.mapkeys_MapKeys_Inaccessible, MapModule_ClassKeyFactory.create()).build());
+      this.mapOfClassOfAndIntegerProvider = LazyClassKeyMap.MapFactory.<Integer>of(MapFactory.<String, Integer>builder(1).put(LazyClassKeyProvider.mapkeys_MapKeys_Inaccessible, MapModule_ClassKeyFactory.create()).build());
       this.mapOfComplexKeyAndIntegerProvider = MapFactory.<MapKeys.ComplexKey, Integer>builder(3).put(MapModule_ComplexKeyWithInaccessibleValueMapKey.create(), MapModule_ComplexKeyWithInaccessibleValueFactory.create()).put(MapModule_ComplexKeyWithInaccessibleArrayValueMapKey.create(), MapModule_ComplexKeyWithInaccessibleArrayValueFactory.create()).put(MapModule_ComplexKeyWithInaccessibleAnnotationValueMapKey.create(), MapModule_ComplexKeyWithInaccessibleAnnotationValueFactory.create()).build();
     }