Migrate AOT tests to use GeneratedClasses and refine/polish AOT APIs

Migrate all AOT tests to make use of `GeneratedClasses` rather than
directly generating Java files. This commit also refines and polishes
AOT APIs to being greater consistency.

Specifically:

	- The `MethodGenerator` interface has been removed in favor of
	  working directly with `GeneratedMethods`.
	- The visibility of several constructors and methods has been
	  reduced to package-private.
	- The `using(...)` and `builder` methods have been removed in
	  favor of setting the `Consumer` callbacks directly as
	  constructor arguments.
	- Variable names for builders are now named `type` or `method`
	  depending on what they're building.

Closes gh-28831

spring-context/src/main/java/org/springframework/context/annotation/ConfigurationClassPostProcessor.java

@@ -534,10 +534,9 @@ public class ConfigurationClassPostProcessor implements BeanDefinitionRegistryPo
 			Map<String, String> mappings = buildImportAwareMappings();
 			if (!mappings.isEmpty()) {
 				GeneratedMethod generatedMethod = beanFactoryInitializationCode
-						.getMethodGenerator()
-						.generateMethod("addImportAwareBeanPostProcessors")
-						.using(builder -> generateAddPostProcessorMethod(builder,
-								mappings));
+						.getMethods()
+						.add("addImportAwareBeanPostProcessors", method ->
+								generateAddPostProcessorMethod(method, mappings));
 				beanFactoryInitializationCode
 						.addInitializer(MethodReference.of(generatedMethod.getName()));
 				ResourceHints hints = generationContext.getRuntimeHints().resources();
@@ -546,14 +545,14 @@ public class ConfigurationClassPostProcessor implements BeanDefinitionRegistryPo
 			}
 		}
 
