Allow AOT contributions to customize code fragments.

Update the `BeanRegistrationAotContribution` interface to allow
it to customize `BeanRegistrationCodeFragments`. This change
allows us to drop the `BeanRegistrationCodeFragmentsCustomizer`
interface since an `BeanRegistrationAotProcessor` can now be
used instead.

Closes gh-28557

Author: philwebb

spring-aop/src/main/java/org/springframework/aop/scope/ScopedProxyBeanRegistrationCodeFragmentsCustomizer.java renamed to spring-aop/src/main/java/org/springframework/aop/scope/ScopedProxyBeanRegistrationAotProcessor.java

@@ -26,9 +26,10 @@
 
 import org.springframework.aot.generate.GeneratedMethod;
 import org.springframework.aot.generate.GenerationContext;
+import org.springframework.beans.factory.aot.BeanRegistrationAotContribution;
+import org.springframework.beans.factory.aot.BeanRegistrationAotProcessor;
 import org.springframework.beans.factory.aot.BeanRegistrationCode;
 import org.springframework.beans.factory.aot.BeanRegistrationCodeFragments;
-import org.springframework.beans.factory.aot.BeanRegistrationCodeFragmentsCustomizer;
 import org.springframework.beans.factory.config.BeanDefinition;
 import org.springframework.beans.factory.config.ConfigurableBeanFactory;
 import org.springframework.beans.factory.support.InstanceSupplier;
@@ -39,39 +40,37 @@
 import org.springframework.lang.Nullable;
 
 /**
- * {@link BeanRegistrationCodeFragmentsCustomizer} for
- * {@link ScopedProxyFactoryBean}.
+ * {@link BeanRegistrationAotProcessor} for {@link ScopedProxyFactoryBean}.
  *
  * @author Stephane Nicoll
  * @author Phillip Webb
  * @since 6.0
  */
