From 29d9a8ef34e0d576f365f156cc32d5c7f4ae8cf1 Mon Sep 17 00:00:00 2001
From: Wanying Ding <wanyingd@google.com>
Date: Fri, 22 Mar 2024 10:30:43 -0700
Subject: [PATCH] Make LazyClassKeyMap accept both MapFactory and
 MapProviderFactory

Will support producer in a follow up cl.

Fixes #4254

RELNOTES=n/a
PiperOrigin-RevId: 618220537
---
 java/dagger/internal/LazyClassKeyMap.java     | 13 +--
 ...ssKeyMapBindingComponentProcessorTest.java | 84 +++++++++++++++++++
 2 files changed, 92 insertions(+), 5 deletions(-)

diff --git a/java/dagger/internal/LazyClassKeyMap.java b/java/dagger/internal/LazyClassKeyMap.java
index dabf86f9132..defa3d9c7f8 100644
--- a/java/dagger/internal/LazyClassKeyMap.java
+++ b/java/dagger/internal/LazyClassKeyMap.java
@@ -111,21 +111,24 @@ 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>> {
-    MapFactory<String, V> delegate;
+    Provider<?> delegate;
 
-    public static <V> Factory<V> of(MapFactory<String, V> delegate) {
-      return new Factory<>(delegate);
+    // MapProviderFactory or MapFactory
+    public static <V> Factory<V> of(Provider<?> delegate) {
+      return new Factory<V>(delegate);
     }
 
-    private Factory(MapFactory<String, V> delegate) {
+    private Factory(Provider<?> delegate) {
       this.delegate = delegate;
     }
 
     @Override
     public Map<Class<?>, V> get() {
-      return LazyClassKeyMap.of(delegate.get());
+      return LazyClassKeyMap.of((Map<String, V>) delegate.get());
     }
   }
 }
diff --git a/javatests/dagger/internal/codegen/LazyClassKeyMapBindingComponentProcessorTest.java b/javatests/dagger/internal/codegen/LazyClassKeyMapBindingComponentProcessorTest.java
index 9937d05ed09..ba3fab665b6 100644
--- a/javatests/dagger/internal/codegen/LazyClassKeyMapBindingComponentProcessorTest.java
+++ b/javatests/dagger/internal/codegen/LazyClassKeyMapBindingComponentProcessorTest.java
@@ -188,4 +188,88 @@ public void lazyClassKeySimilarQualifiedName_doesNotConflict() throws Exception
         .withProcessingOptions(compilerMode.processorOptions())
         .compile(subject -> subject.hasErrorCount(0));
   }
+
+  @Test
+  public void lazyClassKeyProvider_compilesSuccessfully() throws Exception {
+    Source fooBar =
+        CompilerTests.javaSource("test.Foo_Bar", "package test;", "", "interface Foo_Bar {}");
+    Source fooBar2 =
+        CompilerTests.javaSource(
+            "test.Foo", "package test;", "", "interface Foo { interface Bar {} }");
+    Source mapKeyBindingsModule =
+        CompilerTests.javaSource(
+            "test.MapKeyBindingsModule",
+            "package test;",
+            "",
+            "import dagger.Module;",
+            "import dagger.Provides;",
+            "import dagger.multibindings.LazyClassKey;",
+            "import dagger.multibindings.IntoMap;",
+            "",
+            "@Module",
+            "public interface MapKeyBindingsModule {",
+            " @Provides @IntoMap @LazyClassKey(test.Foo_Bar.class)",
+            " static int classKey() { return 1; }",
+            "}");
+
+    Source componentFile =
+        CompilerTests.javaSource(
+            "test.TestComponent",
+            "package test;",
+            "",
+            "import dagger.Component;",
+            "import javax.inject.Provider;",
+            "import java.util.Map;",
+            "",
+            "@Component(modules = MapKeyBindingsModule.class)",
+            "interface TestComponent {",
+            "  Map<Class<?>, Provider<Integer>> classKey();",
+            "}");
+    CompilerTests.daggerCompiler(fooBar, fooBar2, mapKeyBindingsModule, componentFile)
+        .withProcessingOptions(compilerMode.processorOptions())
+        .compile(subject -> subject.hasErrorCount(0));
+  }
+
+  @Test
+  public void scopedLazyClassKeyProvider_compilesSuccessfully() throws Exception {
+    Source fooBar =
+        CompilerTests.javaSource("test.Foo_Bar", "package test;", "", "interface Foo_Bar {}");
+    Source fooBar2 =
+        CompilerTests.javaSource(
+            "test.Foo", "package test;", "", "interface Foo { interface Bar {} }");
+    Source mapKeyBindingsModule =
+        CompilerTests.javaSource(
+            "test.MapKeyBindingsModule",
+            "package test;",
+            "",
+            "import dagger.Module;",
+            "import dagger.Provides;",
+            "import dagger.multibindings.LazyClassKey;",
+            "import dagger.multibindings.IntoMap;",
+            "",
+            "@Module",
+            "public interface MapKeyBindingsModule {",
+            " @Provides @IntoMap @LazyClassKey(test.Foo_Bar.class)",
+            " static int classKey() { return 1; }",
+            "}");
+
+    Source componentFile =
+        CompilerTests.javaSource(
+            "test.TestComponent",
+            "package test;",
+            "",
+            "import dagger.Component;",
+            "import javax.inject.Singleton;",
+            "import javax.inject.Provider;",
+            "import java.util.Map;",
+            "",
+            "@Component(modules = MapKeyBindingsModule.class)",
+            "@Singleton",
+            "interface TestComponent {",
+            "  Provider<Map<Class<?>, Provider<Integer>>> classKey();",
+            "}");
+    CompilerTests.daggerCompiler(fooBar, fooBar2, mapKeyBindingsModule, componentFile)
+        .withProcessingOptions(compilerMode.processorOptions())
+        .compile(subject -> subject.hasErrorCount(0));
+  }
 }