Skip to content

Commit 6a76141

Browse files
snicollsnicoll
committed
Simplify BeanRegistrationsAotProcessor
Closes gh-32944
1 parent 34eccbe commit 6a76141

File tree

5 files changed

+38
-59
lines changed

5 files changed

+38
-59
lines changed

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

-28
This file was deleted.

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

+20-12
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616

1717
package org.springframework.beans.factory.aot;
1818

19-
import java.util.Map;
19+
import java.util.List;
2020

2121
import javax.lang.model.element.Modifier;
2222

@@ -30,6 +30,7 @@
3030
import org.springframework.aot.hint.ReflectionHints;
3131
import org.springframework.aot.hint.RuntimeHints;
3232
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
33+
import org.springframework.beans.factory.support.RegisteredBean;
3334
import org.springframework.javapoet.ClassName;
3435
import org.springframework.javapoet.CodeBlock;
3536
import org.springframework.javapoet.MethodSpec;
@@ -50,10 +51,10 @@ class BeanRegistrationsAotContribution
5051

5152
private static final String BEAN_FACTORY_PARAMETER_NAME = "beanFactory";
5253

53-
private final Map<BeanRegistrationKey, Registration> registrations;
54+
private final List<Registration> registrations;
5455

5556

56-
BeanRegistrationsAotContribution(Map<BeanRegistrationKey, Registration> registrations) {
57+
BeanRegistrationsAotContribution(List<Registration> registrations) {
5758
this.registrations = registrations;
5859
}
5960

@@ -84,13 +85,13 @@ private void generateRegisterBeanDefinitionsMethod(MethodSpec.Builder method,
8485
method.addModifiers(Modifier.PUBLIC);
8586
method.addParameter(DefaultListableBeanFactory.class, BEAN_FACTORY_PARAMETER_NAME);
8687
CodeBlock.Builder code = CodeBlock.builder();
87-
this.registrations.forEach((registeredBean, registration) -> {
88+
this.registrations.forEach(registration -> {
8889
MethodReference beanDefinitionMethod = registration.methodGenerator
8990
.generateBeanDefinitionMethod(generationContext, beanRegistrationsCode);
9091
CodeBlock methodInvocation = beanDefinitionMethod.toInvokeCodeBlock(
9192
ArgumentCodeGenerator.none(), beanRegistrationsCode.getClassName());
9293
code.addStatement("$L.registerBeanDefinition($S, $L)",
93-
BEAN_FACTORY_PARAMETER_NAME, registeredBean.beanName(), methodInvocation);
94+
BEAN_FACTORY_PARAMETER_NAME, registration.beanName(), methodInvocation);
9495
});
9596
method.addCode(code.build());
9697
}
@@ -100,30 +101,37 @@ private void generateRegisterAliasesMethod(MethodSpec.Builder method) {
100101
method.addModifiers(Modifier.PUBLIC);
101102
method.addParameter(DefaultListableBeanFactory.class, BEAN_FACTORY_PARAMETER_NAME);
102103
CodeBlock.Builder code = CodeBlock.builder();
103-
this.registrations.forEach((registeredBean, registration) -> {
104-
for (String alias : registration.aliases) {
104+
this.registrations.forEach(registration -> {
105+
for (String alias : registration.aliases()) {
105106
code.addStatement("$L.registerAlias($S, $S)", BEAN_FACTORY_PARAMETER_NAME,
106-
registeredBean.beanName(), alias);
107+
registration.beanName(), alias);
107108
}
108109
});
109110
method.addCode(code.build());
110111
}
111112

112-
private void generateRegisterHints(RuntimeHints runtimeHints, Map<BeanRegistrationKey, Registration> registrations) {
113-
registrations.keySet().forEach(beanRegistrationKey -> {
113+
private void generateRegisterHints(RuntimeHints runtimeHints, List<Registration> registrations) {
114+
registrations.forEach(registration -> {
114115
ReflectionHints hints = runtimeHints.reflection();
115-
Class<?> beanClass = beanRegistrationKey.beanClass();
116+
Class<?> beanClass = registration.registeredBean.getBeanClass();
116117
hints.registerType(beanClass, MemberCategory.INTROSPECT_PUBLIC_METHODS, MemberCategory.INTROSPECT_DECLARED_METHODS);
117118
hints.registerForInterfaces(beanClass, typeHint -> typeHint.withMembers(MemberCategory.INTROSPECT_PUBLIC_METHODS));
118119
});
119120
}
120121

121122
/**
122123
* Gather the necessary information to register a particular bean.
124+
* @param registeredBean the bean to register
123125
* @param methodGenerator the {@link BeanDefinitionMethodGenerator} to use
124126
* @param aliases the bean aliases, if any
125127
*/
126-
record Registration(BeanDefinitionMethodGenerator methodGenerator, String[] aliases) {}
128+
record Registration(RegisteredBean registeredBean, BeanDefinitionMethodGenerator methodGenerator, String[] aliases) {
129+
130+
String beanName() {
131+
return this.registeredBean.getBeanName();
132+
}
133+
134+
}
127135

128136

129137
/**

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

+6-6
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2002-2023 the original author or authors.
2+
* Copyright 2002-2024 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -16,8 +16,8 @@
1616

1717
package org.springframework.beans.factory.aot;
1818

19-
import java.util.LinkedHashMap;
20-
import java.util.Map;
19+
import java.util.ArrayList;
20+
import java.util.List;
2121

2222
import org.springframework.beans.factory.aot.BeanRegistrationsAotContribution.Registration;
2323
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
@@ -41,15 +41,15 @@ class BeanRegistrationsAotProcessor implements BeanFactoryInitializationAotProce
4141
public BeanRegistrationsAotContribution processAheadOfTime(ConfigurableListableBeanFactory beanFactory) {
4242
BeanDefinitionMethodGeneratorFactory beanDefinitionMethodGeneratorFactory =
4343
new BeanDefinitionMethodGeneratorFactory(beanFactory);
44-
Map<BeanRegistrationKey, Registration> registrations = new LinkedHashMap<>();
44+
List<Registration> registrations = new ArrayList<>();
4545

4646
for (String beanName : beanFactory.getBeanDefinitionNames()) {
4747
RegisteredBean registeredBean = RegisteredBean.of(beanFactory, beanName);
4848
BeanDefinitionMethodGenerator beanDefinitionMethodGenerator =
4949
beanDefinitionMethodGeneratorFactory.getBeanDefinitionMethodGenerator(registeredBean);
5050
if (beanDefinitionMethodGenerator != null) {
51-
registrations.put(new BeanRegistrationKey(beanName, registeredBean.getBeanClass()),
52-
new Registration(beanDefinitionMethodGenerator, beanFactory.getAliases(beanName)));
51+
registrations.add(new Registration(registeredBean, beanDefinitionMethodGenerator,
52+
beanFactory.getAliases(beanName)));
5353
}
5454
}
5555

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

+8-9
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2002-2023 the original author or authors.
2+
* Copyright 2002-2024 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -18,7 +18,6 @@
1818

1919
import java.util.ArrayList;
2020
import java.util.List;
21-
import java.util.Map;
2221
import java.util.function.BiConsumer;
2322
import java.util.function.Consumer;
2423

@@ -78,7 +77,7 @@ void applyToAppliesContribution() {
7877
RegisteredBean registeredBean = registerBean(new RootBeanDefinition(TestBean.class));
7978
BeanDefinitionMethodGenerator generator = new BeanDefinitionMethodGenerator(this.methodGeneratorFactory,
8079
registeredBean, null, List.of());
81-
BeanRegistrationsAotContribution contribution = createContribution(TestBean.class, generator);
80+
BeanRegistrationsAotContribution contribution = createContribution(registeredBean, generator);
8281
contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
8382
compile((consumer, compiled) -> {
8483
DefaultListableBeanFactory freshBeanFactory = new DefaultListableBeanFactory();
@@ -92,7 +91,7 @@ void applyToAppliesContributionWithAliases() {
9291
RegisteredBean registeredBean = registerBean(new RootBeanDefinition(TestBean.class));
9392
BeanDefinitionMethodGenerator generator = new BeanDefinitionMethodGenerator(this.methodGeneratorFactory,
9493
registeredBean, null, List.of());
95-
BeanRegistrationsAotContribution contribution = createContribution(TestBean.class, generator, "testAlias");
94+
BeanRegistrationsAotContribution contribution = createContribution(registeredBean, generator, "testAlias");
9695
contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
9796
compile((consumer, compiled) -> {
9897
DefaultListableBeanFactory freshBeanFactory = new DefaultListableBeanFactory();
@@ -109,7 +108,7 @@ void applyToWhenHasNameGeneratesPrefixedFeatureName() {
109108
RegisteredBean registeredBean = registerBean(new RootBeanDefinition(TestBean.class));
110109
BeanDefinitionMethodGenerator generator = new BeanDefinitionMethodGenerator(this.methodGeneratorFactory,
111110
registeredBean, null, List.of());
112-
BeanRegistrationsAotContribution contribution = createContribution(TestBean.class, generator);
111+
BeanRegistrationsAotContribution contribution = createContribution(registeredBean, generator);
113112
contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
114113
compile((consumer, compiled) -> {
115114
SourceFile sourceFile = compiled.getSourceFile(".*BeanDefinitions");
@@ -132,7 +131,7 @@ MethodReference generateBeanDefinitionMethod(GenerationContext generationContext
132131
}
133132

134133
};
135-
BeanRegistrationsAotContribution contribution = createContribution(TestBean.class, generator);
134+
BeanRegistrationsAotContribution contribution = createContribution(registeredBean, generator);
136135
contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
137136
assertThat(beanRegistrationsCodes).hasSize(1);
138137
BeanRegistrationsCode actual = beanRegistrationsCodes.get(0);
@@ -144,7 +143,7 @@ void applyToRegisterReflectionHints() {
144143
RegisteredBean registeredBean = registerBean(new RootBeanDefinition(Employee.class));
145144
BeanDefinitionMethodGenerator generator = new BeanDefinitionMethodGenerator(this.methodGeneratorFactory,
146145
registeredBean, null, List.of());
147-
BeanRegistrationsAotContribution contribution = createContribution(Employee.class, generator);
146+
BeanRegistrationsAotContribution contribution = createContribution(registeredBean, generator);
148147
contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
149148
assertThat(reflection().onType(Employee.class)
150149
.withMemberCategories(MemberCategory.INTROSPECT_PUBLIC_METHODS, MemberCategory.INTROSPECT_DECLARED_METHODS))
@@ -186,10 +185,10 @@ private void compile(BiConsumer<Consumer<DefaultListableBeanFactory>, Compiled>
186185
result.accept(compiled.getInstance(Consumer.class), compiled));
187186
}
188187

189-
private BeanRegistrationsAotContribution createContribution(Class<?> beanClass,
188+
private BeanRegistrationsAotContribution createContribution(RegisteredBean registeredBean,
190189
BeanDefinitionMethodGenerator methodGenerator,String... aliases) {
191190
return new BeanRegistrationsAotContribution(
192-
Map.of(new BeanRegistrationKey("testBean", beanClass), new Registration(methodGenerator, aliases)));
191+
List.of(new Registration(registeredBean, methodGenerator, aliases)));
193192
}
194193

195194
}

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

+4-4
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2002-2023 the original author or authors.
2+
* Copyright 2002-2024 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -51,7 +51,7 @@ void processAheadOfTimeReturnsBeanRegistrationsAotContributionWithRegistrations(
5151
BeanRegistrationsAotContribution contribution = processor
5252
.processAheadOfTime(beanFactory);
5353
assertThat(contribution).extracting("registrations")
54-
.asInstanceOf(InstanceOfAssertFactories.MAP).hasSize(2);
54+
.asInstanceOf(InstanceOfAssertFactories.LIST).hasSize(2);
5555
}
5656

5757
@Test
@@ -62,8 +62,8 @@ void processAheadOfTimeReturnsBeanRegistrationsAotContributionWithAliases() {
6262
beanFactory.registerAlias("test", "testAlias");
6363
BeanRegistrationsAotContribution contribution = processor
6464
.processAheadOfTime(beanFactory);
65-
assertThat(contribution).extracting("registrations").asInstanceOf(InstanceOfAssertFactories.MAP)
66-
.hasEntrySatisfying(new BeanRegistrationKey("test", TestBean.class), registration ->
65+
assertThat(contribution).extracting("registrations").asInstanceOf(InstanceOfAssertFactories.LIST)
66+
.singleElement().satisfies(registration ->
6767
assertThat(registration).extracting("aliases").asInstanceOf(InstanceOfAssertFactories.ARRAY)
6868
.singleElement().isEqualTo("testAlias"));
6969
}

0 commit comments

Comments
 (0)