Add methodhandle support to InternalFactoryToScopedProviderAdapter

lukesandberg · Guice Team · commit a7104775f9b2 · 2025-03-24T06:13:59.000-07:00
We can take advantage of the fact that `setDependency` is a no-op when `disableCircularProxies` is set, but otherwise this is a pretty straightforward translation.

For Singleton scoped factories, we can follow the pattern of `Initializable` and create a 'patchable' callsite which should enable constant folding of resolved singletons.

PiperOrigin-RevId: 739912256
diff --git a/core/src/com/google/inject/internal/InternalFactoryToScopedProviderAdapter.java b/core/src/com/google/inject/internal/InternalFactoryToScopedProviderAdapter.java
@@ -17,11 +17,19 @@
 package com.google.inject.internal;
 
 import static com.google.common.base.Preconditions.checkNotNull;
+import static com.google.inject.internal.InternalMethodHandles.castReturnTo;
+import static com.google.inject.internal.InternalMethodHandles.castReturnToObject;
+import static com.google.inject.internal.InternalMethodHandles.findVirtualOrDie;
+import static java.lang.invoke.MethodType.methodType;
 
+import com.google.errorprone.annotations.Keep;
 import com.google.errorprone.annotations.concurrent.LazyInit;
 import com.google.inject.Provider;
 import com.google.inject.Scope;
 import com.google.inject.spi.Dependency;
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MutableCallSite;
 
 /**
  * A factory that wraps a provider that has been scoped.
@@ -66,6 +74,37 @@ public T get(InternalContext context, Dependency<?> dependency, boolean linked)
     }
   }
 
+  @Override
+  public MethodHandle getHandle(LinkageContext context, Dependency<?> dependency, boolean linked) {
+    // Call Provider.get() on the constant provider
+    // ()->Object
+    var invokeProvider =
+        InternalMethodHandles.getProvider(
+            MethodHandles.constant(jakarta.inject.Provider.class, provider));
+
+    // null check the result using the dependency.
+    // ()->Object
+    invokeProvider = InternalMethodHandles.nullCheckResult(invokeProvider, source, dependency);
+    // Catch any RuntimeException as an InternalProvisionException.
+    // ()->Object
+    invokeProvider =
+        InternalMethodHandles.catchErrorInProviderAndRethrowWithSource(invokeProvider, source);
+    // (InternalContext) -> Object
+    // Add an InternalContext argument to match the factory protocol.
+    invokeProvider = MethodHandles.dropArguments(invokeProvider, 0, InternalContext.class);
+    // We need to call 'setDependency' so it is available to scope implementations and scope
+    // delegate providers. See comment in `get` method for more details.
+    invokeProvider =
+        MethodHandles.foldArguments(
+            invokeProvider,
+            MethodHandles.insertArguments(INTERNAL_CONTEXT_SET_DEPENDENCY_HANDLE, 1, dependency));
+    return castReturnTo(invokeProvider, dependency.getKey().getTypeLiteral().getRawType());
+  }
+
+  private static final MethodHandle INTERNAL_CONTEXT_SET_DEPENDENCY_HANDLE =
+      InternalMethodHandles.findVirtualOrDie(
+          InternalContext.class, "setDependency", methodType(void.class, Dependency.class));
+
   @Override
   public String toString() {
     return provider.toString();
@@ -91,15 +130,101 @@ static final class ForSingletonScope<T> extends InternalFactoryToScopedProviderA
       super(provider, source);
     }
 
+    @Override
+    public T get(InternalContext context, Dependency<?> dependency, boolean linked)
+        throws InternalProvisionException {
+      Object value = this.value;
+      if (value != UNINITIALIZED_VALUE) {
+        // safe because we only store values of T or UNINITIALIZED_VALUE
+        @SuppressWarnings("unchecked")
+        T typedValue = (T) value;
+        if (typedValue == null && !dependency.isNullable()) {
+          InternalProvisionException.onNullInjectedIntoNonNullableDependency(source, dependency);
+        }
+        return typedValue;
+      }
+      T t = super.get(context, dependency, linked);
+      if (!context.areCircularProxiesEnabled() || !BytecodeGen.isCircularProxy(t)) {
+        // Avoid caching circular proxies.
+        this.value = t;
+      }
+      return t;
+    }
+
+    @Override
+    public MethodHandle getHandle(
+        LinkageContext context, Dependency<?> dependency, boolean linked) {
+      // If it is somehow already initialized, we can return a constant handle.
+      Object value = this.value;
+      if (value != UNINITIALIZED_VALUE) {
+        return getHandleForConstant(dependency, source, value);
+      }
+      // Otherwise we bind to a callsite that will patch itself once it is initialized.
+      return new SingletonCallSite(super.getHandle(context, dependency, linked), dependency, source)
+          .dynamicInvoker();
+    }
+
+    private static MethodHandle getHandleForConstant(
+        Dependency<?> dependency, Object source, Object value) {
+      var handle = InternalMethodHandles.constantFactoryGetHandle(dependency, value);
+      handle = InternalMethodHandles.nullCheckResult(handle, source, dependency);
+      return handle;
+    }
+
+    /**
+     * A special callsite that allows us to take advantage of the singleton scope semantics.
+     *
+     * <p>After initialization we can patch the callsite with a constant, which should unlock some
+     * inlining opportunities.
+     */
+    static final class SingletonCallSite extends MutableCallSite {
+      static final MethodHandle BOOTSTRAP_CALL_MH =
+          findVirtualOrDie(
+              SingletonCallSite.class,
+              "boostrapCallSite",
+              methodType(Object.class, InternalContext.class, Object.class));
+
+      private final Dependency<?> dependency;
+      private final Object source;
+
+      SingletonCallSite(MethodHandle actualGetHandle, Dependency<?> dependency, Object source) {
+        super(actualGetHandle.type());
+        this.dependency = dependency;
+        this.source = source;
+        // Invoke the 'actual' handle and then pass the result to the `boostrapCallSite` method.
+        // This will allow us to eventually 'fold' the result into the callsite.
+        // (InternalContext, InternalContext) -> Object
+        var invokeBootstrap =
+            MethodHandles.filterArguments(
+                BOOTSTRAP_CALL_MH.bindTo(this), 1, castReturnToObject(actualGetHandle));
+
+        // (InternalContext, InternalContext) -> T
+        invokeBootstrap = castReturnTo(invokeBootstrap, actualGetHandle.type().returnType());
+        // (InternalContext) -> T
+        invokeBootstrap =
+            MethodHandles.permuteArguments(invokeBootstrap, actualGetHandle.type(), new int[2]);
+        setTarget(invokeBootstrap);
+      }
+
+      @Keep
+      Object boostrapCallSite(InternalContext context, Object result) {
+        // Don't cache circular proxies.
+        if (!context.areCircularProxiesEnabled() || !BytecodeGen.isCircularProxy(result)) {
+          setTarget(getHandleForConstant(dependency, source, result));
+          // This ensures that other threads will see the new target.  This isn't strictly necessary
+          // since the underlying provider is both ThreadSafe and idempotent, but it should improve
+          // performance by giving the JIT and easy optimization opportunity.
+          MutableCallSite.syncAll(new MutableCallSite[] {this});
+        }
+        // otherwise we shouldn't cache the result.
+        return result;
+      }
+    }
+
     private T getAndCache(InjectorImpl injector, Dependency<?> dependency)
         throws InternalProvisionException {
       try (InternalContext context = injector.enterContext()) {
-        T t = get(context, dependency, /* linked= */ false);
-        if (!context.areCircularProxiesEnabled() || !BytecodeGen.isCircularProxy(t)) {
-          // Avoid caching circular proxies.
-          this.value = t;
-        }
-        return t;
+        return get(context, dependency, /* linked= */ false);
       }
     }
 
