Revisit BeanRegistrationCodeFragments

This commit revisit BeanRegistrationCodeFragments to separate the
responsibility between the default implementation and the delegates. It
also reviews how customization are applied by improving the Javadoc and
the method name.

Closes gh-28865

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

@@ -30,6 +30,7 @@ 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.BeanRegistrationCodeFragmentsDecorator;
 import org.springframework.beans.factory.config.BeanDefinition;
 import org.springframework.beans.factory.config.ConfigurableBeanFactory;
 import org.springframework.beans.factory.support.InstanceSupplier;
@@ -63,9 +64,9 @@ class ScopedProxyBeanRegistrationAotProcessor implements BeanRegistrationAotProc
 						": no target bean definition found with name " + targetBeanName);
 				return null;
 			}
-			return BeanRegistrationAotContribution.ofBeanRegistrationCodeFragmentsCustomizer(codeFragments ->
-				new ScopedProxyBeanRegistrationCodeFragments(codeFragments, registeredBean,
-						targetBeanName, targetBeanDefinition));
+			return BeanRegistrationAotContribution.withCustomCodeFragments(codeFragments ->
+					new ScopedProxyBeanRegistrationCodeFragments(codeFragments, registeredBean,
+							targetBeanName, targetBeanDefinition));
 		}
 		return null;
 	}
@@ -87,7 +88,7 @@ class ScopedProxyBeanRegistrationAotProcessor implements BeanRegistrationAotProc
 	}
 
 
