Add AOT processing of bean aliases

This commit adds AOT processing of bean aliases.

Closes gh-29391

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

@@ -30,12 +30,14 @@ import org.springframework.beans.factory.support.DefaultListableBeanFactory;
 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.
+ * register bean definitions and aliases.
  *
  * @author Phillip Webb
+ * @author Sebastien Deleuze
  * @since 6.0
  * @see BeanRegistrationsAotProcessor
  */
@@ -46,11 +48,14 @@ class BeanRegistrationsAotContribution
 
 	private final Map<String, BeanDefinitionMethodGenerator> registrations;
 
+	private final MultiValueMap<String, String> aliases;
+
 
 	BeanRegistrationsAotContribution(
-			Map<String, BeanDefinitionMethodGenerator> registrations) {
+			Map<String, BeanDefinitionMethodGenerator> registrations, MultiValueMap<String, String> aliases) {
 
 		this.registrations = registrations;
+		this.aliases = aliases;
 	}
 
 
@@ -64,12 +69,15 @@ class BeanRegistrationsAotContribution
 					type.addModifiers(Modifier.PUBLIC);
 				});
 		BeanRegistrationsCodeGenerator codeGenerator = new BeanRegistrationsCodeGenerator(generatedClass);
-		GeneratedMethod generatedMethod = codeGenerator.getMethods().add("registerBeanDefinitions", method ->
-				generateRegisterMethod(method, generationContext, codeGenerator));
-		beanFactoryInitializationCode.addInitializer(generatedMethod.toMethodReference());
+		GeneratedMethod generatedBeanDefinitionsMethod = codeGenerator.getMethods().add("registerBeanDefinitions", method ->
+				generateRegisterBeanDefinitionsMethod(method, generationContext, codeGenerator));
+		beanFactoryInitializationCode.addInitializer(generatedBeanDefinitionsMethod.toMethodReference());
+		GeneratedMethod generatedAliasesMethod = codeGenerator.getMethods().add("registerAliases",
+				this::generateRegisterAliasesMethod);
+		beanFactoryInitializationCode.addInitializer(generatedAliasesMethod.toMethodReference());
 	}
 
-	private void generateRegisterMethod(MethodSpec.Builder method,
+	private void generateRegisterBeanDefinitionsMethod(MethodSpec.Builder method,
 			GenerationContext generationContext,
 			BeanRegistrationsCode beanRegistrationsCode) {
 
@@ -91,6 +99,18 @@ class BeanRegistrationsAotContribution
 		method.addCode(code.build());
 	}
 
+	private void generateRegisterAliasesMethod(MethodSpec.Builder method) {
+		method.addJavadoc("Register the aliases.");
+		method.addModifiers(Modifier.PUBLIC);
+		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)));
+		method.addCode(code.build());
+	}
+
 
 	/**
 	 * {@link BeanRegistrationsCode} with generation support.

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

@@ -21,12 +21,15 @@ import java.util.Map;
 
 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
  * @since 6.0
  */
 class BeanRegistrationsAotProcessor implements BeanFactoryInitializationAotProcessor {
@@ -36,6 +39,7 @@ class BeanRegistrationsAotProcessor implements BeanFactoryInitializationAotProce
 		BeanDefinitionMethodGeneratorFactory beanDefinitionMethodGeneratorFactory =
 				new BeanDefinitionMethodGeneratorFactory(beanFactory);
 		Map<String, BeanDefinitionMethodGenerator> registrations = new LinkedHashMap<>();
+		MultiValueMap<String, String> aliases = new LinkedMultiValueMap<>();
 		for (String beanName : beanFactory.getBeanDefinitionNames()) {
 			RegisteredBean registeredBean = RegisteredBean.of(beanFactory, beanName);
 			BeanDefinitionMethodGenerator beanDefinitionMethodGenerator = beanDefinitionMethodGeneratorFactory
@@ -43,11 +47,14 @@ class BeanRegistrationsAotProcessor implements BeanFactoryInitializationAotProce
 			if (beanDefinitionMethodGenerator != null) {
 				registrations.put(beanName, beanDefinitionMethodGenerator);
 			}
+			for (String alias : beanFactory.getAliases(beanName)) {
+				aliases.add(beanName, alias);
+			}
 		}
 		if (registrations.isEmpty()) {
 			return null;
 		}
-		return new BeanRegistrationsAotContribution(registrations);
+		return new BeanRegistrationsAotContribution(registrations, aliases);
 	}
 
 }

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

@@ -42,9 +42,12 @@ import org.springframework.core.test.io.support.MockSpringFactoriesLoader;
 import org.springframework.core.test.tools.Compiled;
 import org.springframework.core.test.tools.SourceFile;
 import org.springframework.core.test.tools.TestCompiler;
+import org.springframework.javapoet.ClassName;
 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;
 
@@ -52,6 +55,7 @@ import static org.assertj.core.api.Assertions.assertThat;
  * Tests for {@link BeanRegistrationsAotContribution}.
  *
  * @author Phillip Webb
