Polishing.

Remove test-only methods and move these into the tests. Improve encapsulation.

Move AotRepositoryBeanDefinitionPropertiesDecorator from config to aot.generate package due to its strong coupling with aot.generate packages.

See: #3344
Original pull request: #3351

Author: mp911de

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

@@ -52,7 +52,20 @@ public class AotQueryMethodGenerationContext {
 	private final ExpressionMarker expressionMarker;
 
 	protected AotQueryMethodGenerationContext(RepositoryInformation repositoryInformation, Method method,
-			QueryMethod queryMethod, AotRepositoryFragmentMetadata targetTypeMetadata) {
+			QueryMethod queryMethod) {
+
+		this.method = method;
+		this.annotations = MergedAnnotations.from(method);
+		this.queryMethod = queryMethod;
+		this.repositoryInformation = repositoryInformation;
+		this.targetTypeMetadata = new AotRepositoryFragmentMetadata();
+		this.targetMethodMetadata = new MethodMetadata(repositoryInformation, method);
+		this.variableNameFactory = LocalVariableNameFactory.forMethod(targetMethodMetadata);
+		this.expressionMarker = new ExpressionMarker();
+	}
+
+	AotQueryMethodGenerationContext(RepositoryInformation repositoryInformation, Method method, QueryMethod queryMethod,
+			AotRepositoryFragmentMetadata targetTypeMetadata) {
 
 		this.method = method;
 		this.annotations = MergedAnnotations.from(method);
@@ -352,4 +365,5 @@ public String localVariable(String variableName) {
 	public ExpressionMarker getExpressionMarker() {
 		return expressionMarker;
 	}
+
 }

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

@@ -13,20 +13,23 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.springframework.data.repository.config;
+package org.springframework.data.repository.aot.generate;
 
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.function.Supplier;
 
+import javax.lang.model.element.Modifier;
+
 import org.springframework.beans.factory.BeanFactory;
 import org.springframework.core.ResolvableType;
-import org.springframework.data.repository.aot.generate.RepositoryContributor;
 import org.springframework.data.repository.core.support.RepositoryComposition;
 import org.springframework.data.repository.core.support.RepositoryFactoryBeanSupport;
 import org.springframework.javapoet.CodeBlock;
+import org.springframework.javapoet.MethodSpec;
 import org.springframework.javapoet.TypeName;
+import org.springframework.javapoet.TypeSpec;
 import org.springframework.util.Assert;
 import org.springframework.util.StringUtils;
 
@@ -39,7 +42,7 @@
  * @author Christoph Strobl
  * @since 4.0
  */
-class AotRepositoryBeanDefinitionPropertiesDecorator {
+public class AotRepositoryBeanDefinitionPropertiesDecorator {
 
 	private static final Map<ResolvableType, String> RESERVED_TYPES;
 
@@ -57,10 +60,13 @@ class AotRepositoryBeanDefinitionPropertiesDecorator {
 	 * @param inheritedProperties bean definition code (containing properties and such) already added via another
 	 *          component.
 	 * @param repositoryContributor the contributor providing the actual AOT repository implementation.
+	 * @throws IllegalArgumentException if {@link RepositoryContributor#getContributedTypeName()} is not set.
 	 */
 	public AotRepositoryBeanDefinitionPropertiesDecorator(Supplier<CodeBlock> inheritedProperties,
 			RepositoryContributor repositoryContributor) {
 
+		Assert.notNull(repositoryContributor.getContributedTypeName(), "Contributed type name must not be null");
+
 		this.inheritedProperties = inheritedProperties;
 		this.repositoryContributor = repositoryContributor;
 	}
@@ -73,54 +79,58 @@ public AotRepositoryBeanDefinitionPropertiesDecorator(Supplier<CodeBlock> inheri
 	 * needs to have potential constructor arguments resolved.
 	 *
 	 * @return the decorated code block.
-	 * @throws IllegalArgumentException if {@link RepositoryContributor#getContributedTypeName()} is not set.
 	 */
 	public CodeBlock decorate() {
 
-		Assert.notNull(repositoryContributor.getContributedTypeName(), "Contributed type name must not be null");
-
 		CodeBlock.Builder builder = CodeBlock.builder();
+
 		// bring in properties as usual
 		builder.add(inheritedProperties.get());
 
-		builder.add("beanDefinition.getPropertyValues().addPropertyValue(\"repositoryFragmentsFunction\", new $T() {\n",
-				RepositoryFactoryBeanSupport.RepositoryFragmentsFunction.class);
-		builder.indent();
+		MethodSpec.Builder callbackMethod = MethodSpec.methodBuilder("getRepositoryFragments").addModifiers(Modifier.PUBLIC)
+				.returns(RepositoryComposition.RepositoryFragments.class);
 
-		builder.add("public $T getRepositoryFragments(", RepositoryComposition.RepositoryFragments.class);
-		int counter = 0;
 		for (Entry<ResolvableType, String> entry : RESERVED_TYPES.entrySet()) {
-			builder.add("$T $L", entry.getKey().toClass(), entry.getValue());
-			if (++counter < RESERVED_TYPES.size()) {
-				builder.add(", ");
-			}
+			callbackMethod.addParameter(entry.getKey().toClass(), entry.getValue());
 		}
-		builder.add(") {\n");
 
-		builder.indent();
+		callbackMethod.addCode(buildCallbackBody());
+
+		TypeSpec repositoryFragmentsFunction = TypeSpec.anonymousClassBuilder("")
+				.superclass(RepositoryFactoryBeanSupport.RepositoryFragmentsFunction.class).addMethod(callbackMethod.build())
+				.build();
 
-		for (Map.Entry<String, ResolvableType> entry : repositoryContributor.requiredArgs().entrySet()) {
+		builder.addStatement("beanDefinition.getPropertyValues().addPropertyValue($S, $L)", "repositoryFragmentsFunction",
+				repositoryFragmentsFunction);
+
+		return builder.build();
+	}
+
+	private CodeBlock buildCallbackBody() {
+
+		CodeBlock.Builder callback = CodeBlock.builder();
+
+		for (Entry<String, ResolvableType> entry : repositoryContributor.requiredArgs().entrySet()) {
 
 			TypeName argumentType = TypeName.get(entry.getValue().getType());
 			String reservedArgumentName = RESERVED_TYPES.get(entry.getValue());
 			if (reservedArgumentName == null) {
-				builder.addStatement("$1T $2L = beanFactory.getBean($1T.class)", argumentType, entry.getKey());
+				callback.addStatement("$1T $2L = beanFactory.getBean($1T.class)", argumentType, entry.getKey());
 			} else {
-				if (!reservedArgumentName.equals(entry.getKey())) {
-					builder.addStatement("$T $L = $L", argumentType, entry.getKey(), reservedArgumentName);
+
+				if (reservedArgumentName.equals(entry.getKey())) {
+					continue;
 				}
+
+				callback.addStatement("$T $L = $L", argumentType, entry.getKey(), reservedArgumentName);
 			}
 		}
 
-		builder.addStatement("return RepositoryComposition.RepositoryFragments.just(new $L($L))",
+		callback.addStatement("return $T.just(new $L($L))", RepositoryComposition.RepositoryFragments.class,
 				repositoryContributor.getContributedTypeName().getCanonicalName(),
 				StringUtils.collectionToDelimitedString(repositoryContributor.requiredArgs().keySet(), ", "));
-		builder.unindent();
-		builder.add("}\n");
-		builder.unindent();
-		builder.add("});\n");
 
-		return builder.build();
+		return callback.build();
 	}
 
 }

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

@@ -17,7 +17,6 @@
 
 import org.springframework.core.ResolvableType;
 import org.springframework.javapoet.CodeBlock;
-import org.springframework.javapoet.MethodSpec;
 
 /**
  * Builder for AOT Repository Constructors.
@@ -34,7 +33,9 @@ public interface AotRepositoryConstructorBuilder {
 	 * @param parameterName name of the parameter.
 	 * @param type parameter type.
 	 */
-	void addParameter(String parameterName, Class<?> type);
+	default void addParameter(String parameterName, Class<?> type) {
+		addParameter(parameterName, ResolvableType.forClass(type));
+	}
 
 	/**
 	 * Add constructor parameter and create a field storing its value.
@@ -51,26 +52,26 @@ default void addParameter(String parameterName, ResolvableType type) {
 	 *
 	 * @param parameterName name of the parameter.
 	 * @param type parameter type.
-	 * @param createField whether to create a field for the parameter and assign its value to the field.
+	 * @param bindToField whether to create a field for the parameter and assign its value to the field.
 	 */
-	default void addParameter(String parameterName, Class<?> type, boolean createField) {
-		addParameter(parameterName, ResolvableType.forClass(type), createField);
+	default void addParameter(String parameterName, Class<?> type, boolean bindToField) {
+		addParameter(parameterName, ResolvableType.forClass(type), bindToField);
 	}
 
 	/**
 	 * Add constructor parameter.
 	 *
 	 * @param parameterName name of the parameter.
 	 * @param type parameter type.
-	 * @param createField whether to create a field for the parameter and assign its value to the field.
+	 * @param bindToField whether to create a field for the parameter and assign its value to the field.
 	 */
-	void addParameter(String parameterName, ResolvableType type, boolean createField);
+	void addParameter(String parameterName, ResolvableType type, boolean bindToField);
 
 	/**
-	 * Add constructor customizer. Customizer is invoked after adding constructor arguments and before assigning
+	 * Add constructor body customizer. The customizer is invoked after adding constructor arguments and before assigning
 	 * constructor arguments to fields.
 	 *
-	 * @param customizer the customizer with direct access to the {@link MethodSpec.Builder constructor builder}.
+	 * @param customizer the customizer with direct access to the {@link CodeBlock.Builder constructor builder}.
 	 */
 	void customize(ConstructorCustomizer customizer);