Link factoryMethod consistently in AOT-generated bean definitions

Closes gh-28748

spring-beans/src/main/java/org/springframework/beans/factory/aot/AutowiredArguments.java

@@ -26,7 +26,7 @@ import org.springframework.util.ClassUtils;
  * @author Phillip Webb
  * @author Stephane Nicoll
  * @since 6.0
- * @see AutowiredInstantiationArgumentsResolver
+ * @see BeanInstanceSupplier
  * @see AutowiredMethodArgumentsResolver
  */
 @FunctionalInterface

spring-beans/src/main/java/org/springframework/beans/factory/aot/AutowiredMethodArgumentsResolver.java

@@ -105,7 +105,7 @@ public final class AutowiredMethodArgumentsResolver extends AutowiredElementReso
 	}
 
 	/**
-	 * Return a new {@link AutowiredInstantiationArgumentsResolver} instance
+	 * Return a new {@link AutowiredMethodArgumentsResolver} instance
 	 * that uses direct bean name injection shortcuts for specific parameters.
 	 * @param beanNames the bean names to use as shortcuts (aligned with the
 	 * method parameters)

spring-beans/src/main/java/org/springframework/beans/factory/aot/BeanInstanceSupplier.java

@@ -40,6 +40,7 @@ import org.springframework.beans.factory.config.ConstructorArgumentValues.ValueH
 import org.springframework.beans.factory.config.DependencyDescriptor;
 import org.springframework.beans.factory.support.AbstractAutowireCapableBeanFactory;
 import org.springframework.beans.factory.support.BeanDefinitionValueResolver;
+import org.springframework.beans.factory.support.InstanceSupplier;
 import org.springframework.beans.factory.support.RegisteredBean;
 import org.springframework.beans.factory.support.RootBeanDefinition;
 import org.springframework.core.CollectionFactory;
@@ -49,67 +50,81 @@ import org.springframework.util.Assert;
 import org.springframework.util.ClassUtils;
 import org.springframework.util.ObjectUtils;
 import org.springframework.util.ReflectionUtils;
+import org.springframework.util.function.ThrowingBiFunction;
 import org.springframework.util.function.ThrowingFunction;
+import org.springframework.util.function.ThrowingSupplier;
 
 /**
- * Resolver used to support the autowiring of constructors or factory methods.
- * Typically used in AOT-processed applications as a targeted alternative to the
- * reflection based injection.
+ * Specialized {@link InstanceSupplier} that provides the factory {@link Method}
+ * used to instantiate the underlying bean instance, if any. Transparently
+ * handles resolution of {@link AutowiredArguments} if necessary. Typically used
+ * in AOT-processed applications as a targeted alternative to the reflection
+ * based injection.
  * <p>
- * When resolving arguments in a native image, the {@link Constructor} or
- * {@link Method} being used must be marked with an
- * {@link ExecutableMode#INTROSPECT introspection} hint so that parameter
- * annotations can be read. Full {@link ExecutableMode#INVOKE invocation} hints
- * are only required if the {@code resolveAndInstantiate} methods of this class
- * are being used (typically to support private constructors, methods or
- * classes).
+ * If no {@code generator} is provided, reflection is used to instantiate the
+ * bean instance, and full {@link ExecutableMode#INVOKE invocation} hints are
+ * contributed. Multiple generator callback styles are supported:
+ * <ul>
+ * <li>A function with the {@code registeredBean} and resolved {@code arguments}
+ * for executables that require arguments resolution. An
+ * {@link ExecutableMode#INTROSPECT introspection} hint is added so that
+ * parameter annotations can be read </li>
+ * <li>A function with only the {@code registeredBean} for simpler cases that
+ * do not require resolution of arguments</li>
+ * <li>A supplier when a method reference can be used</li>
+ * </ul>
+ * Generator callbacks handle checked exceptions so that the caller does not
+ * have to deal with it.
  *
  * @author Phillip Webb
  * @author Stephane Nicoll
  * @since 6.0
  * @see AutowiredArguments
  */