-		private void generateAddPostProcessorMethod(MethodSpec.Builder builder,
+		private void generateAddPostProcessorMethod(MethodSpec.Builder method,
 				Map<String, String> mappings) {
 
-			builder.addJavadoc(
+			method.addJavadoc(
 					"Add ImportAwareBeanPostProcessor to support ImportAware beans");
-			builder.addModifiers(Modifier.PRIVATE);
-			builder.addParameter(DefaultListableBeanFactory.class, BEAN_FACTORY_VARIABLE);
-			builder.addCode(generateAddPostProcessorCode(mappings));
+			method.addModifiers(Modifier.PRIVATE);
+			method.addParameter(DefaultListableBeanFactory.class, BEAN_FACTORY_VARIABLE);
+			method.addCode(generateAddPostProcessorCode(mappings));
 		}
 
 		private CodeBlock generateAddPostProcessorCode(Map<String, String> mappings) {

spring-context/src/main/java/org/springframework/context/aot/ApplicationContextAotGenerator.java

@@ -16,7 +16,6 @@
 
 package org.springframework.context.aot;
 
-import org.springframework.aot.generate.GeneratedClass;
 import org.springframework.aot.generate.GenerationContext;
 import org.springframework.beans.factory.BeanFactory;
 import org.springframework.beans.factory.support.DefaultListableBeanFactory;
@@ -47,15 +46,11 @@ public class ApplicationContextAotGenerator {
 	public ClassName generateApplicationContext(GenericApplicationContext applicationContext,
 			GenerationContext generationContext) {
 		applicationContext.refreshForAotProcessing();
-		DefaultListableBeanFactory beanFactory = applicationContext
-				.getDefaultListableBeanFactory();
-		ApplicationContextInitializationCodeGenerator codeGenerator = new ApplicationContextInitializationCodeGenerator();
-		new BeanFactoryInitializationAotContributions(beanFactory).applyTo(generationContext,
-				codeGenerator);
-		GeneratedClass applicationContextInitializer = generationContext.getGeneratedClasses()
-				.forFeature("ApplicationContextInitializer")
-				.generate(codeGenerator.generateJavaFile());
-		return applicationContextInitializer.getName();
+		DefaultListableBeanFactory beanFactory = applicationContext.getDefaultListableBeanFactory();
+		ApplicationContextInitializationCodeGenerator codeGenerator =
+				new ApplicationContextInitializationCodeGenerator(generationContext);
+		new BeanFactoryInitializationAotContributions(beanFactory).applyTo(generationContext, codeGenerator);
+		return codeGenerator.getGeneratedClass().getName();
 	}
 
 }

spring-context/src/main/java/org/springframework/context/aot/ApplicationContextInitializationCodeGenerator.java

@@ -18,12 +18,12 @@ package org.springframework.context.aot;
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.function.Consumer;
 
 import javax.lang.model.element.Modifier;
 
+import org.springframework.aot.generate.GeneratedClass;
 import org.springframework.aot.generate.GeneratedMethods;
-import org.springframework.aot.generate.MethodGenerator;
+import org.springframework.aot.generate.GenerationContext;
 import org.springframework.aot.generate.MethodReference;
 import org.springframework.beans.factory.aot.BeanFactoryInitializationCode;
 import org.springframework.beans.factory.support.DefaultListableBeanFactory;
@@ -44,45 +44,41 @@ import org.springframework.javapoet.TypeSpec;
 class ApplicationContextInitializationCodeGenerator
 		implements BeanFactoryInitializationCode {
 
+	private static final String INITIALIZE_METHOD = "initialize";
+
 	private static final String APPLICATION_CONTEXT_VARIABLE = "applicationContext";
 
-
-	private final GeneratedMethods generatedMethods = new GeneratedMethods();
-
 	private final List<MethodReference> initializers = new ArrayList<>();
 
+	private final GeneratedClass generatedClass;
 
-	@Override
-	public MethodGenerator getMethodGenerator() {
-		return this.generatedMethods;
+
+	ApplicationContextInitializationCodeGenerator(GenerationContext generationContext) {
+		this.generatedClass = generationContext.getGeneratedClasses()
+				.addForFeature("ApplicationContextInitializer", this::generateType);
+		this.generatedClass.reserveMethodNames(INITIALIZE_METHOD);
 	}
 
-	@Override
-	public void addInitializer(MethodReference methodReference) {
-		this.initializers.add(methodReference);
+
+	private void generateType(TypeSpec.Builder type) {
+		type.addJavadoc(
+				"{@link $T} to restore an application context based on previous AOT processing.",
+				ApplicationContextInitializer.class);
+		type.addModifiers(Modifier.PUBLIC);
+		type.addSuperinterface(ParameterizedTypeName.get(
+				ApplicationContextInitializer.class, GenericApplicationContext.class));
+		type.addMethod(generateInitializeMethod());
 	}
 
-	Consumer<TypeSpec.Builder> generateJavaFile() {
-		return builder -> {
-			builder.addJavadoc(
-					"{@link $T} to restore an application context based on previous AOT processing.",
-					ApplicationContextInitializer.class);
-			builder.addModifiers(Modifier.PUBLIC);
-			builder.addSuperinterface(ParameterizedTypeName.get(
-					ApplicationContextInitializer.class, GenericApplicationContext.class));
-			builder.addMethod(generateInitializeMethod());
-			this.generatedMethods.doWithMethodSpecs(builder::addMethod);
-		};
-	}
 
 	private MethodSpec generateInitializeMethod() {
-		MethodSpec.Builder builder = MethodSpec.methodBuilder("initialize");
-		builder.addAnnotation(Override.class);
-		builder.addModifiers(Modifier.PUBLIC);
-		builder.addParameter(GenericApplicationContext.class,
+		MethodSpec.Builder method = MethodSpec.methodBuilder(INITIALIZE_METHOD);
+		method.addAnnotation(Override.class);
+		method.addModifiers(Modifier.PUBLIC);
+		method.addParameter(GenericApplicationContext.class,
 				APPLICATION_CONTEXT_VARIABLE);
-		builder.addCode(generateInitializeCode());
-		return builder.build();
+		method.addCode(generateInitializeCode());
+		return method.build();
 	}
 
 	private CodeBlock generateInitializeCode() {
@@ -93,10 +89,23 @@ class ApplicationContextInitializationCodeGenerator
 		builder.addStatement("$L.setAutowireCandidateResolver(new $T())",
 				BEAN_FACTORY_VARIABLE, ContextAnnotationAutowireCandidateResolver.class);
 		for (MethodReference initializer : this.initializers) {
-			builder.addStatement(
-					initializer.toInvokeCodeBlock(CodeBlock.of(BEAN_FACTORY_VARIABLE)));
+			builder.addStatement(initializer.toInvokeCodeBlock(CodeBlock.of(BEAN_FACTORY_VARIABLE)));
 		}
 		return builder.build();
 	}
 
+	GeneratedClass getGeneratedClass() {
+		return this.generatedClass;
+	}
+
+	@Override
+	public GeneratedMethods getMethods() {
+		return this.generatedClass.getMethods();
+	}
+
+	@Override
+	public void addInitializer(MethodReference methodReference) {
+		this.initializers.add(methodReference);
+	}
+
 }

spring-context/src/test/java/org/springframework/context/annotation/ConfigurationClassPostProcessorAotContributionTests.java

@@ -39,10 +39,8 @@ import org.springframework.context.testfixture.context.generator.annotation.Impo
 import org.springframework.context.testfixture.context.generator.annotation.ImportConfiguration;
 import org.springframework.core.testfixture.aot.generate.TestGenerationContext;
 import org.springframework.javapoet.CodeBlock;
-import org.springframework.javapoet.JavaFile;
 import org.springframework.javapoet.MethodSpec;
 import org.springframework.javapoet.ParameterizedTypeName;
-import org.springframework.javapoet.TypeSpec;
 import org.springframework.lang.Nullable;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -56,21 +54,26 @@ import static org.assertj.core.api.Assertions.entry;
  */
 class ConfigurationClassPostProcessorAotContributionTests {
 
-	private DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();
+	private final InMemoryGeneratedFiles generatedFiles;
 
-	private InMemoryGeneratedFiles generatedFiles = new InMemoryGeneratedFiles();
+	private final DefaultGenerationContext generationContext;
 
-	private DefaultGenerationContext generationContext = new TestGenerationContext(
-			this.generatedFiles);
+	private final MockBeanFactoryInitializationCode beanFactoryInitializationCode;
+
+
+	ConfigurationClassPostProcessorAotContributionTests() {
+		this.generatedFiles = new InMemoryGeneratedFiles();
+		this.generationContext = new TestGenerationContext(this.generatedFiles);
+		this.beanFactoryInitializationCode = new MockBeanFactoryInitializationCode(this.generationContext);
+	}
 
-	private MockBeanFactoryInitializationCode beanFactoryInitializationCode = new MockBeanFactoryInitializationCode();
 
 	@Test
 	void applyToWhenHasImportAwareConfigurationRegistersBeanPostProcessorWithMapEntry() {
 		BeanFactoryInitializationAotContribution contribution = getContribution(
 				ImportConfiguration.class);
 		contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
-		testCompiledResult((initializer, compiled) -> {
+		compile((initializer, compiled) -> {
 			DefaultListableBeanFactory freshBeanFactory = new DefaultListableBeanFactory();
 			initializer.accept(freshBeanFactory);
 			ImportAwareAotBeanPostProcessor postProcessor = (ImportAwareAotBeanPostProcessor) freshBeanFactory
@@ -89,8 +92,8 @@ class ConfigurationClassPostProcessorAotContributionTests {
 				.singleElement()
 				.satisfies(resourceHint -> assertThat(resourceHint.getIncludes())
 						.map(ResourcePatternHint::getPattern)
-						.containsOnly(
-								"org/springframework/context/testfixture/context/generator/annotation/ImportConfiguration.class"));
+						.containsOnly("org/springframework/context/testfixture/context/generator/annotation/"
+								+ "ImportConfiguration.class"));
 	}
 
 	@Test
@@ -100,38 +103,28 @@ class ConfigurationClassPostProcessorAotContributionTests {
 
 	@Nullable
 	private BeanFactoryInitializationAotContribution getContribution(Class<?> type) {
-		this.beanFactory.registerBeanDefinition("configuration",
-				new RootBeanDefinition(type));
+		DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();
+		beanFactory.registerBeanDefinition("configuration", new RootBeanDefinition(type));
 		ConfigurationClassPostProcessor postProcessor = new ConfigurationClassPostProcessor();
-		postProcessor.postProcessBeanFactory(this.beanFactory);
-		return postProcessor.processAheadOfTime(this.beanFactory);
+		postProcessor.postProcessBeanFactory(beanFactory);
+		return postProcessor.processAheadOfTime(beanFactory);
 	}
 
 	@SuppressWarnings("unchecked")
-	private void testCompiledResult(
-			BiConsumer<Consumer<DefaultListableBeanFactory>, Compiled> result) {
-		JavaFile javaFile = createJavaFile();
+	private void compile(BiConsumer<Consumer<DefaultListableBeanFactory>, Compiled> result) {
+		MethodReference methodReference = this.beanFactoryInitializationCode
+				.getInitializers().get(0);
+		this.beanFactoryInitializationCode.getTypeBuilder().set(type -> {
+			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(methodReference.toInvokeCodeBlock(CodeBlock.of("beanFactory")))
+					.build());
+		});
 		this.generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(this.generatedFiles).compile(javaFile::writeTo,
-				compiled -> result.accept(compiled.getInstance(Consumer.class),
-						compiled));
-	}
-
-	private JavaFile createJavaFile() {
-		MethodReference methodReference = this.beanFactoryInitializationCode.getInitializers()
-				.get(0);
-		TypeSpec.Builder builder = TypeSpec.classBuilder("TestConsumer");
-		builder.addModifiers(Modifier.PUBLIC);
-		builder.addSuperinterface(ParameterizedTypeName.get(Consumer.class,
-				DefaultListableBeanFactory.class));
-		builder.addMethod(MethodSpec.methodBuilder("accept").addModifiers(Modifier.PUBLIC)
-				.addParameter(DefaultListableBeanFactory.class, "beanFactory")
-				.addStatement(
-						methodReference.toInvokeCodeBlock(CodeBlock.of("beanFactory")))
-				.build());
-		this.beanFactoryInitializationCode.getMethodGenerator()
-				.doWithMethodSpecs(builder::addMethod);
-		return JavaFile.builder("__", builder.build()).build();
+		TestCompiler.forSystem().withFiles(this.generatedFiles).compile(compiled ->
+				result.accept(compiled.getInstance(Consumer.class), compiled));
 	}
 
 	private void assertPostProcessorEntry(ImportAwareAotBeanPostProcessor postProcessor,

spring-context/src/test/java/org/springframework/context/aot/ApplicationContextAotGeneratorTests.java

@@ -188,14 +188,11 @@ class ApplicationContextAotGeneratorTests {
 			BiConsumer<ApplicationContextInitializer<GenericApplicationContext>, Compiled> result) {
 		ApplicationContextAotGenerator generator = new ApplicationContextAotGenerator();
 		InMemoryGeneratedFiles generatedFiles = new InMemoryGeneratedFiles();
-		DefaultGenerationContext generationContext = new TestGenerationContext(
-				generatedFiles);
+		DefaultGenerationContext generationContext = new TestGenerationContext(generatedFiles);
 		generator.generateApplicationContext(applicationContext, generationContext);
 		generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(generatedFiles)
-				.compile(compiled -> result.accept(
-						compiled.getInstance(ApplicationContextInitializer.class),
-						compiled));
+		TestCompiler.forSystem().withFiles(generatedFiles).compile(compiled ->
+				result.accept(compiled.getInstance(ApplicationContextInitializer.class), compiled));
 	}
 
 	private GenericApplicationContext toFreshApplicationContext(

spring-context/src/test/java/org/springframework/context/aot/ReflectiveProcessorBeanRegistrationAotProcessorTests.java

@@ -25,6 +25,7 @@ import java.lang.annotation.Target;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.aot.generate.GenerationContext;
+import org.springframework.aot.generate.InMemoryGeneratedFiles;
 import org.springframework.aot.hint.MemberCategory;
 import org.springframework.aot.hint.RuntimeHints;
 import org.springframework.aot.hint.TypeReference;
@@ -53,7 +54,7 @@ class ReflectiveProcessorBeanRegistrationAotProcessorTests {
 
 	private final ReflectiveProcessorBeanRegistrationAotProcessor processor = new ReflectiveProcessorBeanRegistrationAotProcessor();
 
-	private final GenerationContext generationContext = new TestGenerationContext();
+	private final GenerationContext generationContext = new TestGenerationContext(new InMemoryGeneratedFiles());
 
 	@Test
 	void shouldIgnoreNonAnnotatedType() {

spring-context/src/test/java/org/springframework/context/aot/RuntimeHintsBeanFactoryInitializationAotProcessorTests.java

@@ -26,6 +26,7 @@ import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.aot.generate.GenerationContext;
+import org.springframework.aot.generate.InMemoryGeneratedFiles;
 import org.springframework.aot.hint.ResourceBundleHint;
 import org.springframework.aot.hint.RuntimeHints;
 import org.springframework.aot.hint.RuntimeHintsRegistrar;
@@ -55,7 +56,7 @@ class RuntimeHintsBeanFactoryInitializationAotProcessorTests {
 
 	@BeforeEach
 	void setup() {
-		this.generationContext = new TestGenerationContext();
+		this.generationContext = new TestGenerationContext(new InMemoryGeneratedFiles());
 		this.generator = new ApplicationContextAotGenerator();
 	}