-	private static class ScopedProxyBeanRegistrationCodeFragments extends BeanRegistrationCodeFragments {
+	private static class ScopedProxyBeanRegistrationCodeFragments extends BeanRegistrationCodeFragmentsDecorator {
 
 		private static final String REGISTERED_BEAN_PARAMETER_NAME = "registeredBean";
 
@@ -97,10 +98,10 @@ class ScopedProxyBeanRegistrationAotProcessor implements BeanRegistrationAotProc
 
 		private final BeanDefinition targetBeanDefinition;
 
-		ScopedProxyBeanRegistrationCodeFragments(BeanRegistrationCodeFragments codeGenerator,
+		ScopedProxyBeanRegistrationCodeFragments(BeanRegistrationCodeFragments delegate,
 				RegisteredBean registeredBean, String targetBeanName, BeanDefinition targetBeanDefinition) {
 
-			super(codeGenerator);
+			super(delegate);
 			this.registeredBean = registeredBean;
 			this.targetBeanName = targetBeanName;
 			this.targetBeanDefinition = targetBeanDefinition;

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

@@ -26,6 +26,7 @@ import org.springframework.util.Assert;
  * a single bean definition.
  *
  * @author Phillip Webb
+ * @author Stephane Nicoll
  * @since 6.0
  * @see BeanRegistrationAotProcessor
  */
@@ -55,24 +56,23 @@ public interface BeanRegistrationAotContribution {
 			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
+	 * Create a {@link BeanRegistrationAotContribution} that customizes
+	 * the {@link BeanRegistrationCodeFragments}. Typically used in
+	 * conjunction with an extension of {@link BeanRegistrationCodeFragmentsDecorator}
+	 * that overrides a specific callback.
+	 * @param defaultCodeFragments the default code fragments
 	 * @return a new {@link BeanRegistrationAotContribution} instance
-	 * @see #customizeBeanRegistrationCodeFragments(GenerationContext, BeanRegistrationCodeFragments)
+	 * @see BeanRegistrationCodeFragmentsDecorator
 	 */
-	static BeanRegistrationAotContribution ofBeanRegistrationCodeFragmentsCustomizer(
-			UnaryOperator<BeanRegistrationCodeFragments> beanRegistrationCodeFragmentsCustomizer) {
-		Assert.notNull(beanRegistrationCodeFragmentsCustomizer,
-				"BeanRegistrationCodeFragmentsCustomizer must not be null");
+	static BeanRegistrationAotContribution withCustomCodeFragments(
+			UnaryOperator<BeanRegistrationCodeFragments> defaultCodeFragments) {
+		Assert.notNull(defaultCodeFragments, "'defaultCodeFragments' must not be null");
 		return new BeanRegistrationAotContribution() {
 
 			@Override
 			public BeanRegistrationCodeFragments customizeBeanRegistrationCodeFragments(
 					GenerationContext generationContext, BeanRegistrationCodeFragments codeFragments) {
-				return beanRegistrationCodeFragmentsCustomizer.apply(codeFragments);
+				return defaultCodeFragments.apply(codeFragments);
 			}
 
 			@Override

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

@@ -27,46 +27,26 @@ import org.springframework.beans.factory.support.RegisteredBean;
 import org.springframework.beans.factory.support.RootBeanDefinition;
 import org.springframework.core.ResolvableType;
 import org.springframework.javapoet.CodeBlock;
-import org.springframework.util.Assert;
 
 /**
- * Class used to generate the various fragments of code needed to register a
- * bean.
+ * Generate the various fragments of code needed to register a bean.
  *
  * @author Phillip Webb
  * @since 6.0
  */
-public abstract class BeanRegistrationCodeFragments {
+public interface BeanRegistrationCodeFragments {
 
 	/**
 	 * The variable name to used when creating the bean definition.
 	 */
-	protected static final String BEAN_DEFINITION_VARIABLE = "beanDefinition";
+	String BEAN_DEFINITION_VARIABLE = "beanDefinition";
 
 	/**
 	 * The variable name to used when creating the bean definition.
 	 */
-	protected static final String INSTANCE_SUPPLIER_VARIABLE = "instanceSupplier";
-
-
-	private final BeanRegistrationCodeFragments codeFragments;
-
-
-	protected BeanRegistrationCodeFragments(BeanRegistrationCodeFragments codeFragments) {
-		Assert.notNull(codeFragments, "'codeFragments' must not be null");
-		this.codeFragments = codeFragments;
-	}
+	String INSTANCE_SUPPLIER_VARIABLE = "instanceSupplier";
 
 
-	/**
-	 * Package-private constructor exclusively for
-	 * {@link DefaultBeanRegistrationCodeFragments}. All methods are overridden
-	 * so {@code this.codeFragments} is never actually used.
-	 */
-	BeanRegistrationCodeFragments() {
-		this.codeFragments = this;
-	}
-
 	/**
 	 * Return the target for the registration. Used to determine where to write
 	 * the code.
@@ -74,11 +54,8 @@ public abstract class BeanRegistrationCodeFragments {
 	 * @param constructorOrFactoryMethod the constructor or factory method
 	 * @return the target class
 	 */
-	public Class<?> getTarget(RegisteredBean registeredBean,
-			Executable constructorOrFactoryMethod) {
-
-		return this.codeFragments.getTarget(registeredBean, constructorOrFactoryMethod);
-	}
+	Class<?> getTarget(RegisteredBean registeredBean,
+			Executable constructorOrFactoryMethod);
 
 	/**
 	 * Generate the code that defines the new bean definition instance.
@@ -87,13 +64,8 @@ public abstract class BeanRegistrationCodeFragments {
 	 * @param beanRegistrationCode the bean registration code
 	 * @return the generated code
 	 */
-	public CodeBlock generateNewBeanDefinitionCode(GenerationContext generationContext,
-			ResolvableType beanType, BeanRegistrationCode beanRegistrationCode) {
-
-		return this.codeFragments.generateNewBeanDefinitionCode(generationContext,
-				beanType, beanRegistrationCode);
-
-	}
+	CodeBlock generateNewBeanDefinitionCode(GenerationContext generationContext,
+			ResolvableType beanType, BeanRegistrationCode beanRegistrationCode);
 
 	/**
 	 * Generate the code that sets the properties of the bean definition.
@@ -102,14 +74,9 @@ public abstract class BeanRegistrationCodeFragments {
 	 * @param attributeFilter any attribute filtering that should be applied
 	 * @return the generated code
 	 */
-	public CodeBlock generateSetBeanDefinitionPropertiesCode(
+	CodeBlock generateSetBeanDefinitionPropertiesCode(
 			GenerationContext generationContext, BeanRegistrationCode beanRegistrationCode,
-			RootBeanDefinition beanDefinition, Predicate<String> attributeFilter) {
-
-		return this.codeFragments.generateSetBeanDefinitionPropertiesCode(
-				generationContext, beanRegistrationCode, beanDefinition, attributeFilter);
-
-	}
+			RootBeanDefinition beanDefinition, Predicate<String> attributeFilter);
 
 	/**
 	 * Generate the code that sets the instance supplier on the bean definition.
@@ -120,13 +87,9 @@ public abstract class BeanRegistrationCodeFragments {
 	 * @return the generated code
 	 * @see #generateInstanceSupplierCode
 	 */
-	public CodeBlock generateSetBeanInstanceSupplierCode(
+	CodeBlock generateSetBeanInstanceSupplierCode(
 			GenerationContext generationContext, BeanRegistrationCode beanRegistrationCode,
-			CodeBlock instanceSupplierCode, List<MethodReference> postProcessors) {
-
-		return this.codeFragments.generateSetBeanInstanceSupplierCode(generationContext,
-				beanRegistrationCode, instanceSupplierCode, postProcessors);
-	}
+			CodeBlock instanceSupplierCode, List<MethodReference> postProcessors);
 
 	/**
 	 * Generate the instance supplier code.
@@ -138,13 +101,9 @@ public abstract class BeanRegistrationCodeFragments {
 	 * than always needing an {@link InstanceSupplier}
 	 * @return the generated code
 	 */
-	public CodeBlock generateInstanceSupplierCode(
+	CodeBlock generateInstanceSupplierCode(
 			GenerationContext generationContext, BeanRegistrationCode beanRegistrationCode,
-			Executable constructorOrFactoryMethod, boolean allowDirectSupplierShortcut) {
-
-		return this.codeFragments.generateInstanceSupplierCode(generationContext,
-				beanRegistrationCode, constructorOrFactoryMethod, allowDirectSupplierShortcut);
-	}
+			Executable constructorOrFactoryMethod, boolean allowDirectSupplierShortcut);
 
 	/**
 	 * Generate the return statement.
@@ -152,10 +111,7 @@ public abstract class BeanRegistrationCodeFragments {
 	 * @param beanRegistrationCode the bean registration code
 	 * @return the generated code
 	 */
-	public CodeBlock generateReturnCode(
-			GenerationContext generationContext, BeanRegistrationCode beanRegistrationCode) {
-
-		return this.codeFragments.generateReturnCode(generationContext, beanRegistrationCode);
-	}
+	CodeBlock generateReturnCode(
+			GenerationContext generationContext, BeanRegistrationCode beanRegistrationCode);
 
 }

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

@@ -0,0 +1,108 @@
+/*
+ * 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 java.lang.reflect.Executable;
+import java.util.List;
+import java.util.function.Predicate;
+import java.util.function.UnaryOperator;
+
+import org.springframework.aot.generate.GenerationContext;
+import org.springframework.aot.generate.MethodReference;
+import org.springframework.beans.factory.support.RegisteredBean;
+import org.springframework.beans.factory.support.RootBeanDefinition;
+import org.springframework.core.ResolvableType;
+import org.springframework.javapoet.CodeBlock;
+import org.springframework.util.Assert;
+
+/**
+ * A {@link BeanRegistrationCodeFragments} decorator implementation. Typically
+ * used when part of the default code fragments have to customized, by extending
+ * this class and use it as part of
+ * {@link BeanRegistrationAotContribution#withCustomCodeFragments(UnaryOperator)}.
+ *
+ * @author Phillip Webb
+ * @author Stephane Nicoll
+ * @since 6.0
+ */
+public class BeanRegistrationCodeFragmentsDecorator implements BeanRegistrationCodeFragments {
+
+
+	private final BeanRegistrationCodeFragments delegate;
+
+
+	protected BeanRegistrationCodeFragmentsDecorator(BeanRegistrationCodeFragments delegate) {
+		Assert.notNull(delegate, "Delegate must not be null");
+		this.delegate = delegate;
+	}
+
+	@Override
+	public Class<?> getTarget(RegisteredBean registeredBean,
+			Executable constructorOrFactoryMethod) {
+
+		return this.delegate.getTarget(registeredBean, constructorOrFactoryMethod);
+	}
+
+	@Override
+	public CodeBlock generateNewBeanDefinitionCode(GenerationContext generationContext,
+			ResolvableType beanType, BeanRegistrationCode beanRegistrationCode) {
+
+		return this.delegate.generateNewBeanDefinitionCode(generationContext,
+				beanType, beanRegistrationCode);
+
+	}
+
+	@Override
+	public CodeBlock generateSetBeanDefinitionPropertiesCode(
+			GenerationContext generationContext,
+			BeanRegistrationCode beanRegistrationCode, RootBeanDefinition beanDefinition,
+			Predicate<String> attributeFilter) {
+
+		return this.delegate.generateSetBeanDefinitionPropertiesCode(
+				generationContext, beanRegistrationCode, beanDefinition, attributeFilter);
+
+	}
+
+	@Override
+	public CodeBlock generateSetBeanInstanceSupplierCode(
+			GenerationContext generationContext,
+			BeanRegistrationCode beanRegistrationCode, CodeBlock instanceSupplierCode,
+			List<MethodReference> postProcessors) {
+
+		return this.delegate.generateSetBeanInstanceSupplierCode(generationContext,
+				beanRegistrationCode, instanceSupplierCode, postProcessors);
+	}
+
+	@Override
+	public CodeBlock generateInstanceSupplierCode(GenerationContext generationContext,
+			BeanRegistrationCode beanRegistrationCode,
+			Executable constructorOrFactoryMethod, boolean allowDirectSupplierShortcut) {
+
+		return this.delegate.generateInstanceSupplierCode(generationContext,
+				beanRegistrationCode, constructorOrFactoryMethod,
+				allowDirectSupplierShortcut);
+	}
+
+	@Override
+	public CodeBlock generateReturnCode(GenerationContext generationContext,
+			BeanRegistrationCode beanRegistrationCode) {
+
+		return this.delegate.generateReturnCode(generationContext,
+				beanRegistrationCode);
+	}
+
+}

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

@@ -44,7 +44,7 @@ import org.springframework.util.ClassUtils;
  *
  * @author Phillip Webb
  */
-class DefaultBeanRegistrationCodeFragments extends BeanRegistrationCodeFragments {
+class DefaultBeanRegistrationCodeFragments implements BeanRegistrationCodeFragments {
 
 	/**
 	 * The variable name used to hold the bean type.

spring-beans/src/test/java/org/springframework/beans/factory/aot/BeanDefinitionMethodGeneratorTests.java

@@ -196,7 +196,7 @@ class BeanDefinitionMethodGeneratorTests {
 		RegisteredBean registeredBean = registerBean(
 				new RootBeanDefinition(TestBean.class));
 		BeanRegistrationAotContribution aotContribution = BeanRegistrationAotContribution
-				.ofBeanRegistrationCodeFragmentsCustomizer(this::customizeBeanDefinitionCode);
+				.withCustomCodeFragments(this::customizeBeanDefinitionCode);
 		List<BeanRegistrationAotContribution> aotContributions = Collections.singletonList(aotContribution);
 		BeanDefinitionMethodGenerator generator = new BeanDefinitionMethodGenerator(
 				this.methodGeneratorFactory, registeredBean, null, aotContributions);
@@ -211,7 +211,7 @@ class BeanDefinitionMethodGeneratorTests {
 
 	private BeanRegistrationCodeFragments customizeBeanDefinitionCode(
 			BeanRegistrationCodeFragments codeFragments) {
-		return new BeanRegistrationCodeFragments(codeFragments) {
+		return new BeanRegistrationCodeFragmentsDecorator(codeFragments) {
 
 			@Override
 			public CodeBlock generateNewBeanDefinitionCode(
@@ -251,7 +251,7 @@ class BeanDefinitionMethodGeneratorTests {
 		beanDefinition.setAttribute("b", "B");
 		RegisteredBean registeredBean = registerBean(beanDefinition);
 		BeanRegistrationAotContribution aotContribution = BeanRegistrationAotContribution
-				.ofBeanRegistrationCodeFragmentsCustomizer(this::customizeAttributeFilter);
+				.withCustomCodeFragments(this::customizeAttributeFilter);
 		List<BeanRegistrationAotContribution> aotContributions = Collections
 				.singletonList(aotContribution);
 		BeanDefinitionMethodGenerator generator = new BeanDefinitionMethodGenerator(
@@ -267,7 +267,7 @@ class BeanDefinitionMethodGeneratorTests {
 
 	private BeanRegistrationCodeFragments customizeAttributeFilter(
 			BeanRegistrationCodeFragments codeFragments) {
-		return new BeanRegistrationCodeFragments(codeFragments) {
+		return new BeanRegistrationCodeFragmentsDecorator(codeFragments) {
 
 			@Override
 			public CodeBlock generateSetBeanDefinitionPropertiesCode(

spring-orm/src/main/java/org/springframework/orm/jpa/persistenceunit/PersistenceManagedTypesBeanRegistrationAotProcessor.java

@@ -34,6 +34,7 @@ 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.BeanRegistrationCodeFragmentsDecorator;
 import org.springframework.beans.factory.support.RegisteredBean;
 import org.springframework.core.annotation.AnnotationUtils;
 import org.springframework.javapoet.CodeBlock;
@@ -58,13 +59,13 @@ class PersistenceManagedTypesBeanRegistrationAotProcessor implements BeanRegistr
 	@Override
 	public BeanRegistrationAotContribution processAheadOfTime(RegisteredBean registeredBean) {
 		if (PersistenceManagedTypes.class.isAssignableFrom(registeredBean.getBeanClass())) {
-			return BeanRegistrationAotContribution.ofBeanRegistrationCodeFragmentsCustomizer(codeFragments ->
+			return BeanRegistrationAotContribution.withCustomCodeFragments(codeFragments ->
 					new JpaManagedTypesBeanRegistrationCodeFragments(codeFragments, registeredBean));
 		}
 		return null;
 	}
 
-	private static class JpaManagedTypesBeanRegistrationCodeFragments extends BeanRegistrationCodeFragments {
+	private static class JpaManagedTypesBeanRegistrationCodeFragments extends BeanRegistrationCodeFragmentsDecorator {
 
 		private static final ParameterizedTypeName LIST_OF_STRINGS_TYPE = ParameterizedTypeName.get(List.class, String.class);