-public final class AutowiredInstantiationArgumentsResolver extends AutowiredElementResolver {
+public final class BeanInstanceSupplier extends AutowiredElementResolver implements InstanceSupplier<Object> {
 
 	private final ExecutableLookup lookup;
 
+	@Nullable
+	private final ThrowingBiFunction<RegisteredBean, AutowiredArguments, Object> generator;
+
 	@Nullable
 	private final String[] shortcuts;
 
 
-	private AutowiredInstantiationArgumentsResolver(ExecutableLookup lookup,
+	private BeanInstanceSupplier(ExecutableLookup lookup,
+			@Nullable ThrowingBiFunction<RegisteredBean, AutowiredArguments, Object> generator,
 			@Nullable String[] shortcuts) {
-
 		this.lookup = lookup;
+		this.generator = generator;
 		this.shortcuts = shortcuts;
 	}
 
-
 	/**
-	 * Create a {@link AutowiredInstantiationArgumentsResolver} that resolves
+	 * Create a {@link BeanInstanceSupplier} that resolves
 	 * arguments for the specified bean constructor.
 	 * @param parameterTypes the constructor parameter types
-	 * @return a new {@link AutowiredInstantiationArgumentsResolver} instance
+	 * @return a new {@link BeanInstanceSupplier} instance
 	 */
-	public static AutowiredInstantiationArgumentsResolver forConstructor(
+	public static BeanInstanceSupplier forConstructor(
 			Class<?>... parameterTypes) {
 
 		Assert.notNull(parameterTypes, "'parameterTypes' must not be null");
 		Assert.noNullElements(parameterTypes,
 				"'parameterTypes' must not contain null elements");
-		return new AutowiredInstantiationArgumentsResolver(
-				new ConstructorLookup(parameterTypes), null);
+		return new BeanInstanceSupplier(
+				new ConstructorLookup(parameterTypes), null, null);
 	}
 
 	/**
-	 * Create a new {@link AutowiredInstantiationArgumentsResolver} that
+	 * Create a new {@link BeanInstanceSupplier} that
 	 * resolves arguments for the specified factory method.
 	 * @param declaringClass the class that declares the factory method
 	 * @param methodName the factory method name
 	 * @param parameterTypes the factory method parameter types
-	 * @return a new {@link AutowiredInstantiationArgumentsResolver} instance
+	 * @return a new {@link BeanInstanceSupplier} instance
 	 */
-	public static AutowiredInstantiationArgumentsResolver forFactoryMethod(
+	public static BeanInstanceSupplier forFactoryMethod(
 			Class<?> declaringClass, String methodName, Class<?>... parameterTypes) {
 
 		Assert.notNull(declaringClass, "'declaringClass' must not be null");
@@ -117,9 +132,9 @@ public final class AutowiredInstantiationArgumentsResolver extends AutowiredElem
 		Assert.notNull(parameterTypes, "'parameterTypes' must not be null");
 		Assert.noNullElements(parameterTypes,
 				"'parameterTypes' must not contain null elements");
-		return new AutowiredInstantiationArgumentsResolver(
+		return new BeanInstanceSupplier(
 				new FactoryMethodLookup(declaringClass, methodName, parameterTypes),
-				null);
+				null, null);
 	}
 
 
@@ -128,32 +143,82 @@ public final class AutowiredInstantiationArgumentsResolver extends AutowiredElem
 	}
 
 	/**
-	 * Return a new {@link AutowiredInstantiationArgumentsResolver} instance
-	 * that uses direct bean name injection shortcuts for specific parameters.
-	 * @param beanNames the bean names to use as shortcuts (aligned with the
-	 * constructor or factory method parameters)
-	 * @return a new {@link AutowiredInstantiationArgumentsResolver} instance
-	 * that uses the shortcuts
+	 * Return a new {@link BeanInstanceSupplier} instance that uses the specified
+	 * {@code generator} bi-function to instantiate the underlying bean.
+	 * @param generator a {@link ThrowingBiFunction} that uses the
+	 * {@link RegisteredBean} and resolved {@link AutowiredArguments} to
+	 * instantiate the underlying bean
+	 * @return a new {@link BeanInstanceSupplier} instance with the specified
+	 * generator
 	 */
-	public AutowiredInstantiationArgumentsResolver withShortcuts(String... beanNames) {
-		return new AutowiredInstantiationArgumentsResolver(this.lookup, beanNames);
+	public BeanInstanceSupplier withGenerator(
+			ThrowingBiFunction<RegisteredBean, AutowiredArguments, Object> generator) {
+		Assert.notNull(generator, "'generator' must not be null");
+		return new BeanInstanceSupplier(this.lookup, generator, this.shortcuts);
 	}
 
 	/**
-	 * Resolve arguments for the specified registered bean and provide them to
-	 * the given generator in order to return a result.
-	 * @param registeredBean the registered bean
-	 * @param generator the generator to execute with the resolved constructor
-	 * or factory method arguments
+	 * Return a new {@link BeanInstanceSupplier} instance that uses the specified
+	 * {@code generator} function to instantiate the underlying bean.
+	 * @param generator a {@link ThrowingFunction} that uses the
+	 * {@link RegisteredBean} to instantiate the underlying bean
+	 * @return a new {@link BeanInstanceSupplier} instance with the specified
+	 * generator
 	 */
-	public <T> T resolve(RegisteredBean registeredBean,
-			ThrowingFunction<AutowiredArguments, T> generator) {
+	public BeanInstanceSupplier withGenerator(
+			ThrowingFunction<RegisteredBean, Object> generator) {
+		Assert.notNull(generator, "'generator' must not be null");
+		return new BeanInstanceSupplier(this.lookup, (registeredBean, args) ->
+				generator.apply(registeredBean), this.shortcuts);
+	}
 
+	/**
+	 * Return a new {@link BeanInstanceSupplier} instance that uses the specified
+	 * {@code generator} supplier to instantiate the underlying bean.
+	 * @param generator a {@link ThrowingSupplier} to instantiate the underlying
+	 * bean
+	 * @return a new {@link BeanInstanceSupplier} instance with the specified
+	 * generator
+	 */
+	public BeanInstanceSupplier withGenerator(ThrowingSupplier<Object> generator) {
+		Assert.notNull(generator, "'generator' must not be null");
+		return new BeanInstanceSupplier(this.lookup, (registeredBean, args) ->
+				generator.get(), this.shortcuts);
+	}
+
+	/**
+	 * Return a new {@link BeanInstanceSupplier} instance
+	 * that uses direct bean name injection shortcuts for specific parameters.
+	 * @param beanNames the bean names to use as shortcuts (aligned with the
+	 * constructor or factory method parameters)
+	 * @return a new {@link BeanInstanceSupplier} instance
+	 * that uses the shortcuts
+	 */
+	public BeanInstanceSupplier withShortcuts(String... beanNames) {
+		return new BeanInstanceSupplier(this.lookup, this.generator, beanNames);
+	}
+
+	@Override
+	public Object get(RegisteredBean registeredBean) throws Exception {
 		Assert.notNull(registeredBean, "'registeredBean' must not be null");
-		Assert.notNull(generator, "'action' must not be null");
-		AutowiredArguments resolved = resolveArguments(registeredBean,
-				this.lookup.get(registeredBean));
-		return generator.apply(resolved);
+		Executable executable = this.lookup.get(registeredBean);
+		AutowiredArguments arguments = resolveArguments(registeredBean, executable);
+		if (this.generator != null) {
+			return this.generator.apply(registeredBean, arguments);
+		}
+		else {
+			return instantiate(registeredBean.getBeanFactory(), executable,
+					arguments.toArray());
+		}
+	}
+
+	@Nullable
+	@Override
+	public Method getFactoryMethod() {
+		if (this.lookup instanceof FactoryMethodLookup factoryMethodLookup) {
+			return factoryMethodLookup.get();
+		}
+		return null;
 	}
 
 	/**
@@ -161,43 +226,11 @@ public final class AutowiredInstantiationArgumentsResolver extends AutowiredElem
 	 * @param registeredBean the registered bean
 	 * @return the resolved constructor or factory method arguments
 	 */
-	public AutowiredArguments resolve(RegisteredBean registeredBean) {
+	AutowiredArguments resolveArguments(RegisteredBean registeredBean) {
 		Assert.notNull(registeredBean, "'registeredBean' must not be null");
 		return resolveArguments(registeredBean, this.lookup.get(registeredBean));
 	}
 
-	/**
-	 * Resolve arguments for the specified registered bean and instantiate a new
-	 * instance using reflection.
-	 * @param registeredBean the registered bean
-	 * @return an instance of the bean
-	 */
-	@SuppressWarnings("unchecked")
-	public <T> T resolveAndInstantiate(RegisteredBean registeredBean) {
-		return (T) resolveAndInstantiate(registeredBean, Object.class);
-	}
-
-	/**
-	 * Resolve arguments for the specified registered bean and instantiate a new
-	 * instance using reflection.
-	 * @param registeredBean the registered bean
-	 * @param requiredType the required result type
-	 * @return an instance of the bean
-	 */
-	@SuppressWarnings("unchecked")
-	public <T> T resolveAndInstantiate(RegisteredBean registeredBean,
-			Class<T> requiredType) {
-
-		Assert.notNull(registeredBean, "'registeredBean' must not be null");
-		Assert.notNull(registeredBean, "'requiredType' must not be null");
-		Executable executable = this.lookup.get(registeredBean);
-		AutowiredArguments arguments = resolveArguments(registeredBean, executable);
-		Object instance = instantiate(registeredBean.getBeanFactory(), executable,
-				arguments.toArray());
-		Assert.isInstanceOf(requiredType, instance);
-		return (T) instance;
-	}
-
 	private AutowiredArguments resolveArguments(RegisteredBean registeredBean,
 			Executable executable) {
 
@@ -233,9 +266,6 @@ public final class AutowiredInstantiationArgumentsResolver extends AutowiredElem
 					autowiredBeans, parameter, dependencyDescriptor, argumentValue);
 		}
 		registerDependentBeans(beanFactory, beanName, autowiredBeans);
-		if (executable instanceof Method method) {
-			mergedBeanDefinition.setResolvedFactoryMethod(method);
-		}
 		return AutowiredArguments.of(resolved);
 	}
 
@@ -403,10 +433,12 @@ public final class AutowiredInstantiationArgumentsResolver extends AutowiredElem
 
 		private final Class<?>[] parameterTypes;
 
+
 		ConstructorLookup(Class<?>[] parameterTypes) {
 			this.parameterTypes = parameterTypes;
 		}
 
+
 		@Override
 		public Executable get(RegisteredBean registeredBean) {
 			Class<?> beanClass = registeredBean.getBeanClass();
@@ -453,6 +485,10 @@ public final class AutowiredInstantiationArgumentsResolver extends AutowiredElem
 
 		@Override
 		public Executable get(RegisteredBean registeredBean) {
+			return get();
+		}
+
+		Method get() {
 			Method method = ReflectionUtils.findMethod(this.declaringClass,
 					this.methodName, this.parameterTypes);
 			Assert.notNull(method, () -> String.format("%s cannot be found", this));

spring-beans/src/main/java/org/springframework/beans/factory/aot/InstanceSupplierCodeGenerator.java

@@ -35,17 +35,22 @@ import org.springframework.beans.factory.support.RegisteredBean;
 import org.springframework.core.ResolvableType;
 import org.springframework.javapoet.ClassName;
 import org.springframework.javapoet.CodeBlock;
+import org.springframework.javapoet.CodeBlock.Builder;
 import org.springframework.javapoet.MethodSpec;
-import org.springframework.javapoet.MethodSpec.Builder;
 import org.springframework.util.ClassUtils;
 import org.springframework.util.function.ThrowingSupplier;
 
 /**
- * Internal code generator to create an {@link InstanceSupplier}.
+ * Internal code generator to create an {@link InstanceSupplier}, usually in
+ * the form of a {@link BeanInstanceSupplier} that retains the executable
+ * that is used to instantiate the bean.
  * <p>
- * Generates code in the form:<pre class="code">{@code
- * InstanceSupplier.of(TheGeneratedClass::getMyBeanInstance);
- * }</pre>
+ * Generated code is usually a method reference that generate the
+ * {@link BeanInstanceSupplier}, but some shortcut can be used as well such
+ * as:
+ * <pre class="code">
+ * {@code InstanceSupplier.of(TheGeneratedClass::getMyBeanInstance);}
+ * </pre>
  *
  * @author Phillip Webb
  * @author Stephane Nicoll
@@ -55,6 +60,8 @@ class InstanceSupplierCodeGenerator {
 
 	private static final String REGISTERED_BEAN_PARAMETER_NAME = "registeredBean";
 
+	private static final String ARGS_PARAMETER_NAME = "args";
+
 	private static final javax.lang.model.element.Modifier[] PRIVATE_STATIC = {
 			javax.lang.model.element.Modifier.PRIVATE,
 			javax.lang.model.element.Modifier.STATIC };
@@ -109,15 +116,14 @@ class InstanceSupplierCodeGenerator {
 				constructor);
 		if (accessVisibility == AccessVisibility.PUBLIC
 				|| accessVisibility == AccessVisibility.PACKAGE_PRIVATE) {
-			return generateCodeForAccessibleConstructor(name, constructor, declaringClass,
-					dependsOnBean);
+			return generateCodeForAccessibleConstructor(name, constructor, dependsOnBean,
+					declaringClass);
 		}
-		return generateCodeForInaccessibleConstructor(name, constructor, declaringClass,
-				dependsOnBean);
+		return generateCodeForInaccessibleConstructor(name, constructor, dependsOnBean);
 	}
 
 	private CodeBlock generateCodeForAccessibleConstructor(String name,
-			Constructor<?> constructor, Class<?> declaringClass, boolean dependsOnBean) {
+			Constructor<?> constructor, boolean dependsOnBean, Class<?> declaringClass) {
 
 		this.generationContext.getRuntimeHints().reflection()
 				.registerConstructor(constructor, INTROSPECT);
@@ -132,60 +138,47 @@ class InstanceSupplierCodeGenerator {
 			return CodeBlock.of("$T.of($T::new)", ThrowingSupplier.class,
 					declaringClass);
 		}
-		GeneratedMethod generatedMethod = generateGetInstanceMethod(method ->
-				buildGetInstanceMethodForConstructor(method, name, constructor, declaringClass,
-						dependsOnBean, PRIVATE_STATIC));
-		return CodeBlock.of("$T.of($T::$L)", InstanceSupplier.class, this.className,
-				generatedMethod.getName());
+		GeneratedMethod generatedMethod = generateGetInstanceSupplierMethod(method ->
+				buildGetInstanceMethodForConstructor(method, name, constructor,
+						declaringClass, dependsOnBean, PRIVATE_STATIC));
+		return generateReturnStatement(generatedMethod);
 	}
 
 	private CodeBlock generateCodeForInaccessibleConstructor(String name,
-			Constructor<?> constructor, Class<?> declaringClass, boolean dependsOnBean) {
+			Constructor<?> constructor, boolean dependsOnBean) {
 
 		this.generationContext.getRuntimeHints().reflection()
 				.registerConstructor(constructor);
-		GeneratedMethod generatedMethod = generateGetInstanceMethod(method -> {
-			method.addJavadoc("Instantiate the bean instance for '$L'.", name);
+		GeneratedMethod generatedMethod = generateGetInstanceSupplierMethod(method -> {
+			method.addJavadoc("Get the bean instance supplier for '$L'.", name);
 			method.addModifiers(PRIVATE_STATIC);
-			method.returns(declaringClass);
-			method.addParameter(RegisteredBean.class, REGISTERED_BEAN_PARAMETER_NAME);
+			method.returns(BeanInstanceSupplier.class);
 			int parameterOffset = (!dependsOnBean) ? 0 : 1;
 			method.addStatement(
 					generateResolverForConstructor(constructor, parameterOffset));
-			method.addStatement("return resolver.resolveAndInstantiate($L)",
-					REGISTERED_BEAN_PARAMETER_NAME);
 		});
-		return CodeBlock.of("$T.of($T::$L)", InstanceSupplier.class, this.className,
-				generatedMethod.getName());
+		return generateReturnStatement(generatedMethod);
 	}
 
 	private void buildGetInstanceMethodForConstructor(MethodSpec.Builder method,
 			String name, Constructor<?> constructor, Class<?> declaringClass,
 			boolean dependsOnBean, javax.lang.model.element.Modifier... modifiers) {
 
-		method.addJavadoc("Create the bean instance for '$L'.", name);
+		method.addJavadoc("Get the bean instance supplier for '$L'.", name);
 		method.addModifiers(modifiers);
-		method.returns(declaringClass);
-		method.addParameter(RegisteredBean.class, REGISTERED_BEAN_PARAMETER_NAME);
-		if (constructor.getParameterCount() == 0) {
-			CodeBlock instantiationCode = generateNewInstanceCodeForConstructor(
-					dependsOnBean, declaringClass, NO_ARGS);
-			method.addCode(generateReturnStatement(instantiationCode));
-		}
-		else {
-			int parameterOffset = (!dependsOnBean) ? 0 : 1;
-			CodeBlock.Builder code = CodeBlock.builder();
-			code.addStatement(
-					generateResolverForConstructor(constructor, parameterOffset));
-			CodeBlock arguments = new AutowiredArgumentsCodeGenerator(declaringClass,
-					constructor).generateCode(constructor.getParameterTypes(),
-							parameterOffset);
-			CodeBlock newInstance = generateNewInstanceCodeForConstructor(dependsOnBean,
-					declaringClass, arguments);
-			code.addStatement("return resolver.resolve($L, (args) -> $L)",
-					REGISTERED_BEAN_PARAMETER_NAME, newInstance);
-			method.addCode(code.build());
-		}
+		method.returns(BeanInstanceSupplier.class);
+		int parameterOffset = (!dependsOnBean) ? 0 : 1;
+		CodeBlock.Builder code = CodeBlock.builder();
+		code.add(generateResolverForConstructor(constructor, parameterOffset));
+		boolean hasArguments = constructor.getParameterCount() > 0;
+		CodeBlock arguments = hasArguments
+				? new AutowiredArgumentsCodeGenerator(declaringClass, constructor)
+						.generateCode(constructor.getParameterTypes(), parameterOffset)
+				: NO_ARGS;
+		CodeBlock newInstance = generateNewInstanceCodeForConstructor(dependsOnBean,
+				declaringClass, arguments);
+		code.add(generateWithGeneratorCode(hasArguments, newInstance));
+		method.addStatement(code.build());
 	}
 
 	private CodeBlock generateResolverForConstructor(Constructor<?> constructor,
@@ -193,9 +186,8 @@ class InstanceSupplierCodeGenerator {
 
 		CodeBlock parameterTypes = generateParameterTypesCode(
 				constructor.getParameterTypes(), parameterOffset);
-		return CodeBlock.of("$T resolver = $T.forConstructor($L)",
-				AutowiredInstantiationArgumentsResolver.class,
-				AutowiredInstantiationArgumentsResolver.class, parameterTypes);
+		return CodeBlock.of("return $T.forConstructor($L)",
+				BeanInstanceSupplier.class, parameterTypes);
 	}
 
 	private CodeBlock generateNewInstanceCodeForConstructor(boolean dependsOnBean,
@@ -232,21 +224,16 @@ class InstanceSupplierCodeGenerator {
 		this.generationContext.getRuntimeHints().reflection()
 				.registerMethod(factoryMethod, INTROSPECT);
 		if (!dependsOnBean && factoryMethod.getParameterCount() == 0) {
-			if (!this.allowDirectSupplierShortcut) {
-				return CodeBlock.of("$T.using($T::$L)", InstanceSupplier.class,
-						declaringClass, factoryMethod.getName());
-			}
-			if (!isThrowingCheckedException(factoryMethod)) {
-				return CodeBlock.of("$T::$L", declaringClass, factoryMethod.getName());
-			}
-			return CodeBlock.of("$T.of($T::$L)", ThrowingSupplier.class, declaringClass,
-					factoryMethod.getName());
+			CodeBlock.Builder code = CodeBlock.builder();
+			code.add("$T.forFactoryMethod($T.class, $S)", BeanInstanceSupplier.class,
+					declaringClass, factoryMethod.getName());
+			code.add(".withGenerator($T::$L)", declaringClass, factoryMethod.getName());
+			return code.build();
 		}
-		GeneratedMethod generatedMethod = generateGetInstanceMethod(method ->
-				buildGetInstanceMethodForFactoryMethod(method, name, factoryMethod, declaringClass,
-						dependsOnBean, PRIVATE_STATIC));
-		return CodeBlock.of("$T.of($T::$L)", InstanceSupplier.class, this.className,
-				generatedMethod.getName());
+		GeneratedMethod getInstanceMethod = generateGetInstanceSupplierMethod(method ->
+				buildGetInstanceMethodForFactoryMethod(method, name, factoryMethod,
+						declaringClass, dependsOnBean, PRIVATE_STATIC));
+		return generateReturnStatement(getInstanceMethod);
 	}
 
 	private CodeBlock generateCodeForInaccessibleFactoryMethod(String name,
@@ -254,18 +241,14 @@ class InstanceSupplierCodeGenerator {
 
 		this.generationContext.getRuntimeHints().reflection()
 				.registerMethod(factoryMethod);
-		GeneratedMethod generatedMethod = generateGetInstanceMethod(method -> {
-			method.addJavadoc("Instantiate the bean instance for '$L'.", name);
+		GeneratedMethod getInstanceMethod = generateGetInstanceSupplierMethod(method -> {
+			method.addJavadoc("Get the bean instance supplier for '$L'.", name);
 			method.addModifiers(PRIVATE_STATIC);
-			method.returns(factoryMethod.getReturnType());
-			method.addParameter(RegisteredBean.class, REGISTERED_BEAN_PARAMETER_NAME);
-			method.addStatement(generateResolverForFactoryMethod(factoryMethod,
+			method.returns(BeanInstanceSupplier.class);
+			method.addStatement(generateInstanceSupplierForFactoryMethod(factoryMethod,
 					declaringClass, factoryMethod.getName()));
-			method.addStatement("return resolver.resolveAndInstantiate($L)",
-					REGISTERED_BEAN_PARAMETER_NAME);
 		});
-		return CodeBlock.of("$T.of($T::$L)", InstanceSupplier.class, this.className,
-				generatedMethod.getName());
+		return generateReturnStatement(getInstanceMethod);
 	}
 
 	private void buildGetInstanceMethodForFactoryMethod(MethodSpec.Builder method,
@@ -273,46 +256,34 @@ class InstanceSupplierCodeGenerator {
 			boolean dependsOnBean, javax.lang.model.element.Modifier... modifiers) {
 
 		String factoryMethodName = factoryMethod.getName();
-		method.addJavadoc("Get the bean instance for '$L'.", name);
+		method.addJavadoc("Get the bean instance supplier for '$L'.", name);
 		method.addModifiers(modifiers);
-		method.returns(factoryMethod.getReturnType());
-		if (isThrowingCheckedException(factoryMethod)) {
-			method.addException(Exception.class);
-		}
-		method.addParameter(RegisteredBean.class, REGISTERED_BEAN_PARAMETER_NAME);
-		if (factoryMethod.getParameterCount() == 0) {
-			CodeBlock instantiationCode = generateNewInstanceCodeForMethod(dependsOnBean,
-					declaringClass, factoryMethodName, NO_ARGS);
-			method.addCode(generateReturnStatement(instantiationCode));
-		}
-		else {
-			CodeBlock.Builder code = CodeBlock.builder();
-			code.addStatement(generateResolverForFactoryMethod(factoryMethod,
-					declaringClass, factoryMethodName));
-			CodeBlock arguments = new AutowiredArgumentsCodeGenerator(declaringClass,
-					factoryMethod).generateCode(factoryMethod.getParameterTypes());
-			CodeBlock newInstance = generateNewInstanceCodeForMethod(dependsOnBean,
-					declaringClass, factoryMethodName, arguments);
-			code.addStatement("return resolver.resolve($L, (args) -> $L)",
-					REGISTERED_BEAN_PARAMETER_NAME, newInstance);
-			method.addCode(code.build());
-		}
+		method.returns(BeanInstanceSupplier.class);
+		CodeBlock.Builder code = CodeBlock.builder();
+		code.add(generateInstanceSupplierForFactoryMethod(factoryMethod, declaringClass, factoryMethodName));
+		boolean hasArguments = factoryMethod.getParameterCount() > 0;
+		CodeBlock arguments = hasArguments
+				? new AutowiredArgumentsCodeGenerator(declaringClass, factoryMethod)
+						.generateCode(factoryMethod.getParameterTypes())
+				: NO_ARGS;
+		CodeBlock newInstance = generateNewInstanceCodeForMethod(dependsOnBean,
+				declaringClass, factoryMethodName, arguments);
+		code.add(generateWithGeneratorCode(hasArguments, newInstance));
+		method.addStatement(code.build());
 	}
 
-	private CodeBlock generateResolverForFactoryMethod(Method factoryMethod,
+	private CodeBlock generateInstanceSupplierForFactoryMethod(Method factoryMethod,
 			Class<?> declaringClass, String factoryMethodName) {
 
 		if (factoryMethod.getParameterCount() == 0) {
-			return CodeBlock.of("$T resolver = $T.forFactoryMethod($T.class, $S)",
-					AutowiredInstantiationArgumentsResolver.class,
-					AutowiredInstantiationArgumentsResolver.class, declaringClass,
+			return CodeBlock.of("return $T.forFactoryMethod($T.class, $S)",
+					BeanInstanceSupplier.class, declaringClass,
 					factoryMethodName);
 		}
 		CodeBlock parameterTypes = generateParameterTypesCode(
 				factoryMethod.getParameterTypes(), 0);
-		return CodeBlock.of("$T resolver = $T.forFactoryMethod($T.class, $S, $L)",
-				AutowiredInstantiationArgumentsResolver.class,
-				AutowiredInstantiationArgumentsResolver.class, declaringClass,
+		return CodeBlock.of("return $T.forFactoryMethod($T.class, $S, $L)",
+				BeanInstanceSupplier.class, declaringClass,
 				factoryMethodName, parameterTypes);
 	}
 
@@ -326,9 +297,19 @@ class InstanceSupplierCodeGenerator {
 				REGISTERED_BEAN_PARAMETER_NAME, declaringClass, factoryMethodName, args);
 	}
 
-	private CodeBlock generateReturnStatement(CodeBlock instantiationCode) {
-		CodeBlock.Builder code = CodeBlock.builder();
-		code.addStatement("return $L", instantiationCode);
+	private CodeBlock generateReturnStatement(GeneratedMethod getInstanceMethod) {
+		return CodeBlock.of("$T.$L()", this.className, getInstanceMethod.getName());
+	}
+
+	private CodeBlock generateWithGeneratorCode(boolean hasArguments, CodeBlock newInstance) {
+		CodeBlock lambdaArguments = (hasArguments
+				? CodeBlock.of("($L, $L)", REGISTERED_BEAN_PARAMETER_NAME, ARGS_PARAMETER_NAME)
+				: CodeBlock.of("($L)", REGISTERED_BEAN_PARAMETER_NAME));
+		Builder code = CodeBlock.builder();
+		code.add("\n");
+		code.indent().indent();
+		code.add(".withGenerator($L -> $L)", lambdaArguments, newInstance);
+		code.unindent().unindent();
 		return code.build();
 	}
 
@@ -342,16 +323,16 @@ class InstanceSupplierCodeGenerator {
 	}
 
 	private CodeBlock generateParameterTypesCode(Class<?>[] parameterTypes, int offset) {
-		CodeBlock.Builder builder = CodeBlock.builder();
+		CodeBlock.Builder code = CodeBlock.builder();
 		for (int i = offset; i < parameterTypes.length; i++) {
-			builder.add(i != offset ? ", " : "");
-			builder.add("$T.class", parameterTypes[i]);
+			code.add(i != offset ? ", " : "");
+			code.add("$T.class", parameterTypes[i]);
 		}
-		return builder.build();
+		return code.build();
 	}
 
-	private GeneratedMethod generateGetInstanceMethod(Consumer<Builder> method) {
-		return this.generatedMethods.add("getInstance", method);
+	private GeneratedMethod generateGetInstanceSupplierMethod(Consumer<MethodSpec.Builder> method) {
+		return this.generatedMethods.add("getInstanceSupplier", method);
 	}
 
 	private boolean isThrowingCheckedException(Executable executable) {

spring-beans/src/main/java/org/springframework/beans/factory/support/InstanceSupplier.java

@@ -16,8 +16,10 @@
 
 package org.springframework.beans.factory.support;
 
+import java.lang.reflect.Method;
 import java.util.function.Supplier;
 
+import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 import org.springframework.util.function.ThrowingBiFunction;
 import org.springframework.util.function.ThrowingSupplier;
@@ -29,6 +31,7 @@ import org.springframework.util.function.ThrowingSupplier;
  * supply the instance.
  *
  * @author Phillip Webb
+ * @author Stephane Nicoll
  * @since 6.0
  * @param <T> the type of instance supplied by this supplier
  * @see RegisteredBean
@@ -49,6 +52,17 @@ public interface InstanceSupplier<T> extends ThrowingSupplier<T> {
 	 */
 	T get(RegisteredBean registeredBean) throws Exception;
 
+	/**
+	 * Return the factory method that this supplier uses to create the
+	 * instance, or {@code null} if it is not known or this supplier uses
+	 * another mean.
+	 * @return the factory method used to create the instance, or {@code null}
+	 */
+	@Nullable
+	default Method getFactoryMethod() {
+		return null;
+	}
+
 	/**
 	 * Return a composed instance supplier that first obtains the instance from
 	 * this supplier, and then applied the {@code after} function to obtain the

spring-beans/src/main/java/org/springframework/beans/factory/support/RootBeanDefinition.java

@@ -429,6 +429,16 @@ public class RootBeanDefinition extends AbstractBeanDefinition {
 		return this.factoryMethodToIntrospect;
 	}
 
+	@Override
+	public void setInstanceSupplier(@Nullable Supplier<?> instanceSupplier) {
+		super.setInstanceSupplier(instanceSupplier);
+		Method factoryMethod = (instanceSupplier instanceof InstanceSupplier<?>)
+				? ((InstanceSupplier<?>) instanceSupplier).getFactoryMethod() : null;
+		if (factoryMethod != null) {
+			setResolvedFactoryMethod(factoryMethod);
+		}
+	}
+
 	/**
 	 * Register an externally managed configuration method or field.
 	 */