Polishing.

Improve factory methods. Integrate auditing hints in AuditingBeanRegistrationAotProcessor. Reduce visibility, improve naming.

Original Pull Request: #2624

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

@@ -15,10 +15,14 @@
  */
 package org.springframework.data.aot;
 
+import org.springframework.aop.SpringProxy;
+import org.springframework.aop.framework.Advised;
+import org.springframework.aot.hint.RuntimeHints;
+import org.springframework.aot.hint.TypeReference;
 import org.springframework.beans.factory.aot.BeanRegistrationAotContribution;
 import org.springframework.beans.factory.aot.BeanRegistrationAotProcessor;
 import org.springframework.beans.factory.support.RegisteredBean;
-import org.springframework.data.aot.hint.AuditingHints;
+import org.springframework.core.DecoratingProxy;
 import org.springframework.data.domain.AuditorAware;
 import org.springframework.data.domain.ReactiveAuditorAware;
 import org.springframework.data.repository.util.ReactiveWrappers;
@@ -27,31 +31,43 @@ import org.springframework.lang.Nullable;
 import org.springframework.util.ClassUtils;
 
 /**
+ * {@link BeanRegistrationAotProcessor} to register runtime hints for beans that implement auditor-aware interfaces to
+ * enable JDK proxy creation.
+ *
  * @author Christoph Strobl
+ * @author Mark Paluch
  * @since 3.0
  */
