Use resolved factory method return type for supplier code generation

Closes gh-29598

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

@@ -46,16 +46,16 @@ import org.springframework.util.function.ThrowingSupplier;
  * 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>
- * Generated code is usually a method reference that generate the
- * {@link BeanInstanceSupplier}, but some shortcut can be used as well such
- * as:
+ *
+ * <p>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
+ * @author Juergen Hoeller
  * @since 6.0
  */
 class InstanceSupplierCodeGenerator {
@@ -90,9 +90,7 @@ class InstanceSupplierCodeGenerator {
 	}
 
 
-	CodeBlock generateCode(RegisteredBean registeredBean,
-			Executable constructorOrFactoryMethod) {
-
+	CodeBlock generateCode(RegisteredBean registeredBean, Executable constructorOrFactoryMethod) {
 		if (constructorOrFactoryMethod instanceof Constructor<?> constructor) {
 			return generateCodeForConstructor(registeredBean, constructor);
 		}
@@ -108,6 +106,7 @@ class InstanceSupplierCodeGenerator {
 		Class<?> beanClass = registeredBean.getBeanClass();
 		Class<?> declaringClass = constructor.getDeclaringClass();
 		boolean dependsOnBean = ClassUtils.isInnerClass(declaringClass);
+
 		Visibility accessVisibility = getAccessVisibility(registeredBean, constructor);
 		if (accessVisibility != Visibility.PRIVATE) {
 			return generateCodeForAccessibleConstructor(beanName, beanClass, constructor,
@@ -121,6 +120,7 @@ class InstanceSupplierCodeGenerator {
 
 		this.generationContext.getRuntimeHints().reflection().registerConstructor(
 				constructor, ExecutableMode.INTROSPECT);
+
 		if (!dependsOnBean && constructor.getParameterCount() == 0) {
 			if (!this.allowDirectSupplierShortcut) {
 				return CodeBlock.of("$T.using($T::new)", InstanceSupplier.class, declaringClass);
@@ -130,6 +130,7 @@ class InstanceSupplierCodeGenerator {
 			}
 			return CodeBlock.of("$T.of($T::new)", ThrowingSupplier.class, declaringClass);
 		}
+
 		GeneratedMethod generatedMethod = generateGetInstanceSupplierMethod(method ->
 				buildGetInstanceMethodForConstructor(method, beanName, beanClass, constructor,
 						declaringClass, dependsOnBean, PRIVATE_STATIC));
@@ -141,6 +142,7 @@ class InstanceSupplierCodeGenerator {
 
 		this.generationContext.getRuntimeHints().reflection()
 				.registerConstructor(constructor, ExecutableMode.INVOKE);
+
 		GeneratedMethod generatedMethod = generateGetInstanceSupplierMethod(method -> {
 			method.addJavadoc("Get the bean instance supplier for '$L'.", beanName);
 			method.addModifiers(PRIVATE_STATIC);
@@ -148,6 +150,7 @@ class InstanceSupplierCodeGenerator {
 			int parameterOffset = (!dependsOnBean) ? 0 : 1;
 			method.addStatement(generateResolverForConstructor(beanClass, constructor, parameterOffset));
 		});
+
 		return generateReturnStatement(generatedMethod);
 	}
 
@@ -158,14 +161,17 @@ class InstanceSupplierCodeGenerator {
 		method.addJavadoc("Get the bean instance supplier for '$L'.", beanName);
 		method.addModifiers(modifiers);
 		method.returns(ParameterizedTypeName.get(BeanInstanceSupplier.class, beanClass));
+
 		int parameterOffset = (!dependsOnBean) ? 0 : 1;
 		CodeBlock.Builder code = CodeBlock.builder();
 		code.add(generateResolverForConstructor(beanClass, 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());
@@ -184,92 +190,99 @@ class InstanceSupplierCodeGenerator {
 		if (!dependsOnBean) {
 			return CodeBlock.of("new $T($L)", declaringClass, args);
 		}
+
 		return CodeBlock.of("$L.getBeanFactory().getBean($T.class).new $L($L)",
 				REGISTERED_BEAN_PARAMETER_NAME, declaringClass.getEnclosingClass(),
 				declaringClass.getSimpleName(), args);
 	}
 
-	private CodeBlock generateCodeForFactoryMethod(RegisteredBean registeredBean,
-			Method factoryMethod) {
-
+	private CodeBlock generateCodeForFactoryMethod(RegisteredBean registeredBean, Method factoryMethod) {
 		String beanName = registeredBean.getBeanName();
-		Class<?> beanClass = registeredBean.getBeanClass();
-		Class<?> declaringClass = ClassUtils
-				.getUserClass(factoryMethod.getDeclaringClass());
+		Class<?> declaringClass = ClassUtils.getUserClass(factoryMethod.getDeclaringClass());
 		boolean dependsOnBean = !Modifier.isStatic(factoryMethod.getModifiers());
+
 		Visibility accessVisibility = getAccessVisibility(registeredBean, factoryMethod);
 		if (accessVisibility != Visibility.PRIVATE) {
 			return generateCodeForAccessibleFactoryMethod(
-					beanName, beanClass, factoryMethod, declaringClass, dependsOnBean);
+					beanName, factoryMethod, declaringClass, dependsOnBean);
 		}
-		return generateCodeForInaccessibleFactoryMethod(beanName, beanClass, factoryMethod, declaringClass);
+		return generateCodeForInaccessibleFactoryMethod(beanName, factoryMethod, declaringClass);
 	}
 
 	private CodeBlock generateCodeForAccessibleFactoryMethod(String beanName,
-			Class<?> beanClass, Method factoryMethod, Class<?> declaringClass, boolean dependsOnBean) {
+			Method factoryMethod, Class<?> declaringClass, boolean dependsOnBean) {
 
 		this.generationContext.getRuntimeHints().reflection().registerMethod(
 				factoryMethod, ExecutableMode.INTROSPECT);
+
 		if (!dependsOnBean && factoryMethod.getParameterCount() == 0) {
+			Class<?> suppliedType = ClassUtils.resolvePrimitiveIfNecessary(factoryMethod.getReturnType());
 			CodeBlock.Builder code = CodeBlock.builder();
 			code.add("$T.<$T>forFactoryMethod($T.class, $S)", BeanInstanceSupplier.class,
-					beanClass, declaringClass, factoryMethod.getName());
+					suppliedType, declaringClass, factoryMethod.getName());
 			code.add(".withGenerator($T::$L)", declaringClass, factoryMethod.getName());
 			return code.build();
 		}
+
 		GeneratedMethod getInstanceMethod = generateGetInstanceSupplierMethod(method ->
-				buildGetInstanceMethodForFactoryMethod(method, beanName, beanClass, factoryMethod,
+				buildGetInstanceMethodForFactoryMethod(method, beanName, factoryMethod,
 						declaringClass, dependsOnBean, PRIVATE_STATIC));
 		return generateReturnStatement(getInstanceMethod);
 	}
 
-	private CodeBlock generateCodeForInaccessibleFactoryMethod(String beanName, Class<?> beanClass,
-			Method factoryMethod, Class<?> declaringClass) {
+	private CodeBlock generateCodeForInaccessibleFactoryMethod(
+			String beanName, Method factoryMethod, Class<?> declaringClass) {
 
 		this.generationContext.getRuntimeHints().reflection().registerMethod(factoryMethod, ExecutableMode.INVOKE);
 		GeneratedMethod getInstanceMethod = generateGetInstanceSupplierMethod(method -> {
+			Class<?> suppliedType = ClassUtils.resolvePrimitiveIfNecessary(factoryMethod.getReturnType());
 			method.addJavadoc("Get the bean instance supplier for '$L'.", beanName);
 			method.addModifiers(PRIVATE_STATIC);
-			method.returns(ParameterizedTypeName.get(BeanInstanceSupplier.class, beanClass));
+			method.returns(ParameterizedTypeName.get(BeanInstanceSupplier.class, suppliedType));
 			method.addStatement(generateInstanceSupplierForFactoryMethod(
-					beanClass, factoryMethod, declaringClass, factoryMethod.getName()));
+					factoryMethod, suppliedType, declaringClass, factoryMethod.getName()));
 		});
 		return generateReturnStatement(getInstanceMethod);
 	}
 
 	private void buildGetInstanceMethodForFactoryMethod(MethodSpec.Builder method,
-			String beanName, Class<?> beanClass, Method factoryMethod, Class<?> declaringClass,
+			String beanName, Method factoryMethod, Class<?> declaringClass,
 			boolean dependsOnBean, javax.lang.model.element.Modifier... modifiers) {
 
 		String factoryMethodName = factoryMethod.getName();
+		Class<?> suppliedType = ClassUtils.resolvePrimitiveIfNecessary(factoryMethod.getReturnType());
+
 		method.addJavadoc("Get the bean instance supplier for '$L'.", beanName);
 		method.addModifiers(modifiers);
-		method.returns(ParameterizedTypeName.get(BeanInstanceSupplier.class, beanClass));
+		method.returns(ParameterizedTypeName.get(BeanInstanceSupplier.class, suppliedType));
+
 		CodeBlock.Builder code = CodeBlock.builder();
 		code.add(generateInstanceSupplierForFactoryMethod(
-				beanClass, factoryMethod, declaringClass, factoryMethodName));
+				factoryMethod, suppliedType, 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 generateInstanceSupplierForFactoryMethod(Class<?> beanClass,
-			Method factoryMethod, Class<?> declaringClass, String factoryMethodName) {
+	private CodeBlock generateInstanceSupplierForFactoryMethod(Method factoryMethod,
+			Class<?> suppliedType, Class<?> declaringClass, String factoryMethodName) {
 
 		if (factoryMethod.getParameterCount() == 0) {
 			return CodeBlock.of("return $T.<$T>forFactoryMethod($T.class, $S)",
-					BeanInstanceSupplier.class, beanClass, declaringClass,
-					factoryMethodName);
+					BeanInstanceSupplier.class, suppliedType, declaringClass, factoryMethodName);
 		}
+
 		CodeBlock parameterTypes = generateParameterTypesCode(factoryMethod.getParameterTypes(), 0);
 		return CodeBlock.of("return $T.<$T>forFactoryMethod($T.class, $S, $L)",
-				BeanInstanceSupplier.class, beanClass, declaringClass, factoryMethodName, parameterTypes);
+				BeanInstanceSupplier.class, suppliedType, declaringClass, factoryMethodName, parameterTypes);
 	}
 
 	private CodeBlock generateNewInstanceCodeForMethod(boolean dependsOnBean,