Generate custom bean initialization code for types exposed via ManagedTypes during AOT.

We now replace ManagedTypes bean definitions with generated code that contain the discovered types to avoid class path scaning.

Closes: #2680
Original pull request: #2682.

src/test/java/org/springframework/data/aot/AotTestCodeContributionBuilder.java

@@ -0,0 +1,92 @@
+/*
+ * Copyright 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.data.aot;
+
+import java.util.function.Consumer;
+
+import javax.lang.model.element.Modifier;
+
+import org.mockito.Mockito;
+import org.springframework.aot.test.generate.TestGenerationContext;
+import org.springframework.aot.test.generate.compile.Compiled;
+import org.springframework.aot.test.generate.compile.TestCompiler;
+import org.springframework.beans.factory.aot.BeanRegistrationAotContribution;
+import org.springframework.beans.factory.aot.BeanRegistrationCodeFragments;
+import org.springframework.beans.factory.support.InstanceSupplier;
+import org.springframework.javapoet.CodeBlock;
+import org.springframework.javapoet.MethodSpec;
+import org.springframework.javapoet.ParameterizedTypeName;
+
+/**
+ * @author Christoph Strobl
+ */
+public class AotTestCodeContributionBuilder {
+
+	TestGenerationContext generationContext;
+	MockBeanRegistrationCode beanRegistrationCode;
+	BeanRegistrationAotContribution contribution;
+
+	static AotTestCodeContributionBuilder withContextFor(Class<?> type) {
+		return withContext(new TestGenerationContext(type));
+	}
+
+	static AotTestCodeContributionBuilder withContext(TestGenerationContext ctx) {
+
+		AotTestCodeContributionBuilder codeGenerationBuilder = new AotTestCodeContributionBuilder();
+		codeGenerationBuilder.generationContext = ctx;
+		codeGenerationBuilder.beanRegistrationCode = new MockBeanRegistrationCode(ctx);
+		return codeGenerationBuilder;
+	}
+
+	BeanRegistrationCodeFragments getFragments(BeanRegistrationAotContribution contribution) {
+
+		this.contribution = contribution;
+
+		return contribution.customizeBeanRegistrationCodeFragments(generationContext,
+				Mockito.mock(BeanRegistrationCodeFragments.class));
+	}
+
+	AotTestCodeContributionBuilder writeContentFor(BeanRegistrationAotContribution contribution) {
+
+		CodeBlock codeBlock = getFragments(contribution).generateInstanceSupplierCode(generationContext,
+				beanRegistrationCode, null, false);
+
+		Class<?> beanType = Object.class;
+		try {
+			beanType = contribution instanceof RegisteredBeanAotContribution
+					? ((RegisteredBeanAotContribution) contribution).getSource().getBeanClass()
+					: Object.class;
+		} catch (Exception e) {}
+
+		ParameterizedTypeName parameterizedReturnTypeName = ParameterizedTypeName.get(InstanceSupplier.class, beanType);
+		beanRegistrationCode.getTypeBuilder().set(type -> {
+			type.addModifiers(Modifier.PUBLIC);
+			type.addMethod(MethodSpec.methodBuilder("get").addModifiers(Modifier.PUBLIC).returns(parameterizedReturnTypeName)
+					.addStatement("return $L", codeBlock).build());
+		});
+
+		return this;
+	}
+
+	public void compile() {
+		compile(it -> {});
+	}
+
+	public void compile(Consumer<Compiled> compiled) {
+		generationContext.writeGeneratedContent();
+		TestCompiler.forSystem().withFiles(generationContext.getGeneratedFiles()).compile(compiled);
+	}
+}

src/test/java/org/springframework/data/aot/BeanRegistrationContributionAssert.java

@@ -20,7 +20,6 @@ import static org.mockito.Mockito.*;
 import java.util.function.Consumer;
 
 import org.assertj.core.api.AbstractAssert;
-
 import org.springframework.aot.generate.GenerationContext;
 import org.springframework.aot.test.generate.TestGenerationContext;
 import org.springframework.beans.factory.aot.BeanRegistrationAotContribution;

src/test/java/org/springframework/data/aot/DeferredTypeBuilder.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright 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.data.aot;
+
+import java.util.function.Consumer;
+
+import org.springframework.javapoet.TypeSpec;
+import org.springframework.javapoet.TypeSpec.Builder;
+import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
+
+/**
+ * @author Christoph Strobl
+ */
+public class DeferredTypeBuilder implements Consumer<Builder> {
+
+	@Nullable
+	private Consumer<Builder> type;
+
+	@Override
+	public void accept(Builder type) {
+		Assert.notNull(this.type, "No type builder set");
+		this.type.accept(type);
+	}
+
+	public void set(Consumer<Builder> type) {
+		this.type = type;
+	}
+}

src/test/java/org/springframework/data/aot/ManagedTypesBeanRegistrationAotProcessorUnitTests.java

@@ -19,21 +19,32 @@ import static org.assertj.core.api.Assertions.*;
 import static org.mockito.ArgumentMatchers.*;
 import static org.mockito.Mockito.*;
 