-public class AuditingBeanRegistrationAotProcessor implements BeanRegistrationAotProcessor {
+class AuditingBeanRegistrationAotProcessor implements BeanRegistrationAotProcessor {
 
 	@Nullable
 	@Override
 	public BeanRegistrationAotContribution processAheadOfTime(RegisteredBean registeredBean) {
 
 		if (isAuditingHandler(registeredBean)) {
-			return (generationContext, beanRegistrationCode) -> new AuditingHints.AuditingRuntimeHints()
-					.registerHints(generationContext.getRuntimeHints(), registeredBean.getBeanFactory().getBeanClassLoader());
+			return (generationContext, beanRegistrationCode) -> registerSpringProxy(AuditorAware.class,
+					generationContext.getRuntimeHints());
 		}
+
 		if (ReactiveWrappers.isAvailable(ReactiveLibrary.PROJECT_REACTOR) && isReactiveAuditorAware(registeredBean)) {
-			return (generationContext, beanRegistrationCode) -> new AuditingHints.ReactiveAuditingRuntimeHints()
-					.registerHints(generationContext.getRuntimeHints(), registeredBean.getBeanFactory().getBeanClassLoader());
+			return (generationContext, beanRegistrationCode) -> registerSpringProxy(ReactiveAuditorAware.class,
+					generationContext.getRuntimeHints());
 		}
+
 		return null;
 	}
 
-	boolean isAuditingHandler(RegisteredBean bean) {
+	private static boolean isAuditingHandler(RegisteredBean bean) {
 		return ClassUtils.isAssignable(AuditorAware.class, bean.getBeanClass());
 	}
 
-	boolean isReactiveAuditorAware(RegisteredBean bean) {
+	private static boolean isReactiveAuditorAware(RegisteredBean bean) {
 		return ClassUtils.isAssignable(ReactiveAuditorAware.class, bean.getBeanClass());
 	}
+
+	private static void registerSpringProxy(Class<?> type, RuntimeHints runtimeHints) {
+
+		runtimeHints.proxies().registerJdkProxy(TypeReference.of(type), TypeReference.of(SpringProxy.class),
+				TypeReference.of(Advised.class), TypeReference.of(DecoratingProxy.class));
+	}
 }

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

@@ -16,7 +16,6 @@
 package org.springframework.data.aot;
 
 import java.util.Collections;
-import java.util.function.Function;
 import java.util.function.Supplier;
 
 import org.apache.commons.logging.Log;
@@ -41,24 +40,12 @@ import org.springframework.util.ObjectUtils;
  *
  * @author Christoph Strobl
  * @author John Blum
- * @see org.springframework.beans.factory.config.ConfigurableListableBeanFactory
- * @see org.springframework.beans.factory.aot.BeanFactoryInitializationAotContribution
- * @see org.springframework.beans.factory.aot.BeanFactoryInitializationAotProcessor
  * @since 3.0
  */
-public class SpringDataBeanFactoryInitializationAotProcessor implements BeanFactoryInitializationAotProcessor {
+public class ManagedTypesBeanFactoryInitializationAotProcessor implements BeanFactoryInitializationAotProcessor {
 
 	private static final Log logger = LogFactory.getLog(BeanFactoryInitializationAotProcessor.class);
 
-	private static final Function<Object, Object> arrayToListFunction = target ->
-		ObjectUtils.isArray(target) ? CollectionUtils.arrayToList(target) : target;
-
-	private static final Function<Object, Object> asSingletonSetFunction = target ->
-		!(target instanceof Iterable<?>) ? Collections.singleton(target) : target;
-
-	private static final Function<Object, Object> constructorArgumentFunction =
-		arrayToListFunction.andThen(asSingletonSetFunction);
-
 	@Nullable
 	@Override
 	public BeanFactoryInitializationAotContribution processAheadOfTime(ConfigurableListableBeanFactory beanFactory) {
@@ -70,42 +57,61 @@ public class SpringDataBeanFactoryInitializationAotProcessor implements BeanFact
 	private void processManagedTypes(ConfigurableListableBeanFactory beanFactory) {
 
 		if (beanFactory instanceof BeanDefinitionRegistry registry) {
+
 			for (String beanName : beanFactory.getBeanNamesForType(ManagedTypes.class)) {
-
-				BeanDefinition beanDefinition = beanFactory.getBeanDefinition(beanName);
-
-				if (hasConstructorArguments(beanDefinition)) {
-
-					ValueHolder argumentValue = beanDefinition.getConstructorArgumentValues()
-						.getArgumentValue(0, null, null, null);
-
-					if (argumentValue.getValue() instanceof Supplier supplier) {
-
-						if (logger.isDebugEnabled()) {
-							logger.info(String.format("Replacing ManagedType bean definition %s.", beanName));
-						}
-
-						Object value = constructorArgumentFunction.apply(supplier.get());
-
-						BeanDefinition beanDefinitionReplacement = newManagedTypeBeanDefinition(value);
-
-						registry.removeBeanDefinition(beanName);
-						registry.registerBeanDefinition(beanName, beanDefinitionReplacement);
-					}
-				}
+				postProcessManagedTypes(beanFactory, registry, beanName);
 			}
 		}
 	}
 
+	private void postProcessManagedTypes(ConfigurableListableBeanFactory beanFactory, BeanDefinitionRegistry registry,
+			String beanName) {
+
+		BeanDefinition beanDefinition = beanFactory.getBeanDefinition(beanName);
+
+		if (hasConstructorArguments(beanDefinition)) {
+
+			ValueHolder argumentValue = beanDefinition.getConstructorArgumentValues().getArgumentValue(0, null, null, null);
+
+			if (argumentValue.getValue()instanceof Supplier supplier) {
+
+				if (logger.isDebugEnabled()) {
+					logger.info(String.format("Replacing ManagedType bean definition %s.", beanName));
+				}
+
+				Object value = potentiallyWrapToIterable(supplier.get());
+
+				BeanDefinition beanDefinitionReplacement = newManagedTypeBeanDefinition(beanDefinition.getBeanClassName(),
+						value);
+
+				registry.removeBeanDefinition(beanName);
+				registry.registerBeanDefinition(beanName, beanDefinitionReplacement);
+			}
+		}
+	}
+
+	private static Object potentiallyWrapToIterable(Object value) {
+
+		if (ObjectUtils.isArray(value)) {
+			return CollectionUtils.arrayToList(value);
+		}
+
+		if (value instanceof Iterable<?>) {
+			return value;
+		}
+
+		return Collections.singleton(value);
+	}
+
 	private boolean hasConstructorArguments(BeanDefinition beanDefinition) {
 		return !beanDefinition.getConstructorArgumentValues().isEmpty();
 	}
 
-	private BeanDefinition newManagedTypeBeanDefinition(Object constructorArgumentValue) {
+	private BeanDefinition newManagedTypeBeanDefinition(String managedTypesClassName, Object constructorArgumentValue) {
 
-		return BeanDefinitionBuilder.rootBeanDefinition(ManagedTypes.class)
-			.setFactoryMethod("fromIterable")
-			.addConstructorArgValue(constructorArgumentValue)
-			.getBeanDefinition();
+		return BeanDefinitionBuilder.rootBeanDefinition(managedTypesClassName) //
+				.setFactoryMethod("fromIterable") //
+				.addConstructorArgValue(constructorArgumentValue) //
+				.getBeanDefinition();
 	}
 }

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

@@ -43,7 +43,7 @@ import org.springframework.util.StringUtils;
 public class ManagedTypesBeanRegistrationAotProcessor implements BeanRegistrationAotProcessor {
 
 	private final Log logger = LogFactory.getLog(getClass());
-	@Nullable private String moduleIdentifier;
+	private @Nullable String moduleIdentifier;
 
 	@Override
 	public BeanRegistrationAotContribution processAheadOfTime(RegisteredBean registeredBean) {

src/main/java/org/springframework/data/aot/hint/AuditingHints.java

@@ -1,74 +0,0 @@
-/*
- * 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.hint;
-
-import org.springframework.aop.SpringProxy;
-import org.springframework.aop.framework.Advised;
-import org.springframework.aot.hint.RuntimeHints;
-import org.springframework.aot.hint.RuntimeHintsRegistrar;
-import org.springframework.aot.hint.TypeReference;
-import org.springframework.core.DecoratingProxy;
-import org.springframework.data.domain.AuditorAware;
-import org.springframework.data.domain.ReactiveAuditorAware;
-import org.springframework.lang.Nullable;
-
-/**
- * Runtime hint registrars for Spring Data modules that provide auditing functionality.
- * <p>
- * Mainly for internal use within the framework.
- *
- * @author Christoph Strobl
- * @since 3.0
- * @see RuntimeHintsRegistrar
- */
-public class AuditingHints {
-
-	/**
-	 * {@link RuntimeHintsRegistrar} to be applied when auditing is enabled. Can be used along with
-	 * {@link org.springframework.context.annotation.ImportRuntimeHints @ImportRuntimeHints} for conditional registration.
-	 *
-	 * @author Christoph Strobl
-	 * @since 3.0
-	 */
-	public static class AuditingRuntimeHints implements RuntimeHintsRegistrar {
-
-		@Override
-		public void registerHints(RuntimeHints hints, @Nullable ClassLoader classLoader) {
-			registerSpringProxy(AuditorAware.class, hints);
-		}
-	}
-
-	/**
-	 * {@link RuntimeHintsRegistrar} to be applied when reactive auditing is enabled. Can be used along with
-	 * {@link org.springframework.context.annotation.ImportRuntimeHints @ImportRuntimeHints} for conditional registration.
-	 *
-	 * @author Christoph Strobl
-	 * @since 3.0
-	 */
-	public static class ReactiveAuditingRuntimeHints implements RuntimeHintsRegistrar {
-
-		@Override
-		public void registerHints(RuntimeHints hints, @Nullable ClassLoader classLoader) {
-			registerSpringProxy(ReactiveAuditorAware.class, hints);
-		}
-	}
-
-	private static void registerSpringProxy(Class<?> type, RuntimeHints runtimeHints) {
-
-		runtimeHints.proxies().registerJdkProxy(TypeReference.of(type), TypeReference.of(SpringProxy.class),
-				TypeReference.of(Advised.class), TypeReference.of(DecoratingProxy.class));
-	}
-}

src/main/java/org/springframework/data/aot/hint/RepositoryRuntimeHints.java

@@ -40,7 +40,7 @@ import org.springframework.lang.Nullable;
  * @author Christoph Strobl
  * @since 3.0
  */
-public class RepositoryRuntimeHints implements RuntimeHintsRegistrar {
+class RepositoryRuntimeHints implements RuntimeHintsRegistrar {
 
 	@Override
 	public void registerHints(RuntimeHints hints, @Nullable ClassLoader classLoader) {

src/main/java/org/springframework/data/domain/ManagedTypes.java

@@ -16,6 +16,7 @@
 package org.springframework.data.domain;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.function.Consumer;
@@ -31,6 +32,7 @@ import org.springframework.util.Assert;
  *
  * @author Christoph Strobl
  * @author John Blum
+ * @author Mark Paluch
  * @since 3.0
  */
 @FunctionalInterface
@@ -47,6 +49,16 @@ public interface ManagedTypes {
 		return fromIterable(Collections.emptySet());
 	}
 
+	/**
+	 * Factory method used to construct {@link ManagedTypes} from the given array of {@link Class types}.
+	 *
+	 * @param types array of {@link Class types} used to initialize the {@link ManagedTypes}; must not be {@literal null}.
+	 * @return new instance of {@link ManagedTypes} initialized from {@link Class types}.
+	 */
+	static ManagedTypes from(Class<?>... types) {
+		return fromIterable(Arrays.asList(types));
+	}
+
 	/**
 	 * Factory method used to construct {@link ManagedTypes} from the given, required {@link Iterable} of {@link Class
 	 * types}.