-class ScopedProxyBeanRegistrationCodeFragmentsCustomizer
-		implements BeanRegistrationCodeFragmentsCustomizer {
+class ScopedProxyBeanRegistrationAotProcessor
+		implements BeanRegistrationAotProcessor {
 
 	private static final Log logger = LogFactory
-			.getLog(ScopedProxyBeanRegistrationCodeFragmentsCustomizer.class);
+			.getLog(ScopedProxyBeanRegistrationAotProcessor.class);
 
 
 	@Override
-	public BeanRegistrationCodeFragments customizeBeanRegistrationCodeFragments(
-			RegisteredBean registeredBean, BeanRegistrationCodeFragments codeFragments) {
-
+	public BeanRegistrationAotContribution processAheadOfTime(RegisteredBean registeredBean) {
 		Class<?> beanType = registeredBean.getBeanType().toClass();
-		if (!beanType.equals(ScopedProxyFactoryBean.class)) {
-			return codeFragments;
-		}
-		String targetBeanName = getTargetBeanName(
-				registeredBean.getMergedBeanDefinition());
-		BeanDefinition targetBeanDefinition = getTargetBeanDefinition(
-				registeredBean.getBeanFactory(), targetBeanName);
-		if (targetBeanDefinition == null) {
-			logger.warn("Could not handle " + ScopedProxyFactoryBean.class.getSimpleName()
-					+ ": no target bean definition found with name " + targetBeanName);
-			return codeFragments;
+		if (beanType.equals(ScopedProxyFactoryBean.class)) {
+			String targetBeanName = getTargetBeanName(
+					registeredBean.getMergedBeanDefinition());
+			BeanDefinition targetBeanDefinition = getTargetBeanDefinition(
+					registeredBean.getBeanFactory(), targetBeanName);
+			if (targetBeanDefinition == null) {
+				logger.warn("Could not handle " + ScopedProxyFactoryBean.class.getSimpleName()
+						+ ": no target bean definition found with name " + targetBeanName);
+				return null;
+			}
+			return BeanRegistrationAotContribution.ofBeanRegistrationCodeFragmentsCustomizer(codeFragments ->
+				new ScopedProxyBeanRegistrationCodeFragments(codeFragments, registeredBean,
+						targetBeanName, targetBeanDefinition));
 		}
-		return new ScopedProxyBeanRegistrationCodeFragments(codeFragments, registeredBean,
-				targetBeanName, targetBeanDefinition);
+		return null;
 	}
 
 	@Nullable

spring-aop/src/main/resources/META-INF/spring/aot.factories

@@ -1,2 +1,2 @@
-org.springframework.beans.factory.aot.registration.BeanRegistrationCodeFragmentsCustomizer=\
-org.springframework.aop.scope.ScopedProxyBeanRegistrationCodeFragmentsCustomizer
+org.springframework.beans.factory.aot.registration.BeanRegistrationAotProcessor=\
+org.springframework.aop.scope.ScopedProxyBeanRegistrationAotProcessor

spring-aop/src/test/java/org/springframework/aop/scope/ScopedProxyBeanRegistrationCodeFragmentsCustomizerTests.java renamed to spring-aop/src/test/java/org/springframework/aop/scope/ScopedProxyBeanRegistrationAotProcessorTests.java

@@ -50,13 +50,13 @@
 import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
 
 /**
- * Tests for {@link ScopedProxyBeanRegistrationCodeFragmentsCustomizer}.
+ * Tests for {@link ScopedProxyBeanRegistrationAotProcessor}.
  *
  * @author Stephane Nicoll
  * @author Phillip Webb
  * @since 6.0
  */
-class ScopedProxyBeanRegistrationCodeFragmentsCustomizerTests {
+class ScopedProxyBeanRegistrationAotProcessorTests {
 
 	private DefaultListableBeanFactory beanFactory;
 

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

@@ -56,30 +56,25 @@ class BeanDefinitionMethodGenerator {
 
 	private final List<BeanRegistrationAotContribution> aotContributions;
 
-	private final List<BeanRegistrationCodeFragmentsCustomizer> codeFragmentsCustomizers;
-
 
 	/**
 	 * Create a new {@link BeanDefinitionMethodGenerator} instance.
 	 * @param methodGeneratorFactory the method generator factory
 	 * @param registeredBean the registered bean
 	 * @param innerBeanPropertyName the inner bean property name
 	 * @param aotContributions the AOT contributions
-	 * @param codeFragmentsCustomizers the code fragments customizers
 	 */
 	BeanDefinitionMethodGenerator(
 			BeanDefinitionMethodGeneratorFactory methodGeneratorFactory,
 			RegisteredBean registeredBean, @Nullable String innerBeanPropertyName,
-			List<BeanRegistrationAotContribution> aotContributions,
-			List<BeanRegistrationCodeFragmentsCustomizer> codeFragmentsCustomizers) {
+			List<BeanRegistrationAotContribution> aotContributions) {
 
 		this.methodGeneratorFactory = methodGeneratorFactory;
 		this.registeredBean = registeredBean;
 		this.constructorOrFactoryMethod = ConstructorOrFactoryMethodResolver
 				.resolve(registeredBean);
 		this.innerBeanPropertyName = innerBeanPropertyName;
 		this.aotContributions = aotContributions;
-		this.codeFragmentsCustomizers = codeFragmentsCustomizers;
 	}
 
 	/**
@@ -92,8 +87,7 @@ class BeanDefinitionMethodGenerator {
 	MethodReference generateBeanDefinitionMethod(GenerationContext generationContext,
 			BeanRegistrationsCode beanRegistrationsCode) {
 
-		BeanRegistrationCodeFragments codeFragments = getCodeFragments(
-				beanRegistrationsCode);
+		BeanRegistrationCodeFragments codeFragments = getCodeFragments(beanRegistrationsCode);
 		Class<?> target = codeFragments.getTarget(this.registeredBean,
 				this.constructorOrFactoryMethod);
 		if (!target.getName().startsWith("java.")) {
@@ -115,7 +109,17 @@ MethodReference generateBeanDefinitionMethod(GenerationContext generationContext
 				Modifier.PRIVATE);
 		return MethodReference.ofStatic(beanRegistrationsCode.getClassName(),
 				generatedMethod.getName().toString());
+	}
+
+	private BeanRegistrationCodeFragments getCodeFragments(
+			BeanRegistrationsCode beanRegistrationsCode) {
 
+		BeanRegistrationCodeFragments codeFragments = new DefaultBeanRegistrationCodeFragments(
+				beanRegistrationsCode, this.registeredBean, this.methodGeneratorFactory);
+		for (BeanRegistrationAotContribution aotContribution : this.aotContributions) {
+			codeFragments = aotContribution.customizeBeanRegistrationCodeFragments(codeFragments);
+		}
+		return codeFragments;
 	}
 
 	private GeneratedMethod generateBeanDefinitionMethod(
@@ -126,8 +130,7 @@ private GeneratedMethod generateBeanDefinitionMethod(
 		BeanRegistrationCodeGenerator codeGenerator = new BeanRegistrationCodeGenerator(
 				className, methodGenerator, this.registeredBean,
 				this.constructorOrFactoryMethod, codeFragments);
-		GeneratedMethod method = methodGenerator.generateMethod("get", "bean",
-				"definition");
+		GeneratedMethod method = methodGenerator.generateMethod("get", "bean", "definition");
 		this.aotContributions.forEach(aotContribution -> aotContribution
 				.applyTo(generationContext, codeGenerator));
 		return method.using(builder -> {
@@ -140,18 +143,6 @@ private GeneratedMethod generateBeanDefinitionMethod(
 		});
 	}
 
-	private BeanRegistrationCodeFragments getCodeFragments(
-			BeanRegistrationsCode beanRegistrationsCode) {
-
-		BeanRegistrationCodeFragments codeFragments = new DefaultBeanRegistrationCodeFragments(
-				beanRegistrationsCode, this.registeredBean, this.methodGeneratorFactory);
-		for (BeanRegistrationCodeFragmentsCustomizer customizer : this.codeFragmentsCustomizers) {
-			codeFragments = customizer.customizeBeanRegistrationCodeFragments(
-					this.registeredBean, codeFragments);
-		}
-		return codeFragments;
-	}
-
 	private String getName() {
 		if (this.innerBeanPropertyName != null) {
 			return this.innerBeanPropertyName;

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

@@ -47,8 +47,6 @@ class BeanDefinitionMethodGeneratorFactory {
 
 	private final List<BeanRegistrationExcludeFilter> excludeFilters;
 
-	private final List<BeanRegistrationCodeFragmentsCustomizer> codeGenerationCustomizers;
-
 
 	/**
 	 * Create a new {@link BeanDefinitionMethodGeneratorFactory} backed by the
@@ -67,8 +65,6 @@ class BeanDefinitionMethodGeneratorFactory {
 	BeanDefinitionMethodGeneratorFactory(AotFactoriesLoader loader) {
 		this.aotProcessors = loader.load(BeanRegistrationAotProcessor.class);
 		this.excludeFilters = loader.load(BeanRegistrationExcludeFilter.class);
-		this.codeGenerationCustomizers = loader
-				.load(BeanRegistrationCodeFragmentsCustomizer.class);
 	}
 
 
@@ -92,7 +88,7 @@ BeanDefinitionMethodGenerator getBeanDefinitionMethodGenerator(
 		List<BeanRegistrationAotContribution> contributions = getAotContributions(
 				registeredBean);
 		return new BeanDefinitionMethodGenerator(this, registeredBean,
-				innerBeanPropertyName, contributions, this.codeGenerationCustomizers);
+				innerBeanPropertyName, contributions);
 	}
 
 	private boolean isExcluded(RegisteredBean registeredBean) {

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

@@ -16,7 +16,10 @@
 
 package org.springframework.beans.factory.aot;
 
+import java.util.function.UnaryOperator;
+
 import org.springframework.aot.generate.GenerationContext;
+import org.springframework.util.Assert;
 
 /**
  * AOT contribution from a {@link BeanRegistrationAotProcessor} used to register
@@ -29,6 +32,19 @@
 @FunctionalInterface
 public interface BeanRegistrationAotContribution {
 
+	/**
+	 * Customize the {@link BeanRegistrationCodeFragments} that will be used to
+	 * generate the bean registration code. Custom code fragments can be used if
+	 * default code generation isn't suitable.
+	 * @param codeFragments the existing code fragments
+	 * @return the code fragments to use, may be the original instance or a
+	 * wrapper
+	 */
+	default BeanRegistrationCodeFragments customizeBeanRegistrationCodeFragments(
+			BeanRegistrationCodeFragments codeFragments) {
+		return codeFragments;
+	}
+
 	/**
 	 * Apply this contribution to the given {@link BeanRegistrationCode}.
 	 * @param generationContext the active generation context
@@ -37,4 +53,33 @@ public interface BeanRegistrationAotContribution {
 	void applyTo(GenerationContext generationContext,
 			BeanRegistrationCode beanRegistrationCode);
 
+	/**
+	 * Factory method that can be used to create a
+	 * {@link BeanRegistrationAotContribution} that applies the given
+	 * {@link BeanRegistrationCodeFragments} customizer.
+	 * @param beanRegistrationCodeFragmentsCustomizer the
+	 * {@link BeanRegistrationCodeFragments} customizer
+	 * @return a new {@link BeanRegistrationAotContribution} instance
+	 * @see #customizeBeanRegistrationCodeFragments(BeanRegistrationCodeFragments)
+	 */
+	static BeanRegistrationAotContribution ofBeanRegistrationCodeFragmentsCustomizer(
+			UnaryOperator<BeanRegistrationCodeFragments> beanRegistrationCodeFragmentsCustomizer) {
+		Assert.notNull(beanRegistrationCodeFragmentsCustomizer,
+				"BeanRegistrationCodeFragmentsCustomizer must not be null");
+		return new BeanRegistrationAotContribution() {
+
+			@Override
+			public BeanRegistrationCodeFragments customizeBeanRegistrationCodeFragments(
+					BeanRegistrationCodeFragments codeFragments) {
+				return beanRegistrationCodeFragmentsCustomizer.apply(codeFragments);
+			}
+
+			@Override
+			public void applyTo(GenerationContext generationContext,
+					BeanRegistrationCode beanRegistrationCode) {
+			}
+
+		};
+	}
+
 }

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

@@ -89,8 +89,7 @@ CodeBlock generateCode(GenerationContext generationContext) {
 		CodeBlock instanceSupplierCode = this.codeFragments.generateInstanceSupplierCode(
 				generationContext, this, this.constructorOrFactoryMethod,
 				this.instancePostProcessors.isEmpty());
-		builder.add(
-				this.codeFragments.generateSetBeanInstanceSupplierCode(generationContext,
+		builder.add(this.codeFragments.generateSetBeanInstanceSupplierCode(generationContext,
 						this, instanceSupplierCode, this.instancePostProcessors));
 		builder.add(this.codeFragments.generateReturnCode(generationContext, this));
 		return builder.build();

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

@@ -182,11 +182,6 @@ private static class MockBeanRegistrationCode implements BeanRegistrationCode {
 
 		private final List<MethodReference> instancePostProcessors = new ArrayList<>();
 
-		@Override
-		public void addInstancePostProcessor(MethodReference methodReference) {
-			this.instancePostProcessors.add(methodReference);
-		}
-
 		@Override
 		public ClassName getClassName() {
 			return null;
@@ -197,6 +192,11 @@ public MethodGenerator getMethodGenerator() {
 			return null;
 		}
 
+		@Override
+		public void addInstancePostProcessor(MethodReference methodReference) {
+			this.instancePostProcessors.add(methodReference);
+		}
+
 	}
 
 }