+ * @author Sebastien Deleuze
  */
 class BeanRegistrationsAotContributionTests {
 
@@ -84,7 +88,7 @@ class BeanRegistrationsAotContributionTests {
 				Collections.emptyList());
 		registrations.put("testBean", generator);
 		BeanRegistrationsAotContribution contribution = new BeanRegistrationsAotContribution(
-				registrations);
+				registrations, new LinkedMultiValueMap<>());
 		contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
 		compile((consumer, compiled) -> {
 			DefaultListableBeanFactory freshBeanFactory = new DefaultListableBeanFactory();
@@ -93,6 +97,27 @@ class BeanRegistrationsAotContributionTests {
 		});
 	}
 
+	@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);
+		contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
+		compile((consumer, compiled) -> {
+			DefaultListableBeanFactory freshBeanFactory = new DefaultListableBeanFactory();
+			consumer.accept(freshBeanFactory);
+			assertThat(freshBeanFactory.getAliases("testBean")).containsExactly("testAlias");
+		});
+	}
+
 	@Test
 	void applyToWhenHasNameGeneratesPrefixedFeatureName() {
 		this.generationContext = new TestGenerationContext(
@@ -106,7 +131,7 @@ class BeanRegistrationsAotContributionTests {
 				Collections.emptyList());
 		registrations.put("testBean", generator);
 		BeanRegistrationsAotContribution contribution = new BeanRegistrationsAotContribution(
-				registrations);
+				registrations, new LinkedMultiValueMap<>());
 		contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
 		compile((consumer, compiled) -> {
 			SourceFile sourceFile = compiled.getSourceFile(".*BeanDefinitions");
@@ -136,7 +161,7 @@ class BeanRegistrationsAotContributionTests {
 		};
 		registrations.put("testBean", generator);
 		BeanRegistrationsAotContribution contribution = new BeanRegistrationsAotContribution(
-				registrations);
+				registrations, new LinkedMultiValueMap<>());
 		contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
 		assertThat(beanRegistrationsCodes).hasSize(1);
 		BeanRegistrationsCode actual = beanRegistrationsCodes.get(0);
@@ -152,17 +177,21 @@ class BeanRegistrationsAotContributionTests {
 	@SuppressWarnings({ "unchecked", "cast" })
 	private void compile(
 			BiConsumer<Consumer<DefaultListableBeanFactory>, Compiled> result) {
-		MethodReference methodReference = this.beanFactoryInitializationCode
+		MethodReference beanRegistrationsMethodReference = this.beanFactoryInitializationCode
 				.getInitializers().get(0);
+		MethodReference aliasesMethodReference = this.beanFactoryInitializationCode
+				.getInitializers().get(1);
 		this.beanFactoryInitializationCode.getTypeBuilder().set(type -> {
-			CodeBlock methodInvocation = methodReference.toInvokeCodeBlock(
-					ArgumentCodeGenerator.of(DefaultListableBeanFactory.class, "beanFactory"),
-					this.beanFactoryInitializationCode.getClassName());
+			ArgumentCodeGenerator beanFactory = ArgumentCodeGenerator.of(DefaultListableBeanFactory.class, "beanFactory");
+			ClassName className = this.beanFactoryInitializationCode.getClassName();
+			CodeBlock beanRegistrationsMethodInvocation = beanRegistrationsMethodReference.toInvokeCodeBlock(beanFactory, className);
+			CodeBlock aliasesMethodInvocation = aliasesMethodReference.toInvokeCodeBlock(beanFactory, className);
 			type.addModifiers(Modifier.PUBLIC);
 			type.addSuperinterface(ParameterizedTypeName.get(Consumer.class, DefaultListableBeanFactory.class));
 			type.addMethod(MethodSpec.methodBuilder("accept").addModifiers(Modifier.PUBLIC)
 					.addParameter(DefaultListableBeanFactory.class, "beanFactory")
-					.addStatement(methodInvocation)
+					.addStatement(beanRegistrationsMethodInvocation)
+					.addStatement(aliasesMethodInvocation)
 					.build());
 		});
 		this.generationContext.writeGeneratedContent();

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

@@ -30,6 +30,7 @@ import static org.assertj.core.api.Assertions.assertThat;
  * Tests for {@link BeanRegistrationsAotProcessor}.
  *
  * @author Phillip Webb
+ * @author Sebastien Deleuze
  */
 class BeanRegistrationsAotProcessorTests {
 
@@ -53,4 +54,17 @@ class BeanRegistrationsAotProcessorTests {
 				.asInstanceOf(InstanceOfAssertFactories.MAP).containsKeys("b1", "b2");
 	}
 
+	@Test
+	void processAheadOfTimeReturnsBeanRegistrationsAotContributionWithAliases() {
+		BeanRegistrationsAotProcessor processor = new BeanRegistrationsAotProcessor();
+		DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();
+		beanFactory.registerBeanDefinition("test", new RootBeanDefinition(TestBean.class));
+		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"));
+	}
+
 }