Share a single SwitchingProvider instance across SingleCheck and DoubleCheck providers in fastInit mode.

Previously in Dagger's fastInit mode, each scoped provider (SingleCheck, DoubleCheck) or assisted factory instantiated a dedicated SwitchingProvider instance:
```java
    this.fooFactoryProvider = SingleCheck.provider(new SwitchingProvider<>(testComponentImpl, 0));
```
This change introduces `dagger.internal.SwitchingProvider` and shares a single `SwitchingProvider` instance across all scoped and assisted factory bindings within a component:
1. Introduce the `dagger.internal.SwitchingProvider<T>` interface with a parameterized `T get(int id)` method.
2. Add `provider(SwitchingProvider<?>, int)` overloads to `SingleCheck` and `DoubleCheck`.
3. In component implementations, generate a single shared `private SwitchingProvider<Object> switchingProvider` instance initialized in `initialize()` and passed to `SingleCheck.provider` and `DoubleCheck.provider`.
4. Unscoped bindings continue to instantiate `SwitchingProvider` directly as a `Provider<T>`.

This reduces redundant SwitchingProvider object allocations during component initialization in fastInit mode.

RELNOTES=n/a
PiperOrigin-RevId: 944582946

Author: trexem

dagger-compiler/main/java/dagger/internal/codegen/writing/BUILD

@@ -39,7 +39,6 @@ java_library(
         "//third_party/java/error_prone:annotations",
         "//third_party/java/guava/base",
         "//third_party/java/guava/collect",
-        "//third_party/java/javapoet",
         "//third_party/java/jsr330_inject:jsr330_inject_without_dagger",
     ],
 )

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

