Add ExpressionMarker abstraction for obtaining Method in AOT generated code.

ExpressionMarker is a stateful abstraction that helps creating local classes used to obtain the enclosing method. The code generation will only add the local class when needed. Prior to this change markers had been added unconditionally to each and every method.

Closes #3338

Author: christophstrobl

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

@@ -21,7 +21,6 @@
 import java.util.List;
 
 import org.jspecify.annotations.Nullable;
-
 import org.springframework.core.MethodParameter;
 import org.springframework.core.ResolvableType;
 import org.springframework.core.annotation.MergedAnnotation;
@@ -50,6 +49,7 @@ public class AotQueryMethodGenerationContext {
 	private final AotRepositoryFragmentMetadata targetTypeMetadata;
 	private final MethodMetadata targetMethodMetadata;
 	private final VariableNameFactory variableNameFactory;
+	private final ExpressionMarker expressionMarker;
 
 	protected AotQueryMethodGenerationContext(RepositoryInformation repositoryInformation, Method method,
 			QueryMethod queryMethod, AotRepositoryFragmentMetadata targetTypeMetadata) {
@@ -61,6 +61,7 @@ protected AotQueryMethodGenerationContext(RepositoryInformation repositoryInform
 		this.targetTypeMetadata = targetTypeMetadata;
 		this.targetMethodMetadata = new MethodMetadata(repositoryInformation, method);
 		this.variableNameFactory = LocalVariableNameFactory.forMethod(targetMethodMetadata);
+		this.expressionMarker = new ExpressionMarker();
 	}
 
 	MethodMetadata getTargetMethodMetadata() {
@@ -342,4 +343,13 @@ public String localVariable(String variableName) {
 		return getParameterName(queryMethod.getParameters().getScoreRangeIndex());
 	}
 
+	/**
+	 * Obtain the {@link ExpressionMarker} for the current method. Will add a local class within the method that can be
+	 * referenced via {@link ExpressionMarker#enclosingMethod()}.
+	 * 
+	 * @return the {@link ExpressionMarker} for this particular method.
+	 */
+	public ExpressionMarker getExpressionMarker() {
+		return expressionMarker;
+	}
 }

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

@@ -96,6 +96,9 @@ public MethodSpec buildMethod() {
 				context.getMethod().getName(), StringUtils.collectionToCommaDelimitedString(context.getTargetMethodMetadata()
 						.getMethodArguments().values().stream().map(it -> it.type().toString()).collect(Collectors.toList())));
 		context.getTargetMethodMetadata().getMethodArguments().forEach((name, spec) -> builder.addParameter(spec));
+		if(context.getExpressionMarker().isInUse()) {
+			builder.addCode(context.getExpressionMarker().declaration());
+		}
 		builder.addCode(methodBody);
 		customizer.accept(context, builder);
 

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

@@ -0,0 +1,85 @@
+/*
+ * Copyright 2025-present 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.repository.aot.generate;
+
+import org.springframework.javapoet.CodeBlock;
+
+/**
+ * ExpressionMarker is used to add a dedicated type to AOT generated methods that can be used to determine the current
+ * method by calling {@link Class#getEnclosingMethod()} on it. This can be useful when working with expressions (eg.
+ * SpEL) that need to be evaluated in a given context.
+ * <p>
+ * {@link ExpressionMarker} is intended to be used via {@link AotQueryMethodGenerationContext} to maintain usage info,
+ * making sure the code is only added ({@link #isInUse()}) when {@link #enclosingMethod()} was called for generating
+ * code.
+ * 
+ * <pre class="code">
+ * ExpressionMarker marker = context.getExpressionMarker();
+ * CodeBlock.builder().add("evaluate($L, $S, $L)", marker.enclosingMethod(), queryString, parameters);
+ * </pre>
+ * 
+ * @author Christoph Strobl
+ * @since 4.0
+ */
+public class ExpressionMarker {
+
+	private final String typeName;
+	private boolean inUse = false;
+
+	ExpressionMarker() {
+		this("ExpressionMarker");
+	}
+
+	ExpressionMarker(String typeName) {
+		this.typeName = typeName;
+	}
+
+	/**
+	 * @return {@code class ExpressionMarker}.
+	 */
+	CodeBlock declaration() {
+		return CodeBlock.of("class $L{};\n", typeName);
+	}
+
+	/**
+	 * Calling this method sets the {@link ExpressionMarker} as {@link #isInUse() in-use}.
+	 *
+	 * @return {@code ExpressionMarker.class}.
+	 */
+	public CodeBlock marker() {
+
+		if (!inUse) {
+			inUse = true;
+		}
+		return CodeBlock.of("$L.class", typeName);
+	}
+
+	/**
+	 * Calling this method sets the {@link ExpressionMarker} as {@link #isInUse() in-use}.
+	 *
+	 * @return {@code ExpressionMarker.class.getEnclosingMethod()}
+	 */
+	public CodeBlock enclosingMethod() {
+		return CodeBlock.of("$L.getEnclosingMethod()", marker());
+	}
+
+	/**
+	 * @return if the marker is in use.
+	 */
+	public boolean isInUse() {
+		return inUse;
+	}
+}

src/test/java/org/springframework/data/repository/aot/generate/AotRepositoryMethodBuilderUnitTests.java

@@ -15,20 +15,24 @@
  */
 package org.springframework.data.repository.aot.generate;
 
-import static org.assertj.core.api.Assertions.*;
-import static org.mockito.ArgumentMatchers.*;
-import static org.mockito.Mockito.*;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.when;
 
 import example.UserRepository;
 import example.UserRepository.User;
 
 import java.lang.reflect.Method;
 import java.util.List;
+import java.util.stream.Stream;
 
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
 import org.mockito.Mockito;
-
 import org.springframework.core.ResolvableType;
 import org.springframework.data.repository.core.RepositoryInformation;
 import org.springframework.data.util.TypeInformation;
@@ -45,10 +49,12 @@ class AotRepositoryMethodBuilderUnitTests {
 
 	@BeforeEach
 	void beforeEach() {
+
 		repositoryInformation = Mockito.mock(RepositoryInformation.class);
 		methodGenerationContext = Mockito.mock(AotQueryMethodGenerationContext.class);
 
 		when(methodGenerationContext.getRepositoryInformation()).thenReturn(repositoryInformation);
+		when(methodGenerationContext.getExpressionMarker()).thenReturn(new ExpressionMarker());
 	}
 
 	@Test // GH-3279
@@ -87,4 +93,37 @@ void generatesMethodWithGenerics() throws NoSuchMethodException {
 				.containsPattern("public .*List<.*User> findByFirstnameIn\\(") //
 				.containsPattern(".*List<.*String> firstnames\\)");
 	}
+
+	@ParameterizedTest // GH-3279
+	@MethodSource(value = { "expressionMarkers" })
+	void generatesExpressionMarkerIfInUse(ExpressionMarker expressionMarker) throws NoSuchMethodException {
+
+		Method method = UserRepository.class.getMethod("findByFirstname", String.class);
+		when(methodGenerationContext.getMethod()).thenReturn(method);
+		when(methodGenerationContext.getReturnType()).thenReturn(ResolvableType.forClass(User.class));
+		doReturn(TypeInformation.of(User.class)).when(repositoryInformation).getReturnType(any());
+		doReturn(TypeInformation.of(User.class)).when(repositoryInformation).getReturnedDomainTypeInformation(any());
+		MethodMetadata methodMetadata = new MethodMetadata(repositoryInformation, method);
+		methodMetadata.addParameter(ParameterSpec.builder(String.class, "firstname").build());
+		when(methodGenerationContext.getTargetMethodMetadata()).thenReturn(methodMetadata);
+		when(methodGenerationContext.getExpressionMarker()).thenReturn(expressionMarker);
+
+		AotRepositoryMethodBuilder builder = new AotRepositoryMethodBuilder(methodGenerationContext);
+		String methodCode = builder.buildMethod().toString();
+		if (expressionMarker.isInUse()) {
+			assertThat(methodCode).contains("class ExpressionMarker{};");
+		} else {
+			assertThat(methodCode).doesNotContain("class ExpressionMarker{};");
+		}
+	}
+
+	static Stream<Arguments> expressionMarkers() {
+
+		ExpressionMarker unused = new ExpressionMarker();
+
+		ExpressionMarker used = new ExpressionMarker();
+		used.marker();
+
+		return Stream.of(Arguments.of(unused), Arguments.of(used));
+	}
 }