Polish

See gh-29391

Author: snicoll

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

@@ -30,14 +30,14 @@
 import org.springframework.javapoet.ClassName;
 import org.springframework.javapoet.CodeBlock;
 import org.springframework.javapoet.MethodSpec;
-import org.springframework.util.MultiValueMap;
 
 /**
  * AOT contribution from a {@link BeanRegistrationsAotProcessor} used to
  * register bean definitions and aliases.
  *
  * @author Phillip Webb
  * @author Sebastien Deleuze
+ * @author Stephane Nicoll
  * @since 6.0
  * @see BeanRegistrationsAotProcessor
  */
@@ -46,16 +46,10 @@ class BeanRegistrationsAotContribution
 
 	private static final String BEAN_FACTORY_PARAMETER_NAME = "beanFactory";
 
-	private final Map<String, BeanDefinitionMethodGenerator> registrations;
-
-	private final MultiValueMap<String, String> aliases;
-
-
-	BeanRegistrationsAotContribution(
-			Map<String, BeanDefinitionMethodGenerator> registrations, MultiValueMap<String, String> aliases) {
+	private final Map<String, Registration> registrations;
 
+	BeanRegistrationsAotContribution(Map<String, Registration> registrations) {
 		this.registrations = registrations;
-		this.aliases = aliases;
 	}
 
 
@@ -86,8 +80,8 @@ private void generateRegisterBeanDefinitionsMethod(MethodSpec.Builder method,
 		method.addParameter(DefaultListableBeanFactory.class,
 				BEAN_FACTORY_PARAMETER_NAME);
 		CodeBlock.Builder code = CodeBlock.builder();
-		this.registrations.forEach((beanName, beanDefinitionMethodGenerator) -> {
-			MethodReference beanDefinitionMethod = beanDefinitionMethodGenerator
+		this.registrations.forEach((beanName, registration) -> {
+			MethodReference beanDefinitionMethod = registration.methodGenerator
 					.generateBeanDefinitionMethod(generationContext,
 							beanRegistrationsCode);
 			CodeBlock methodInvocation = beanDefinitionMethod.toInvokeCodeBlock(
@@ -105,12 +99,22 @@ private void generateRegisterAliasesMethod(MethodSpec.Builder method) {
 		method.addParameter(DefaultListableBeanFactory.class,
 				BEAN_FACTORY_PARAMETER_NAME);
 		CodeBlock.Builder code = CodeBlock.builder();
-		this.aliases.forEach((beanName, beanAliases) ->
-				beanAliases.forEach(alias -> code.addStatement("$L.registerAlias($S, $S)", BEAN_FACTORY_PARAMETER_NAME,
-						beanName, alias)));
+		this.registrations.forEach((beanName, registration) -> {
+			for (String alias : registration.aliases) {
+				code.addStatement("$L.registerAlias($S, $S)",
+						BEAN_FACTORY_PARAMETER_NAME, beanName, alias);
+			}
+		});
 		method.addCode(code.build());
 	}
 
+	/**
+	 * Gather the necessary information to register a particular bean.
+	 * @param methodGenerator the {@link BeanDefinitionMethodGenerator} to use
+	 * @param aliases the bean aliases, if any
+	 */
+	record Registration(BeanDefinitionMethodGenerator methodGenerator, String[] aliases) {}
+
 
 	/**
 	 * {@link BeanRegistrationsCode} with generation support.

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

@@ -19,17 +19,17 @@
 import java.util.LinkedHashMap;
 import java.util.Map;
 
+import org.springframework.beans.factory.aot.BeanRegistrationsAotContribution.Registration;
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
 import org.springframework.beans.factory.support.RegisteredBean;
-import org.springframework.util.LinkedMultiValueMap;
-import org.springframework.util.MultiValueMap;
 
 /**
  * {@link BeanFactoryInitializationAotProcessor} that contributes code to
  * register beans.
  *
  * @author Phillip Webb
  * @author Sebastien Deleuze
+ * @author Stephane Nicoll
  * @since 6.0
  */
 class BeanRegistrationsAotProcessor implements BeanFactoryInitializationAotProcessor {
@@ -38,23 +38,20 @@ class BeanRegistrationsAotProcessor implements BeanFactoryInitializationAotProce
 	public BeanRegistrationsAotContribution processAheadOfTime(ConfigurableListableBeanFactory beanFactory) {
 		BeanDefinitionMethodGeneratorFactory beanDefinitionMethodGeneratorFactory =
 				new BeanDefinitionMethodGeneratorFactory(beanFactory);
-		Map<String, BeanDefinitionMethodGenerator> registrations = new LinkedHashMap<>();
-		MultiValueMap<String, String> aliases = new LinkedMultiValueMap<>();
+		Map<String, Registration> registrations = new LinkedHashMap<>();
 		for (String beanName : beanFactory.getBeanDefinitionNames()) {
 			RegisteredBean registeredBean = RegisteredBean.of(beanFactory, beanName);
 			BeanDefinitionMethodGenerator beanDefinitionMethodGenerator = beanDefinitionMethodGeneratorFactory
 					.getBeanDefinitionMethodGenerator(registeredBean);
 			if (beanDefinitionMethodGenerator != null) {
-				registrations.put(beanName, beanDefinitionMethodGenerator);
-			}
-			for (String alias : beanFactory.getAliases(beanName)) {
-				aliases.add(beanName, alias);
+				registrations.put(beanName, new Registration(beanDefinitionMethodGenerator,
+						beanFactory.getAliases(beanName)));
 			}
 		}
 		if (registrations.isEmpty()) {
 			return null;
 		}
-		return new BeanRegistrationsAotContribution(registrations, aliases);
+		return new BeanRegistrationsAotContribution(registrations);
 	}
 
 }

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

@@ -18,7 +18,6 @@
 
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.function.BiConsumer;
@@ -46,16 +45,16 @@
 import org.springframework.javapoet.CodeBlock;
 import org.springframework.javapoet.MethodSpec;
 import org.springframework.javapoet.ParameterizedTypeName;
-import org.springframework.util.LinkedMultiValueMap;
-import org.springframework.util.MultiValueMap;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.springframework.beans.factory.aot.BeanRegistrationsAotContribution.Registration;
 
 /**
  * Tests for {@link BeanRegistrationsAotContribution}.
  *
  * @author Phillip Webb
  * @author Sebastien Deleuze
+ * @author Stephane Nicoll
  */
 class BeanRegistrationsAotContributionTests {
 
@@ -80,15 +79,12 @@ class BeanRegistrationsAotContributionTests {
 
 	@Test
 	void applyToAppliesContribution() {
-		Map<String, BeanDefinitionMethodGenerator> registrations = new LinkedHashMap<>();
 		RegisteredBean registeredBean = registerBean(
 				new RootBeanDefinition(TestBean.class));
 		BeanDefinitionMethodGenerator generator = new BeanDefinitionMethodGenerator(
 				this.methodGeneratorFactory, registeredBean, null,
 				Collections.emptyList());
-		registrations.put("testBean", generator);
-		BeanRegistrationsAotContribution contribution = new BeanRegistrationsAotContribution(
-				registrations, new LinkedMultiValueMap<>());
+		BeanRegistrationsAotContribution contribution = createContribution(generator);
 		contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
 		compile((consumer, compiled) -> {
 			DefaultListableBeanFactory freshBeanFactory = new DefaultListableBeanFactory();
@@ -99,17 +95,12 @@ void applyToAppliesContribution() {
 
 	@Test
 	void applyToAppliesContributionWithAliases() {
-		Map<String, BeanDefinitionMethodGenerator> registrations = new LinkedHashMap<>();
 		RegisteredBean registeredBean = registerBean(
 				new RootBeanDefinition(TestBean.class));
 		BeanDefinitionMethodGenerator generator = new BeanDefinitionMethodGenerator(
 				this.methodGeneratorFactory, registeredBean, null,
 				Collections.emptyList());
-		registrations.put("testBean", generator);
-		MultiValueMap<String, String> aliases = new LinkedMultiValueMap<>();
-		aliases.add("testBean", "testAlias");
-		BeanRegistrationsAotContribution contribution = new BeanRegistrationsAotContribution(
-				registrations, aliases);
+		BeanRegistrationsAotContribution contribution = createContribution(generator, "testAlias");
 		contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
 		compile((consumer, compiled) -> {
 			DefaultListableBeanFactory freshBeanFactory = new DefaultListableBeanFactory();
@@ -123,15 +114,12 @@ void applyToWhenHasNameGeneratesPrefixedFeatureName() {
 		this.generationContext = new TestGenerationContext(
 				new ClassNameGenerator(TestGenerationContext.TEST_TARGET, "Management"));
 		this.beanFactoryInitializationCode = new MockBeanFactoryInitializationCode(this.generationContext);
-		Map<String, BeanDefinitionMethodGenerator> registrations = new LinkedHashMap<>();
 		RegisteredBean registeredBean = registerBean(
 				new RootBeanDefinition(TestBean.class));
 		BeanDefinitionMethodGenerator generator = new BeanDefinitionMethodGenerator(
 				this.methodGeneratorFactory, registeredBean, null,
 				Collections.emptyList());
-		registrations.put("testBean", generator);
-		BeanRegistrationsAotContribution contribution = new BeanRegistrationsAotContribution(
-				registrations, new LinkedMultiValueMap<>());
+		BeanRegistrationsAotContribution contribution = createContribution(generator);
 		contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
 		compile((consumer, compiled) -> {
 			SourceFile sourceFile = compiled.getSourceFile(".*BeanDefinitions");
@@ -142,7 +130,6 @@ void applyToWhenHasNameGeneratesPrefixedFeatureName() {
 	@Test
 	void applyToCallsRegistrationsWithBeanRegistrationsCode() {
 		List<BeanRegistrationsCode> beanRegistrationsCodes = new ArrayList<>();
-		Map<String, BeanDefinitionMethodGenerator> registrations = new LinkedHashMap<>();
 		RegisteredBean registeredBean = registerBean(
 				new RootBeanDefinition(TestBean.class));
 		BeanDefinitionMethodGenerator generator = new BeanDefinitionMethodGenerator(
@@ -159,9 +146,7 @@ MethodReference generateBeanDefinitionMethod(
 			}
 
 		};
-		registrations.put("testBean", generator);
-		BeanRegistrationsAotContribution contribution = new BeanRegistrationsAotContribution(
-				registrations, new LinkedMultiValueMap<>());
+		BeanRegistrationsAotContribution contribution = createContribution(generator);
 		contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
 		assertThat(beanRegistrationsCodes).hasSize(1);
 		BeanRegistrationsCode actual = beanRegistrationsCodes.get(0);
@@ -199,4 +184,9 @@ private void compile(
 				result.accept(compiled.getInstance(Consumer.class), compiled));
 	}
 
+	private BeanRegistrationsAotContribution createContribution(
+			BeanDefinitionMethodGenerator methodGenerator,String... aliases) {
+		return new BeanRegistrationsAotContribution(Map.of("testBean", new Registration(methodGenerator, aliases)));
+	}
+
 }

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

@@ -62,9 +62,10 @@ void processAheadOfTimeReturnsBeanRegistrationsAotContributionWithAliases() {
 		beanFactory.registerAlias("test", "testAlias");
 		BeanRegistrationsAotContribution contribution = processor
 				.processAheadOfTime(beanFactory);
-		assertThat(contribution).extracting("aliases")
-				.asInstanceOf(InstanceOfAssertFactories.MAP).hasEntrySatisfying("test", value ->
-						assertThat(value).asList().singleElement().isEqualTo("testAlias"));
+		assertThat(contribution).extracting("registrations").asInstanceOf(InstanceOfAssertFactories.MAP)
+				.hasEntrySatisfying("test", registration ->
+						assertThat(registration).extracting("aliases").asInstanceOf(InstanceOfAssertFactories.ARRAY)
+								.singleElement().isEqualTo("testAlias"));
 	}
 
 }