@@ -496,7 +496,7 @@ public final class ShardImplementation implements GeneratedImplementation {
 
     private ShardImplementation(XClassName name) {
       this.name = name;
-      this.switchingProviders = new SwitchingProviders(this, compilerOptions, processingEnv);
+      this.switchingProviders = new SwitchingProviders(this, compilerOptions);
       if (graph.componentDescriptor().isProduction()) {
         claimMethodName(CANCELLATION_LISTENER_METHOD_NAME);
       }

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

@@ -16,26 +16,21 @@
 
 package dagger.internal.codegen.writing;
 
-import androidx.room3.compiler.codegen.XCodeBlock;
 import dagger.assisted.Assisted;
 import dagger.assisted.AssistedFactory;
 import dagger.assisted.AssistedInject;
-import dagger.internal.codegen.binding.Binding;
 import dagger.internal.codegen.binding.BindingGraph;
 import dagger.internal.codegen.binding.ContributionBinding;
 import dagger.internal.codegen.compileroption.CompilerOptions;
-import dagger.internal.codegen.model.BindingKind;
 import dagger.internal.codegen.writing.ComponentImplementation.ShardImplementation;
 import dagger.internal.codegen.writing.FrameworkFieldInitializer.FrameworkInstanceCreationExpression;
-import dagger.internal.codegen.xprocessing.XTypeNames;
 
 /**
  * An object that initializes a framework-type component field for a binding using instances created
  * by switching providers.
  */
 final class SwitchingProviderInstanceSupplier implements FrameworkInstanceSupplier {
   private final FrameworkInstanceSupplier frameworkInstanceSupplier;
-  private final CompilerOptions compilerOptions;
 
   @AssistedInject
   SwitchingProviderInstanceSupplier(
@@ -56,33 +51,14 @@ final class SwitchingProviderInstanceSupplier implements FrameworkInstanceSuppli
             compilerOptions,
             componentImplementation,
             binding,
-            scope(binding, frameworkInstanceCreationExpression));
-    this.compilerOptions = compilerOptions;
+            frameworkInstanceCreationExpression);
   }
 
   @Override
   public MemberSelect memberSelect() {
     return frameworkInstanceSupplier.memberSelect();
   }
 
-  private FrameworkInstanceCreationExpression scope(
-      Binding binding, FrameworkInstanceCreationExpression unscoped) {
-    // Caching assisted factory provider, so that there won't be new factory created for each
-    // provider.get() call.
-    if (!binding.scope().isPresent() && !binding.kind().equals(BindingKind.ASSISTED_FACTORY)) {
-      return unscoped;
-    }
-    return () ->
-        XCodeBlock.of(
-            "%T.provider(%L)",
-            binding.scope().isPresent()
-                ? (binding.scope().get().isReusable()
-                    ? XTypeNames.SINGLE_CHECK
-                    : XTypeNames.DOUBLE_CHECK)
-                : XTypeNames.SINGLE_CHECK,
-            unscoped.creationExpression());
-  }
-
   @AssistedFactory
   static interface Factory {
     SwitchingProviderInstanceSupplier create(ContributionBinding binding);

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

@@ -27,7 +27,6 @@
 import static dagger.internal.codegen.xprocessing.XCodeBlocks.toParametersCodeBlock;
 import static dagger.internal.codegen.xprocessing.XFunSpecs.constructorBuilder;
 import static dagger.internal.codegen.xprocessing.XFunSpecs.methodBuilder;
-import static dagger.internal.codegen.xprocessing.XTypeNames.daggerProviderOf;
 import static javax.lang.model.element.Modifier.FINAL;
 import static javax.lang.model.element.Modifier.PRIVATE;
 import static javax.lang.model.element.Modifier.PUBLIC;
@@ -37,9 +36,9 @@
 import androidx.room3.compiler.codegen.XClassName;
 import androidx.room3.compiler.codegen.XCodeBlock;
 import androidx.room3.compiler.codegen.XFunSpec;
+import androidx.room3.compiler.codegen.XPropertySpec;
 import androidx.room3.compiler.codegen.XTypeName;
 import androidx.room3.compiler.codegen.XTypeSpec;
-import androidx.room3.compiler.processing.XProcessingEnv;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
 import dagger.internal.codegen.binding.ContributionBinding;
@@ -49,7 +48,6 @@
 import dagger.internal.codegen.writing.ComponentImplementation.ShardImplementation;
 import dagger.internal.codegen.writing.FrameworkFieldInitializer.FrameworkInstanceCreationExpression;
 import dagger.internal.codegen.xprocessing.XFunSpecs;
-import dagger.internal.codegen.xprocessing.XProcessingEnvs;
 import dagger.internal.codegen.xprocessing.XPropertySpecs;
 import dagger.internal.codegen.xprocessing.XTypeNames;
 import dagger.internal.codegen.xprocessing.XTypeSpecs;
@@ -75,18 +73,13 @@ final class SwitchingProviders {
 
   private final ShardImplementation shardImplementation;
   private final CompilerOptions compilerOptions;
-  private final XProcessingEnv processingEnv;
   private final XTypeName typeVariable;
   private final int maxCasesPerSwitch;
   private final long maxCasesPerClass;
 
-  SwitchingProviders(
-      ShardImplementation shardImplementation,
-      CompilerOptions compilerOptions,
-      XProcessingEnv processingEnv) {
+  SwitchingProviders(ShardImplementation shardImplementation, CompilerOptions compilerOptions) {
     this.shardImplementation = checkNotNull(shardImplementation);
     this.compilerOptions = checkNotNull(compilerOptions);
-    this.processingEnv = checkNotNull(processingEnv);
     this.maxCasesPerSwitch = compilerOptions.casesPerSwitchingProviderSwitch();
     this.maxCasesPerClass = (long) maxCasesPerSwitch * maxCasesPerSwitch;
     this.typeVariable =
@@ -114,7 +107,7 @@ public XCodeBlock creationExpression() {
 
   private SwitchingProviderBuilder getSwitchingProviderBuilder() {
     if (switchingProviderBuilders.size() % maxCasesPerClass == 0) {
-      String name = shardImplementation.getUniqueClassName("SwitchingProvider");
+      String name = shardImplementation.getUniqueClassName("SwitchingPImpl");
       SwitchingProviderBuilder switchingProviderBuilder =
           new SwitchingProviderBuilder(shardImplementation.name().nestedClass(name));
       shardImplementation.addTypeSupplier(switchingProviderBuilder::build);
@@ -130,6 +123,7 @@ private final class SwitchingProviderBuilder {
     private final Map<Integer, XCodeBlock> switchCases = new TreeMap<>();
     private final Map<Key, Integer> switchIds = new HashMap<>();
     private final XClassName switchingProviderType;
+    private XPropertySpec sharedSwitchingProviderProperty;
 
     SwitchingProviderBuilder(XClassName switchingProviderType) {
       this.switchingProviderType = checkNotNull(switchingProviderType);
@@ -144,32 +138,62 @@ private XCodeBlock getNewInstanceCodeBlock(
         switchCases.put(
             switchId, createSwitchCaseCodeBlock(key, unscopedInstanceRequestRepresentation));
       }
-      return XCodeBlock.of(
-          "new %T<%L>(%L, %L)",
-          switchingProviderType,
-          maybeTypeParameter(binding),
-          shardImplementation.componentFieldsByImplementation().values().stream()
-              .map(field -> XCodeBlock.of("%N", field))
-              .collect(toParametersCodeBlock()),
-          switchIds.get(key));
+      int id = switchIds.get(key);
+
+      if (isWrapped(binding)) {
+        return XCodeBlock.of(
+            "%T.provider(%N, %L)",
+            switchingProviderWrapper(binding),
+            getOrCreateSharedSwitchingProvider(),
+            id);
+      } else {
+        return XCodeBlock.of(
+            "new %T<>(%L, %L)", switchingProviderType, getComponentFieldsCodeBlock(), id);
+      }
+    }
+
+    private XPropertySpec getOrCreateSharedSwitchingProvider() {
+      if (sharedSwitchingProviderProperty == null) {
+        String fieldName = shardImplementation.getUniqueFieldName("sharedSwitchingProvider");
+        XCodeBlock args = XCodeBlock.of("%L, -1", getComponentFieldsCodeBlock());
+        sharedSwitchingProviderProperty =
+            XPropertySpecs.builder(
+                    fieldName, switchingProviderType.parametrizedBy(XTypeName.ANY_OBJECT))
+                .addModifiers(PRIVATE)
+                .build();
+        shardImplementation.addField(
+            ComponentImplementation.FieldSpecKind.FRAMEWORK_FIELD, sharedSwitchingProviderProperty);
+        shardImplementation.addInitialization(
+            XCodeBlock.of(
+                "this.%N = new %T<>(%L);",
+                sharedSwitchingProviderProperty,
+                switchingProviderType,
+                args));
+      }
+      return sharedSwitchingProviderProperty;
+    }
+
+    private XClassName switchingProviderWrapper(ContributionBinding binding) {
+      return binding.scope().isPresent() && !binding.scope().get().isReusable()
+          ? XTypeNames.DOUBLE_CHECK_SWITCHING_PROVIDER
+          : XTypeNames.SINGLE_CHECK_SWITCHING_PROVIDER;
+    }
+
+    private boolean isWrapped(ContributionBinding binding) {
+      return binding.scope().isPresent() || binding.kind().equals(BindingKind.ASSISTED_FACTORY);
     }
 
-    private XCodeBlock maybeTypeParameter(ContributionBinding binding) {
-      // Add the type parameter explicitly when the binding is scoped because Java can't
-      // resolve the type when wrapped. For example, the following will error:
-      //   fooProvider = DoubleCheck.provider(new SwitchingProvider<>(1));
-      return (binding.scope().isPresent()
-              || binding.kind().equals(BindingKind.ASSISTED_FACTORY)
-              || XProcessingEnvs.isPreJava8SourceVersion(processingEnv))
-          ? XCodeBlock.of("%T", shardImplementation.accessibleTypeName(binding.contributedType()))
-          : XCodeBlock.of("");
+    private XCodeBlock getComponentFieldsCodeBlock() {
+      return shardImplementation.componentFieldsByImplementation().values().stream()
+          .map(field -> XCodeBlock.of("%N", field))
+          .collect(toParametersCodeBlock());
     }
 
     private XCodeBlock createSwitchCaseCodeBlock(
         Key key, RequestRepresentation unscopedInstanceRequestRepresentation) {
       // TODO(bcorso): Try to delay calling getDependencyExpression() until we are writing out the
       // SwitchingProvider because calling it here makes FrameworkFieldInitializer think there's a
-      // cycle when initializing SwitchingProviders which adds an uncessary DelegateFactory.
+      // cycle when initializing SwitchingProviders which adds an unnecessary DelegateFactory.
       XCodeBlock instanceCodeBlock =
           unscopedInstanceRequestRepresentation
               .getDependencyExpression(switchingProviderType)
@@ -188,7 +212,7 @@ private XTypeSpec build() {
           XTypeSpecs.classBuilder(switchingProviderType)
               .addModifiers(PRIVATE, FINAL, STATIC)
               .addTypeVariable(typeVariable)
-              .addSuperinterface(daggerProviderOf(typeVariable))
+              .addSuperinterface(XTypeNames.SWITCHING_PROVIDER.parametrizedBy(typeVariable))
               .addFunctions(getMethods());
 
       // The SwitchingProvider constructor lists all component parameters first and switch id last.
@@ -199,7 +223,8 @@ private XTypeSpec build() {
           .forEach(
               field -> {
                 builder.addProperty(field);
-                constructor.addParameter(field.getName(), field.getType()); // SUPPRESS_GET_NAME_CHECK
+                constructor.addParameter(
+                    field.getName(), field.getType()); // SUPPRESS_GET_NAME_CHECK
                 constructor.addStatement("this.%1N = %1N", field);
               });
 
@@ -208,37 +233,52 @@ private XTypeSpec build() {
 
     private ImmutableList<XFunSpec> getMethods() {
       ImmutableList<XCodeBlock> switchCodeBlockPartitions = switchCodeBlockPartitions();
+      ImmutableList.Builder<XFunSpec> getMethods = ImmutableList.builder();
+
+      // Add the no-arg get() method for Provider interface
+      getMethods.add(
+          methodBuilder("get")
+              .isOverride(true)
+              .addModifiers(PUBLIC)
+              .returns(typeVariable)
+              .addStatement("return get(id)")
+              .build());
+
       if (switchCodeBlockPartitions.size() == 1) {
         // The case amount does not exceed maxCasesPerSwitch, so no need for extra get methods.
-        return ImmutableList.of(
-            methodBuilder("get")
-                .isOverride(true)
-                .addModifiers(PUBLIC)
-                .addAnnotation(suppressWarnings(UNCHECKED))
-                .returns(typeVariable)
-                .addCode(getOnlyElement(switchCodeBlockPartitions))
-                .build());
+        return getMethods
+            .add(
+                methodBuilder("get")
+                    .isOverride(true)
+                    .addModifiers(PUBLIC)
+                    .addParameter("id", XTypeName.PRIMITIVE_INT)
+                    .addAnnotation(suppressWarnings(UNCHECKED))
+                    .returns(typeVariable)
+                    .addCode(getOnlyElement(switchCodeBlockPartitions))
+                    .build())
+            .build();
       }
 
       // This is the main public "get" method that will route to private getter methods.
       XFunSpecs.Builder routerMethod =
           methodBuilder("get")
               .isOverride(true)
               .addModifiers(PUBLIC)
+              .addParameter("id", XTypeName.PRIMITIVE_INT)
               .returns(typeVariable)
               .beginControlFlow("switch (id / %L)", maxCasesPerSwitch);
 
-      ImmutableList.Builder<XFunSpec> getMethods = ImmutableList.builder();
       for (int i = 0; i < switchCodeBlockPartitions.size(); i++) {
         XFunSpec method =
             methodBuilder("get" + i)
                 .addModifiers(PRIVATE)
+                .addParameter("id", XTypeName.PRIMITIVE_INT)
                 .addAnnotation(suppressWarnings(UNCHECKED))
                 .returns(typeVariable)
                 .addCode(switchCodeBlockPartitions.get(i))
                 .build();
         getMethods.add(method);
-        routerMethod.addStatement("case %L: return %N()", i, method);
+        routerMethod.addStatement("case %L: return %N(id)", i, method);
       }
 
       routerMethod

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

@@ -87,6 +87,13 @@ object XTypeNames {
 
   @JvmField val DELEGATE_FACTORY = XClassName.get("dagger.internal", "DelegateFactory")
   @JvmField val DOUBLE_CHECK = XClassName.get("dagger.internal", "DoubleCheck")
+  @JvmField val SWITCHING_PROVIDER = XClassName.get("dagger.internal", "SwitchingProvider")
+  @JvmField
+  val DOUBLE_CHECK_SWITCHING_PROVIDER =
+    XClassName.get("dagger.internal", "DoubleCheckSwitchingProvider")
+  @JvmField
+  val SINGLE_CHECK_SWITCHING_PROVIDER =
+    XClassName.get("dagger.internal", "SingleCheckSwitchingProvider")
   @JvmField val DAGGER_PRECONDITIONS = XClassName.get("dagger.internal", "Preconditions")
 
   // TODO(b/404613325): Figure out what to do for calls like java.util.Collections.<T>emptyList()
@@ -378,8 +385,7 @@ object XTypeNames {
     )
   }
 
-  @JvmStatic
-  fun XTypeName.isJavaWildcard(): Boolean = toJavaPoet() is JWildcardTypeName
+  @JvmStatic fun XTypeName.isJavaWildcard(): Boolean = toJavaPoet() is JWildcardTypeName
 
   @JvmStatic
   fun XTypeName.boundsHasSelfReference(): Boolean =

dagger-runtime/api/dagger.api

@@ -94,6 +94,11 @@ public final class dagger/internal/DoubleCheck : dagger/Lazy, dagger/internal/Pr
 	public static fun provider (Ljavax/inject/Provider;)Ljavax/inject/Provider;
 }
 
+public final class dagger/internal/DoubleCheckSwitchingProvider : dagger/Lazy, dagger/internal/Provider {
+	public fun get ()Ljava/lang/Object;
+	public static fun provider (Ldagger/internal/SwitchingProvider;I)Ldagger/internal/Provider;
+}
+
 public abstract interface class dagger/internal/Factory : dagger/internal/Provider {
 }
 
@@ -282,6 +287,16 @@ public final class dagger/internal/SingleCheck : dagger/internal/Provider {
 	public static fun provider (Ljavax/inject/Provider;)Ljavax/inject/Provider;
 }
 
+public final class dagger/internal/SingleCheckSwitchingProvider : dagger/internal/Provider {
+	public fun get ()Ljava/lang/Object;
+	public static fun provider (Ldagger/internal/SwitchingProvider;I)Ldagger/internal/Provider;
+}
+
+public abstract interface class dagger/internal/SwitchingProvider : dagger/internal/Provider {
+	public abstract fun get (I)Ljava/lang/Object;
+	public fun get ()Ljava/lang/Object;
+}
+
 public abstract interface annotation class dagger/multibindings/ClassKey : java/lang/annotation/Annotation {
 	public abstract fun value ()Ljava/lang/Class;
 }

dagger-runtime/main/java/dagger/internal/DoubleCheck.java

@@ -26,6 +26,7 @@
  * A {@link Lazy} and {@link Provider} implementation that memoizes the value returned from a
  * delegate using the double-check idiom described in Item 71 of <i>Effective Java 2</i>.
  */
+
 public final class DoubleCheck<T extends @Nullable Object> implements Provider<T>, Lazy<T> {
   private static final Object UNINITIALIZED = new Object();
 
@@ -80,7 +81,7 @@ public T get() {
   public static <T extends @Nullable Object> dagger.internal.Provider<T> provider(
       dagger.internal.Provider<T> delegate) {
     checkNotNull(delegate);
-    if (delegate instanceof DoubleCheck) {
+    if (delegate instanceof DoubleCheck || delegate instanceof DoubleCheckSwitchingProvider) {
       /* This should be a rare case, but if we have a scoped @Binds that delegates to a scoped
        * binding, we shouldn't cache the value again. */
       return delegate;
@@ -121,3 +122,4 @@ public static <P extends javax.inject.Provider<T>, T> Lazy<T> lazy(P provider) {
     return lazy(asDaggerProvider(provider));
   }
 }
+