Refine AOT Repositories infrastructure

pom.xml

@@ -5,7 +5,7 @@
 
 	<groupId>org.springframework.data</groupId>
 	<artifactId>spring-data-commons</artifactId>
-	<version>4.0.0-SNAPSHOT</version>
+	<version>4.0.0-GH-3267-SNAPSHOT</version>
 
 	<name>Spring Data Core</name>
 	<description>Core Spring concepts underpinning every Spring Data module.</description>

src/main/java/org/springframework/data/aot/AotContext.java

@@ -24,6 +24,8 @@
 import java.util.function.Consumer;
 
 import org.jspecify.annotations.Nullable;
+
+import org.springframework.aot.generate.GenerationContext;
 import org.springframework.beans.factory.BeanFactory;
 import org.springframework.beans.factory.NoSuchBeanDefinitionException;
 import org.springframework.beans.factory.config.BeanDefinition;
@@ -40,9 +42,7 @@
 
 /**
  * The context in which the AOT processing happens. Grants access to the {@link ConfigurableListableBeanFactory
- * beanFactory} and {@link ClassLoader}. Holds a few convenience methods to check if a type
- * {@link TypeIntrospector#isTypePresent() is present} and allows resolution of them through {@link TypeIntrospector}
- * and {@link IntrospectedBeanDefinition}.
+ * beanFactory} and {@link ClassLoader}.
  * <p>
  * Mainly for internal use within the framework.
  *
@@ -99,7 +99,7 @@ static AotContext from(BeanFactory beanFactory, Environment environment) {
 	 * @param moduleName name of the module. Can be {@literal null} or {@literal empty}, in which case it will only check
 	 *          the general {@link #GENERATED_REPOSITORIES_ENABLED} flag.
 	 * @return indicator if repository code generation is enabled.
-	 * @since 5.0
+	 * @since 4.0
 	 */
 	default boolean isGeneratedRepositoriesEnabled(@Nullable String moduleName) {
 
@@ -119,13 +119,13 @@ default boolean isGeneratedRepositoriesEnabled(@Nullable String moduleName) {
 	}
 
 	/**
-	 * Checks if repository metadata file writing is enabled by checking environment variables for general
-	 * enablement ({@link #GENERATED_REPOSITORIES_JSON_ENABLED})
+	 * Checks if repository metadata file writing is enabled by checking environment variables for general enablement
+	 * ({@link #GENERATED_REPOSITORIES_JSON_ENABLED})
 	 * <p>
 	 * Unset properties are considered being {@literal true}.
 	 *
 	 * @return indicator if repository metadata should be written
-	 * @since 5.0
+	 * @since 4.0
 	 */
 	default boolean isGeneratedRepositoriesMetadataEnabled() {
 		return getEnvironment().getProperty(GENERATED_REPOSITORIES_JSON_ENABLED, Boolean.class, true);
@@ -177,8 +177,12 @@ default ClassLoader getRequiredClassLoader() {
 	 *
 	 * @param typeName {@link String name} of the {@link Class type} to evaluate; must not be {@literal null}.
 	 * @return the type introspector for further type-based introspection.
+	 * @deprecated since 4.0 as this isn't widely used and can be easily implemented within user code.
 	 */
-	TypeIntrospector introspectType(String typeName);
+	@Deprecated(since = "4.0", forRemoval = true)
+	default TypeIntrospector introspectType(String typeName) {
+		throw new UnsupportedOperationException(); // preparation for implementation removal.
+	}
 
 	/**
 	 * Returns a new {@link TypeScanner} used to scan for {@link Class types} that will be contributed to the AOT
@@ -201,7 +205,9 @@ default TypeScanner getTypeScanner() {
 	 * @param packageNames {@link Collection} of {@link String package names} to scan.
 	 * @return a {@link Set} of {@link Class types} found during the scan.
 	 * @see #getTypeScanner()
+	 * @deprecated since 4.0, use {@link #getTypeScanner()} directly
 	 */
+	@Deprecated(since = "4.0", forRemoval = true)
 	default Set<Class<?>> scanPackageForTypes(Collection<Class<? extends Annotation>> identifyingAnnotations,
 			Collection<String> packageNames) {
 
@@ -214,7 +220,9 @@ default Set<Class<?>> scanPackageForTypes(Collection<Class<? extends Annotation>
 	 *
 	 * @param reference {@link BeanReference} to the managed bean.
 	 * @return the introspected bean definition.
+	 * @deprecated since 4.0, use {@link #getBeanFactory()} and interact with the bean factory directly.
 	 */
+	@Deprecated(since = "4.0", forRemoval = true)
 	default IntrospectedBeanDefinition introspectBeanDefinition(BeanReference reference) {
 		return introspectBeanDefinition(reference.getBeanName());
 	}
@@ -225,40 +233,50 @@ default IntrospectedBeanDefinition introspectBeanDefinition(BeanReference refere
 	 *
 	 * @param beanName {@link String} containing the {@literal name} of the bean to evaluate; must not be {@literal null}.
 	 * @return the introspected bean definition.
+	 * @deprecated since 4.0, use {@link #getBeanFactory()} and interact with the bean factory directly.
 	 */
-	IntrospectedBeanDefinition introspectBeanDefinition(String beanName);
+	@Deprecated(since = "4.0", forRemoval = true)
+	default IntrospectedBeanDefinition introspectBeanDefinition(String beanName) {
+		throw new UnsupportedOperationException(); // preparation for implementation removal.
+	}
 
 	/**
 	 * Obtain a {@link AotTypeConfiguration} for the given {@link ResolvableType} to customize the AOT processing for the
-	 * given type.
+	 * given type. Repeated calls to the same type will result in merging the configuration.
 	 *
 	 * @param resolvableType the resolvable type to configure.
 	 * @param configurationConsumer configuration consumer function.
+	 * @since 4.0
 	 */
 	default void typeConfiguration(ResolvableType resolvableType, Consumer<AotTypeConfiguration> configurationConsumer) {
 		typeConfiguration(resolvableType.toClass(), configurationConsumer);
 	}
 
 	/**
 	 * Obtain a {@link AotTypeConfiguration} for the given {@link ResolvableType} to customize the AOT processing for the
-	 * given type.
+	 * given type. Repeated calls to the same type will result in merging the configuration.
 	 *
 	 * @param type the type to configure.
 	 * @param configurationConsumer configuration consumer function.
+	 * @since 4.0
 	 */
 	void typeConfiguration(Class<?> type, Consumer<AotTypeConfiguration> configurationConsumer);
 
 	/**
-	 * Return all type configurations registered with this {@link AotContext}.
+	 * Contribute type configurations to the given {@link GenerationContext}. This method is called once per
+	 * {@link GenerationContext} after all type configurations have been registered.
 	 *
-	 * @return all type configurations registered with this {@link AotContext}.
+	 * @param generationContext the context to contribute the type configurations to.
 	 */
-	Collection<AotTypeConfiguration> typeConfigurations();
+	void contributeTypeConfigurations(GenerationContext generationContext);
 
 	/**
 	 * Type-based introspector to resolve {@link Class} from a type name and to introspect the bean factory for presence
 	 * of beans.
+	 *
+	 * @deprecated since 4.0 as this isn't widely used and can be easily implemented within user code.
 	 */
+	@Deprecated(since = "4.0", forRemoval = true)
 	interface TypeIntrospector {
 
 		/**
@@ -318,7 +336,10 @@ default void ifTypePresent(Consumer<Class<?>> action) {
 
 	/**
 	 * Interface defining introspection methods for bean definitions.
+	 *
+	 * @deprecated since 4.0 as this isn't widely used and can be easily implemented within user code.
 	 */
+	@Deprecated(since = "4.0", forRemoval = true)
 	interface IntrospectedBeanDefinition {
 
 		/**

src/main/java/org/springframework/data/aot/AotTypeConfiguration.java

@@ -22,19 +22,15 @@
 
 import org.springframework.aop.SpringProxy;
 import org.springframework.aop.framework.Advised;
-import org.springframework.aot.generate.GenerationContext;
 import org.springframework.aot.hint.MemberCategory;
 import org.springframework.aot.hint.TypeReference;
 import org.springframework.core.DecoratingProxy;
-import org.springframework.core.env.Environment;
 import org.springframework.data.projection.TargetAware;
 
 /**
  * Configuration object that captures various AOT configuration aspects of types within the data context by offering
  * predefined methods to register native configuration necessary for data binding, projection proxy definitions, AOT
  * cglib bytecode generation and other common tasks.
- * <p>
- * On {@link #contribute(Environment, GenerationContext)} the configuration is added to the {@link GenerationContext}.
  *
  * @author Christoph Strobl
  * @since 4.0
@@ -134,11 +130,4 @@ default AotTypeConfiguration proxyInterface(Class<?>... proxyInterfaces) {
 	 */
 	AotTypeConfiguration forQuerydsl();
 
-	/**
-	 * Write the configuration to the given {@link GenerationContext}.
-	 *
-	 * @param environment must not be {@literal null}.
-	 * @param generationContext must not be {@literal null}.
-	 */
-	void contribute(Environment environment, GenerationContext generationContext);
 }

src/main/java/org/springframework/data/aot/DefaultAotContext.java

@@ -17,7 +17,6 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -30,9 +29,11 @@
 import java.util.stream.Stream;
 
 import org.jspecify.annotations.Nullable;
+
 import org.springframework.aot.generate.GenerationContext;
 import org.springframework.aot.hint.MemberCategory;
 import org.springframework.aot.hint.TypeReference;
+import org.springframework.aot.hint.annotation.ReflectiveRuntimeHintsRegistrar;
 import org.springframework.beans.factory.BeanFactory;
 import org.springframework.beans.factory.NoSuchBeanDefinitionException;
 import org.springframework.beans.factory.aot.AotProcessingException;
@@ -55,14 +56,15 @@
  * @author Christoph Strobl
  * @since 3.0
  */
+@SuppressWarnings("removal")
 class DefaultAotContext implements AotContext {
 
 	private final AotMappingContext mappingContext;
 	private final ConfigurableListableBeanFactory factory;
 
-	// TODO: should we reuse the config or potentially have multiple ones with different settings for the same type
-	private final Map<Class<?>, AotTypeConfiguration> typeConfigurations = new HashMap<>();
+	private final Map<Class<?>, ContextualTypeConfiguration> typeConfigurations = new HashMap<>();
 	private final Environment environment;
+	private final ReflectiveRuntimeHintsRegistrar runtimeHintsRegistrar = new ReflectiveRuntimeHintsRegistrar();
 
 	public DefaultAotContext(BeanFactory beanFactory, Environment environment) {
 		this(beanFactory, environment, new AotMappingContext());
@@ -101,10 +103,13 @@ public void typeConfiguration(Class<?> type, Consumer<AotTypeConfiguration> conf
 	}
 
 	@Override
-	public Collection<AotTypeConfiguration> typeConfigurations() {
-		return typeConfigurations.values();
+	public void contributeTypeConfigurations(GenerationContext generationContext) {
+		typeConfigurations.forEach((type, configuration) -> {
+			configuration.contribute(this.environment, generationContext);
+		});
 	}
 
+	@SuppressWarnings("removal")
 	class DefaultTypeIntrospector implements TypeIntrospector {
 
 		private final String typeName;
@@ -144,6 +149,7 @@ public List<String> getBeanNames() {
 		}
 	}
 
+	@SuppressWarnings("removal")
 	class DefaultIntrospectedBeanDefinition implements IntrospectedBeanDefinition {
 
 		private final String beanName;
@@ -227,7 +233,6 @@ public AotTypeConfiguration forQuerydsl() {
 			return this;
 		}
 
-		@Override
 		public void contribute(Environment environment, GenerationContext generationContext) {
 
 			try {
@@ -254,6 +259,7 @@ private void doContribute(Environment environment, GenerationContext generationC
 			}
 
 			if (forDataBinding) {
+				runtimeHintsRegistrar.registerRuntimeHints(generationContext.getRuntimeHints(), type);
 				TypeContributor.contribute(type, Set.of(TypeContributor.DATA_NAMESPACE), generationContext);
 			}
 

src/main/java/org/springframework/data/aot/ManagedTypesBeanRegistrationAotProcessor.java

@@ -119,7 +119,7 @@ private ManagedTypes resolveManagedTypes(RegisteredBean registeredBean) {
 
 	/**
 	 * Hook to provide a customized flavor of {@link BeanRegistrationAotContribution}. By overriding this method calls to
-	 * {@link #contributeType(ResolvableType, GenerationContext, AotContext)} might no longer be issued.
+	 * {@link #registerTypeHints(ResolvableType, AotContext, GenerationContext)} might no longer be issued.
 	 *
 	 * @param aotContext never {@literal null}.
 	 * @param managedTypes never {@literal null}.
@@ -128,7 +128,7 @@ private ManagedTypes resolveManagedTypes(RegisteredBean registeredBean) {
 	protected BeanRegistrationAotContribution contribute(AotContext aotContext, ManagedTypes managedTypes,
 			RegisteredBean registeredBean) {
 		return new ManagedTypesRegistrationAotContribution(aotContext, managedTypes, registeredBean,
-				typeCollectorCustomizer(), this::contributeType);
+				typeCollectorCustomizer(), this::registerTypeHints);
 	}
 
 	/**
@@ -140,13 +140,14 @@ protected BeanRegistrationAotContribution contribute(AotContext aotContext, Mana
 	protected Consumer<TypeCollector> typeCollectorCustomizer() {
 		return typeCollector -> {};
 	}
+
 	/**
 	 * Hook to contribute configuration for a given {@literal type}.
 	 *
 	 * @param type never {@literal null}.
 	 * @param generationContext never {@literal null}.
 	 */
-	protected void contributeType(ResolvableType type, GenerationContext generationContext, AotContext aotContext) {
+	protected void registerTypeHints(ResolvableType type, AotContext aotContext, GenerationContext generationContext) {
 
 		if (logger.isDebugEnabled()) {
 			logger.debug(String.format("Contributing type information for [%s]", type.getType()));
@@ -156,10 +157,6 @@ protected void contributeType(ResolvableType type, GenerationContext generationC
 
 		configureTypeContribution(type.toClass(), aotContext);
 
-		aotContext.typeConfiguration(type, config -> {
-			config.contribute(environment.get(), generationContext);
-		});
-
 		TypeUtils.resolveUsedAnnotations(type.toClass()).forEach(
 				annotation -> TypeContributor.contribute(annotation.getType(), annotationNamespaces, generationContext));
 	}

src/main/java/org/springframework/data/aot/ManagedTypesRegistrationAotContribution.java

@@ -98,8 +98,10 @@ public void applyTo(GenerationContext generationContext, BeanRegistrationCode be
 
 		if (!types.isEmpty()) {
 			TypeCollector.inspect(typeCollectorCustomizer, types)
-					.forEach(type -> contributionAction.register(type, generationContext, aotContext));
+					.forEach(type -> contributionAction.register(type, aotContext, generationContext));
 		}
+
+		aotContext.contributeTypeConfigurations(generationContext);
 	}
 
 	@Override
@@ -117,7 +119,7 @@ public RegisteredBean getSource() {
 	}
 
 	interface TypeRegistration {
-		void register(ResolvableType type, GenerationContext generationContext, AotContext aotContext);
+		void register(ResolvableType type, AotContext aotContext, GenerationContext generationContext);
 	}
 
 	/**

src/main/java/org/springframework/data/javapoet/LordOfTheStrings.java

@@ -858,10 +858,37 @@ public static class TypedReturnBuilder extends ReturnBuilderSupport {
 		 */
 		@Contract("_ -> this")
 		public TypedReturnBuilder number(String resultToReturn) {
+
 			return whenBoxedLong("$1L != null ? $1L.longValue() : null", resultToReturn)
 					.whenLong("$1L != null ? $1L.longValue() : 0L", resultToReturn)
 					.whenBoxedInteger("$1L != null ? $1L.intValue() : null", resultToReturn)
-					.whenInt("$1L != null ? $1L.intValue() : 0", resultToReturn);
+					.whenInt("$1L != null ? $1L.intValue() : 0", resultToReturn)
+					.whenBoxed(Byte.class, "$1L != null ? $1L.byteValue() : null", resultToReturn)
+					.when(byte.class, "$1L != null ? $1L.byteValue() : 0", resultToReturn)
+					.whenBoxed(Short.class, "$1L != null ? $1L.shortValue() : null", resultToReturn)
+					.when(short.class, "$1L != null ? $1L.shortValue() : 0", resultToReturn)
+					.whenBoxed(Double.class, "$1L != null ? $1L.doubleValue() : null", resultToReturn)
+					.when(double.class, "$1L != null ? $1L.doubleValue() : 0", resultToReturn)
+					.whenBoxed(Float.class, "$1L != null ? $1L.floatValue() : null", resultToReturn)
+					.when(float.class, "$1L != null ? $1L.floatValue() : 0f", resultToReturn);
+		}
+
+		/**
+		 * Add return statements for numeric types if the given {@code resultToReturn} points to a non-nullable
+		 * {@link Number}. Considers all primitive numeric types assuming that {@code resultToReturn} is never
+		 * {@literal null}.
+		 *
+		 * @param resultToReturn the argument or variable name holding the result.
+		 * @return {@code this} builder.
+		 */
+		@Contract("_ -> this")
+		public TypedReturnBuilder nonNullableNumber(String resultToReturn) {
+			return whenPrimitiveOrBoxed(long.class, "$1L.longValue()", resultToReturn)
+					.whenPrimitiveOrBoxed(int.class, "$1L.intValue()", resultToReturn)
+					.whenPrimitiveOrBoxed(short.class, "$1L.shortValue()", resultToReturn)
+					.whenPrimitiveOrBoxed(byte.class, "$1L.byteValue()", resultToReturn)
+					.whenPrimitiveOrBoxed(float.class, "$1L.floatValue()", resultToReturn)
+					.whenPrimitiveOrBoxed(double.class, "$1L.doubleValue()", resultToReturn);
 		}
 
 		/**

src/main/java/org/springframework/data/repository/aot/generate/AotRepositoryBeanDefinitionPropertiesDecorator.java

@@ -41,8 +41,8 @@
 
 /**
  * Delegate to decorate AOT {@code BeanDefinition} properties during AOT processing. Adds a {@link CodeBlock} for the
- * fragment function that resolves {@link RepositoryContributor#requiredArgs()} from the {@link BeanFactory} and
- * provides them to the generated repository fragment.
+ * fragment function that resolves {@link RepositoryContributor#getAotFragmentMetadata()} from the {@link BeanFactory}
+ * and provides them to the generated repository fragment.
  *
  * @author Mark Paluch
  * @author Christoph Strobl
@@ -128,7 +128,8 @@ private CodeBlock buildCallbackBody() {
 		CodeBlock.Builder callback = CodeBlock.builder();
 		List<Object> arguments = new ArrayList<>();
 
-		for (Entry<String, ConstructorArgument> entry : repositoryContributor.getConstructorArguments().entrySet()) {
+		for (Entry<String, ConstructorArgument> entry : repositoryContributor.getAotFragmentMetadata()
+				.getConstructorArguments().entrySet()) {
 
 			ConstructorArgument argument = entry.getValue();
 			AotRepositoryConstructorBuilder.ParameterOrigin parameterOrigin = argument.parameterOrigin();