+import java.util.Arrays;
 import java.util.Collections;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Consumer;
 
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
 import org.springframework.aot.generate.GenerationContext;
 import org.springframework.aot.hint.predicate.RuntimeHintsPredicates;
 import org.springframework.aot.test.generate.TestGenerationContext;
 import org.springframework.beans.factory.BeanCreationException;
 import org.springframework.beans.factory.aot.BeanRegistrationAotContribution;
+import org.springframework.beans.factory.aot.BeanRegistrationCodeFragments;
 import org.springframework.beans.factory.support.BeanDefinitionBuilder;
 import org.springframework.beans.factory.support.DefaultListableBeanFactory;
+import org.springframework.beans.factory.support.InstanceSupplier;
 import org.springframework.beans.factory.support.RegisteredBean;
 import org.springframework.beans.factory.support.RootBeanDefinition;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.data.aot.ManagedTypesRegistrationAotContribution.ManagedTypesInstanceCodeFragment;
 import org.springframework.data.domain.ManagedTypes;
+import org.springframework.javapoet.MethodSpec;
+import org.springframework.javapoet.MethodSpec.Builder;
+import org.springframework.test.util.ReflectionTestUtils;
 
 /**
  * @author Christoph Strobl
@@ -47,6 +58,9 @@ class ManagedTypesBeanRegistrationAotProcessorUnitTests {
 	final RootBeanDefinition myManagedTypesDefinition = (RootBeanDefinition) BeanDefinitionBuilder
 			.rootBeanDefinition(MyManagedTypes.class).getBeanDefinition();
 
+	final RootBeanDefinition invocationCountingManagedTypesDefinition = (RootBeanDefinition) BeanDefinitionBuilder
+			.rootBeanDefinition(InvocationRecordingManagedTypes.class).getBeanDefinition();
+
 	DefaultListableBeanFactory beanFactory;
 
 	@BeforeEach
@@ -154,6 +168,79 @@ class ManagedTypesBeanRegistrationAotProcessorUnitTests {
 		verify(beanFactory).getBean(eq("commons.managed-types"), eq(ManagedTypes.class));
 	}
 
+	@Test // GH-2680
+	void generatesInstanceSupplierCodeFragmentToAvoidDuplicateInvocations() {
+
+		beanFactory.registerBeanDefinition("commons.managed-types", invocationCountingManagedTypesDefinition);
+		RegisteredBean registeredBean = RegisteredBean.of(beanFactory, "commons.managed-types");
+
+		BeanRegistrationAotContribution contribution = createPostProcessor("commons")
+				.processAheadOfTime(RegisteredBean.of(beanFactory, "commons.managed-types"));
+
+		AotTestCodeContributionBuilder.withContextFor(this.getClass()).writeContentFor(contribution).compile(it -> {
+
+			InvocationRecordingManagedTypes sourceTypes = beanFactory.getBean(InvocationRecordingManagedTypes.class);
+			assertThat(sourceTypes.getCounter()).isOne();
+
+			InstanceSupplier<InvocationRecordingManagedTypes> types = ReflectionTestUtils
+					.invokeMethod(it.getAllCompiledClasses().iterator().next(), "instance");
+			try {
+				assertThat(types.get(registeredBean).source).isNotSameAs(sourceTypes);
+			} catch (Exception e) {
+				throw new RuntimeException(e);
+			}
+		});
+	}
+
+	@Test // GH-2680
+	void generatesInstanceSupplierCodeFragmentForTypeWithCustomFactoryMethod() {
+
+		beanFactory.registerBeanDefinition("commons.managed-types",
+				BeanDefinitionBuilder.rootBeanDefinition(StoreManagedTypesWithCustomFactoryMethod.class).getBeanDefinition());
+
+		RegisteredBean registeredBean = RegisteredBean.of(beanFactory, "commons.managed-types");
+
+		BeanRegistrationAotContribution contribution = createPostProcessor("commons").processAheadOfTime(registeredBean);
+
+		AotTestCodeContributionBuilder.withContextFor(this.getClass()).writeContentFor(contribution).compile(it -> {
+
+			InstanceSupplier<StoreManagedTypesWithCustomFactoryMethod> types = ReflectionTestUtils
+					.invokeMethod(it.getAllCompiledClasses().iterator().next(), "instance");
+
+			try {
+				assertThat(types.get(registeredBean).toList()).containsExactlyInAnyOrder(A.class, B.class);
+			} catch (Exception e) {
+				throw new RuntimeException(e);
+			}
+		});
+	}
+
+	@Test // GH-2680
+	void canGenerateCodeReturnsTrueIfFactoryMethodPresent() {
+
+		beanFactory.registerBeanDefinition("managed-types", managedTypesDefinition);
+		RegisteredBean registeredBean = RegisteredBean.of(beanFactory, "managed-types");
+
+		ManagedTypesInstanceCodeFragment fragment = new ManagedTypesInstanceCodeFragment(
+				ManagedTypes.from(A.class, B.class), registeredBean, Mockito.mock(BeanRegistrationCodeFragments.class));
+		Builder methodBuilder = MethodSpec.methodBuilder("instance");
+		fragment.generateInstanceFactory(methodBuilder);
+
+		assertThat(fragment.canGenerateCode()).isTrue();
+	}
+
+	@Test // GH-2680
+	void canGenerateCodeReturnsFalseIfNoFactoryMethodPresent() {
+
+		beanFactory.registerBeanDefinition("managed-types", myManagedTypesDefinition);
+		RegisteredBean registeredBean = RegisteredBean.of(beanFactory, "managed-types");
+
+		ManagedTypesInstanceCodeFragment fragment = new ManagedTypesInstanceCodeFragment(
+				ManagedTypes.from(A.class, B.class), registeredBean, Mockito.mock(BeanRegistrationCodeFragments.class));
+
+		assertThat(fragment.canGenerateCode()).isFalse();
+	}
+
 	private ManagedTypesBeanRegistrationAotProcessor createPostProcessor(String moduleIdentifier) {
 		ManagedTypesBeanRegistrationAotProcessor postProcessor = new ManagedTypesBeanRegistrationAotProcessor();
 		postProcessor.setModuleIdentifier(moduleIdentifier);
@@ -172,5 +259,67 @@ class ManagedTypesBeanRegistrationAotProcessorUnitTests {
 		}
 	}
 
+	static class StoreManagedTypesWithFactoryMethodOfClassNames implements ManagedTypes {
+		@Override
+		public void forEach(Consumer<Class<?>> action) {
+			// just do nothing ¯\_(ツ)_/¯
+		}
+	}
+
+	public static class StoreManagedTypesWithCustomFactoryMethod implements ManagedTypes {
+
+		private ManagedTypes source;
+
+		public StoreManagedTypesWithCustomFactoryMethod() {
+			source = it -> Arrays.asList(A.class, B.class).forEach(it);
+		}
+
+		public StoreManagedTypesWithCustomFactoryMethod(ManagedTypes source) {
+			this.source = source;
+		}
+
+		public static StoreManagedTypesWithCustomFactoryMethod of(ManagedTypes source) {
+			return new StoreManagedTypesWithCustomFactoryMethod(source);
+		}
+
+		@Override
+		public void forEach(Consumer<Class<?>> action) {
+			source.forEach(action);
+		}
+	}
+
+	public static class InvocationRecordingManagedTypes implements ManagedTypes {
+
+		private AtomicInteger counter = new AtomicInteger(0);
+		private ManagedTypes source = ManagedTypes.from(A.class, B.class);
+
+		public static InvocationRecordingManagedTypes from(ManagedTypes source) {
+
+			InvocationRecordingManagedTypes newInstance = new InvocationRecordingManagedTypes();
+			newInstance.source = source;
+			return newInstance;
+		}
+
+		@Override
+		public void forEach(Consumer<Class<?>> action) {
+
+			counter.getAndIncrement();
+			source.forEach(action);
+		}
+
+		public int getCounter() {
+			return counter.get();
+		}
+	}
+
 	static class NotManagedTypes {}
+
+	@Configuration(proxyBeanMethods = false)
+	public static class EntityManagerWithPackagesToScanConfiguration {
+
+		@Bean(name = "commons.managed-types")
+		ManagedTypes managedTypes() {
+			return ManagedTypes.from(A.class, B.class);
+		}
+	}
 }

src/test/java/org/springframework/data/aot/MockBeanRegistrationCode.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright 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.data.aot;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.springframework.aot.generate.GeneratedClass;
+import org.springframework.aot.generate.GeneratedMethods;
+import org.springframework.aot.generate.GenerationContext;
+import org.springframework.aot.generate.MethodReference;
+import org.springframework.beans.factory.aot.BeanRegistrationCode;
+import org.springframework.javapoet.ClassName;
+
+/**
+ * @author Christoph Strobl
+ */
+public class MockBeanRegistrationCode implements BeanRegistrationCode {
+
+	private final GeneratedClass generatedClass;
+
+	private final List<MethodReference> instancePostProcessors = new ArrayList<>();
+
+	private final DeferredTypeBuilder typeBuilder = new DeferredTypeBuilder();
+
+	public MockBeanRegistrationCode(GenerationContext generationContext) {
+		this.generatedClass = generationContext.getGeneratedClasses().addForFeature("TestCode", this.typeBuilder);
+	}
+
+	public DeferredTypeBuilder getTypeBuilder() {
+		return this.typeBuilder;
+	}
+
+	@Override
+	public ClassName getClassName() {
+		return this.generatedClass.getName();
+	}
+
+	@Override
+	public GeneratedMethods getMethods() {
+		return this.generatedClass.getMethods();
+	}
+
+	@Override
+	public void addInstancePostProcessor(MethodReference methodReference) {
+		this.instancePostProcessors.add(methodReference);
+	}
+}