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

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

@@ -56,8 +56,6 @@ class BeanDefinitionMethodGenerator {
 
 	private final List<BeanRegistrationAotContribution> aotContributions;
 
-	private final List<BeanRegistrationCodeFragmentsCustomizer> codeFragmentsCustomizers;
-
 
 	/**
 	 * Create a new {@link BeanDefinitionMethodGenerator} instance.
@@ -65,13 +63,11 @@ class BeanDefinitionMethodGenerator {
 	 * @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;
@@ -79,7 +75,6 @@ class BeanDefinitionMethodGenerator {
 				.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 @@ class BeanDefinitionMethodGenerator {
 				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 @@ class BeanDefinitionMethodGenerator {
 		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 @@ class BeanDefinitionMethodGenerator {
 		});
 	}
 
-	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 @@ class BeanDefinitionMethodGeneratorFactory {
 		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 @@ import org.springframework.aot.generate.GenerationContext;
 @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

@@ -1,45 +0,0 @@
-/*
- * Copyright 2002-2022 the original author or authors.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      https://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.springframework.beans.factory.aot;
-
-import org.springframework.beans.factory.support.RegisteredBean;
-
-/**
- * Strategy factory interface that can be used to customize the
- * {@link BeanRegistrationCodeFragments} that us used for a given
- * {@link RegisteredBean}. This interface can be used if default code generation
- * isn't suitable for specific types of {@link RegisteredBean}.
- *
- * @author Phillip Webb
- * @since 6.0
- */
-@FunctionalInterface
-public interface BeanRegistrationCodeFragmentsCustomizer {
-
-	/**
-	 * Apply this {@link BeanRegistrationCodeFragmentsCustomizer} to the given
-	 * {@link BeanRegistrationCodeFragments code fragments generator}. The
-	 * returned code generator my be a
-	 * {@link BeanRegistrationCodeFragmentsWrapper wrapper} around the original.
-	 * @param registeredBean the registered bean
-	 * @param codeFragments the existing code fragments
-	 * @return the code generator to use, either the original or a wrapped one;
-	 */
-	BeanRegistrationCodeFragments customizeBeanRegistrationCodeFragments(
-			RegisteredBean registeredBean, BeanRegistrationCodeFragments codeFragments);
-
-}

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

@@ -89,8 +89,7 @@ class BeanRegistrationCodeGenerator implements BeanRegistrationCode {
 		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();