diff --git a/core/src/com/google/inject/internal/InternalMethodHandles.java b/core/src/com/google/inject/internal/InternalMethodHandles.java
@@ -105,6 +105,20 @@ static MethodHandle castReturnTo(MethodHandle handle, Class<?> returnType) {
     return handle.asType(handle.type().changeReturnType(returnType));
   }
 
+  private static final MethodHandle PROVIDER_GET_HANDLE =
+      findVirtualOrDie(jakarta.inject.Provider.class, "get", methodType(Object.class));
+
+  /**
+   * Returns a method handle that calls {@code <providerHandle>.get()}.
+   *
+   * <p>The returned handle has the same parameters as the delegate and returns {@link Object}.
+   */
+  static MethodHandle getProvider(MethodHandle providerHandle) {
+    return MethodHandles.filterReturnValue(
+        // need to cast to jakarta.inject.Provider to exactly match the signature of Provider.get
+        castReturnTo(providerHandle, jakarta.inject.Provider.class), PROVIDER_GET_HANDLE);
+  }
+
   /** Direct handle for {@link InternalFactory#get} */
   static final MethodHandle INTERNAL_FACTORY_GET_HANDLE =
       findVirtualOrDie(
diff --git a/core/src/com/google/inject/internal/ProviderInternalFactory.java b/core/src/com/google/inject/internal/ProviderInternalFactory.java
@@ -17,7 +17,6 @@
 package com.google.inject.internal;
 
 import static com.google.common.base.Preconditions.checkNotNull;
-import static com.google.inject.internal.InternalMethodHandles.findVirtualOrDie;
 import static java.lang.invoke.MethodType.methodType;
 
 import com.google.inject.Key;
@@ -143,9 +142,6 @@ protected MethodHandle circularGetHandle(
     return createProviderAndPass.asType(InternalMethodHandles.makeFactoryType(dependency));
   }
 
-  private static final MethodHandle PROVIDER_GET_HANDLE =
-      findVirtualOrDie(jakarta.inject.Provider.class, "get", methodType(Object.class));
-
   /**
    * Returns a method handle that calls {@code providerHandle.get()} and catches any
    * RuntimeException as an InternalProvisionException.
@@ -158,7 +154,7 @@ protected MethodHandle provisionHandle(MethodHandle providerHandle, Dependency<?
     // Call Provider.get() and catch any RuntimeException as an InternalProvisionException.
     var invokeProvider =
         InternalMethodHandles.catchErrorInProviderAndRethrowWithSource(
-            MethodHandles.filterReturnValue(providerHandle, PROVIDER_GET_HANDLE), source);
+            InternalMethodHandles.getProvider(providerHandle), source);
     // Wrap in a try..finally.. that calls `finishConstruction` on the context.
     invokeProvider = InternalMethodHandles.finishConstruction(invokeProvider, circularFactoryId);
 
diff --git a/core/test/com/google/inject/BUILD b/core/test/com/google/inject/BUILD
@@ -138,6 +138,7 @@ guice_test_suites(
     name = "gen_tests_stack_trace_%s_use_experimental_method_handles_%s" % (include_stack_trace_option, use_experimental_method_handles_option),
     args = [
         "--guice_include_stack_traces=%s" % include_stack_trace_option,
+        "--guice_use_experimental_method_handles=%s" % use_experimental_method_handles_option,
     ],
     jvm_flags = [
         # those 2 options are required for some tests that checks stack traces
diff --git a/core/test/com/google/inject/errors/BUILD b/core/test/com/google/inject/errors/BUILD
@@ -46,16 +46,20 @@ filegroup(
 )
 
 [guice_test_suites(
-    name = "gen_tests_stack_trace%s" % include_stack_trace_option,
+    name = "gen_tests_stack_trace%s_use_experimental_method_handles_%s" % (include_stack_trace_option, use_experimental_method_handles_option),
     args = [
         "--guice_include_stack_traces=%s" % include_stack_trace_option,
+        "--guice_use_experimental_method_handles=%s" % use_experimental_method_handles_option,
     ],
     sizes = [
         "small",
     ],
-    suffix = "_stack_trace_%s" % include_stack_trace_option,
+    suffix = "_stack_trace_%s_use_experimental_method_handles_%s" % (include_stack_trace_option, use_experimental_method_handles_option),
     deps = [":tests"],
 ) for include_stack_trace_option in [
     "OFF",
     "ONLY_FOR_DECLARING_SOURCE",
+] for use_experimental_method_handles_option in [
+    "NO",
+    "YES",
 ]]
diff --git a/core/test/com/google/inject/errors/NullInjectedIntoNonNullableTest.java b/core/test/com/google/inject/errors/NullInjectedIntoNonNullableTest.java
@@ -1,5 +1,6 @@
 package com.google.inject.errors;
 
+import static com.google.common.truth.Truth.assertThat;
 import static com.google.inject.errors.ErrorMessageTestUtils.assertGuiceErrorEqualsIgnoreLineNumber;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assume.assumeTrue;
@@ -10,6 +11,7 @@
 import com.google.inject.Module;
 import com.google.inject.Provides;
 import com.google.inject.ProvisionException;
+import com.google.inject.Singleton;
 import com.google.inject.internal.InternalFlags;
 import com.google.inject.internal.InternalFlags.IncludeStackTraceOption;
 import java.lang.annotation.Retention;
@@ -60,7 +62,7 @@ protected void configure() {
             try {
               field = IntermediateModule.class.getField("NULL");
             } catch (NoSuchFieldException error) {
-              throw new AssertionError("NULL field missing!");
+              throw new AssertionError("NULL field missing!", error);
             }
             binder()
                 .withSource(field)
@@ -110,4 +112,31 @@ public void nullReturnedFromProviderWithModuleStack() {
     assertGuiceErrorEqualsIgnoreLineNumber(
         exception.getMessage(), "null_returned_from_provider_with_module_stack.txt");
   }
+
+  // Ensure we report the error when the dependency is a singleton, no matter how many times we
+  // try to provision it.
+  @Test
+  public void nullReturnedFromSingletonBinding() {
+    Injector injector =
+        Guice.createInjector(
+            new AbstractModule() {
+              @Override
+              protected void configure() {
+                bind(String.class)
+                    .annotatedWith(Bar.class)
+                    .toProvider(() -> null)
+                    .in(Singleton.class);
+              }
+
+              @Provides
+              String provideString(@Bar String string) {
+                return string;
+              }
+            });
+    var pe = assertThrows(ProvisionException.class, () -> injector.getInstance(String.class));
+    assertThat(pe).hasMessageThat().contains("[Guice/NullInjectedIntoNonNullable]");
+
+    pe = assertThrows(ProvisionException.class, () -> injector.getInstance(String.class));
+    assertThat(pe).hasMessageThat().contains("[Guice/NullInjectedIntoNonNullable]");
+  }
 }
