tiny/spring-state-machine-renderer
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Extend reactive operator support for map and switchMap lambda remapping.

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Kamil Patryk Kozakowski
2026-07-13 20:56:14 +02:00
parent 1146cdbae6
commit 73754cf6ea
3 changed files with 162 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

58
state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/pipeline/ReactiveExpressionSupport.java
View File

@@ -13,11 +13,12 @@ import org.eclipse.jdt.core.dom.ASTNode;
import java.util.Set;
/**
* Extracts terminal payload literals from reactive factory chains ({@code Mono.just}, nested {@code flatMap}).
* Extracts terminal payload literals from reactive factory chains ({@code Mono.just}, nested transforms).
*/
public final class ReactiveExpressionSupport {
private static final Set<String> REACTIVE_FACTORY_METHODS = Set.of("just", "withPayload", "success");
private static final Set<String> REACTIVE_TRANSFORM_METHODS = Set.of("map", "flatMap", "switchMap", "concatMap");
private ReactiveExpressionSupport() {
}
@@ -27,7 +28,7 @@ public final class ReactiveExpressionSupport {
}
/**
* Rewrites {@code p.getEvent()} inside a {@code flatMap(p -> ...)} lambda to {@code payload.getEvent()}
* Rewrites {@code p.getEvent()} inside a reactive transform lambda to {@code payload.getEvent()}
* when {@code p} is fed by {@code Mono.just(payload)} (or a chained reactive source).
*/
public static String remapLambdaParameterGetter(
@@ -44,14 +45,14 @@ public final class ReactiveExpressionSupport {
if (lambda == null) {
return null;
}
MethodInvocation flatMapCall = findEnclosingFlatMap(lambda);
if (flatMapCall == null) {
MethodInvocation transformCall = findEnclosingReactiveTransform(lambda);
if (transformCall == null) {
return null;
}
String mappedReceiver = mapLambdaParameterToSource(
lambda,
paramName.getIdentifier(),
flatMapCall.getExpression(),
transformCall.getExpression(),
constantResolver,
context);
if (mappedReceiver == null) {
@@ -122,10 +123,11 @@ public final class ReactiveExpressionSupport {
return null;
}
private static MethodInvocation findEnclosingFlatMap(ASTNode node) {
private static MethodInvocation findEnclosingReactiveTransform(ASTNode node) {
ASTNode current = node.getParent();
while (current != null) {
if (current instanceof MethodInvocation mi && "flatMap".equals(mi.getName().getIdentifier())) {
if (current instanceof MethodInvocation mi
&& REACTIVE_TRANSFORM_METHODS.contains(mi.getName().getIdentifier())) {
return mi;
}
current = current.getParent();
@@ -143,8 +145,8 @@ public final class ReactiveExpressionSupport {
}
if (expression instanceof MethodInvocation mi) {
String methodName = mi.getName().getIdentifier();
if ("flatMap".equals(methodName) && !mi.arguments().isEmpty()) {
String fromArgument = extractFlatMapArgumentPayload(
if (REACTIVE_TRANSFORM_METHODS.contains(methodName) && !mi.arguments().isEmpty()) {
String fromArgument = extractTransformArgumentPayload(
(Expression) mi.arguments().get(0), mi.getExpression(), constantResolver, context);
if (fromArgument != null) {
return fromArgument;
@@ -167,29 +169,29 @@ public final class ReactiveExpressionSupport {
return null;
}
private static String extractFlatMapArgumentPayload(
private static String extractTransformArgumentPayload(
Expression argument,
Expression flatMapReceiver,
Expression transformReceiver,
ConstantResolver constantResolver,
CodebaseContext context) {
if (argument instanceof LambdaExpression lambda) {
Expression bodyExpression = lambdaBodyExpression(lambda);
if (bodyExpression != null) {
String lambdaPayload = extractLambdaBodyPayload(
bodyExpression, lambda, flatMapReceiver, constantResolver, context);
bodyExpression, lambda, transformReceiver, constantResolver, context);
if (lambdaPayload != null) {
return lambdaPayload;
}
return extractPayload(bodyExpression, flatMapReceiver, constantResolver, context);
return extractPayload(bodyExpression, transformReceiver, constantResolver, context);
}
}
return extractPayload(argument, flatMapReceiver, constantResolver, context);
return extractPayload(argument, transformReceiver, constantResolver, context);
}
private static String extractLambdaBodyPayload(
Expression bodyExpression,
LambdaExpression lambda,
Expression flatMapReceiver,
Expression transformReceiver,
ConstantResolver constantResolver,
CodebaseContext context) {
if (!(bodyExpression instanceof MethodInvocation factoryCall)) {
@@ -207,7 +209,7 @@ public final class ReactiveExpressionSupport {
Expression getterReceiver = getterCall.getExpression();
if (getterReceiver instanceof SimpleName paramName) {
String mappedReceiver = mapLambdaParameterToSource(
lambda, paramName.getIdentifier(), flatMapReceiver, constantResolver, context);
lambda, paramName.getIdentifier(), transformReceiver, constantResolver, context);
if (mappedReceiver != null) {
return mappedReceiver + getterSuffix;
}
@@ -221,7 +223,7 @@ public final class ReactiveExpressionSupport {
private static String mapLambdaParameterToSource(
LambdaExpression lambda,
String paramName,
Expression flatMapReceiver,
Expression transformReceiver,
ConstantResolver constantResolver,
CodebaseContext context) {
if (lambda.parameters().size() != 1) {
@@ -234,36 +236,36 @@ public final class ReactiveExpressionSupport {
if (!paramName.equals(variableDeclaration.getName().getIdentifier())) {
return null;
}
if (flatMapReceiver instanceof MethodInvocation receiverFlatMap
&& "flatMap".equals(receiverFlatMap.getName().getIdentifier())
&& !receiverFlatMap.arguments().isEmpty()
&& receiverFlatMap.arguments().get(0) instanceof LambdaExpression feederLambda
if (transformReceiver instanceof MethodInvocation receiverTransform
&& REACTIVE_TRANSFORM_METHODS.contains(receiverTransform.getName().getIdentifier())
&& !receiverTransform.arguments().isEmpty()
&& receiverTransform.arguments().get(0) instanceof LambdaExpression feederLambda
&& feederLambda != lambda) {
Expression feederBody = lambdaBodyExpression(feederLambda);
if (feederBody != null) {
String feederPayload = extractLambdaBodyPayload(
feederBody, feederLambda, receiverFlatMap.getExpression(), constantResolver, context);
feederBody, feederLambda, receiverTransform.getExpression(), constantResolver, context);
if (feederPayload != null) {
return feederPayload;
}
String extracted = extractPayload(
feederBody, receiverFlatMap.getExpression(), constantResolver, context);
feederBody, receiverTransform.getExpression(), constantResolver, context);
if (extracted != null) {
return extracted;
}
}
}
Expression source = peelJustArgument(flatMapReceiver);
Expression source = peelJustArgument(transformReceiver);
if (source instanceof MethodInvocation getterMi
&& getterMi.getExpression() instanceof SimpleName innerParamName) {
LambdaExpression outerLambda = findEnclosingLambda(lambda);
if (outerLambda != null && outerLambda != lambda) {
MethodInvocation outerFlatMap = findEnclosingFlatMap(outerLambda);
if (outerFlatMap != null) {
MethodInvocation outerTransform = findEnclosingReactiveTransform(outerLambda);
if (outerTransform != null) {
String mappedBase = mapLambdaParameterToSource(
outerLambda,
innerParamName.getIdentifier(),
outerFlatMap.getExpression(),
outerTransform.getExpression(),
constantResolver,
context);
if (mappedBase != null) {
@@ -281,7 +283,7 @@ public final class ReactiveExpressionSupport {
}
return source.toString();
}
return extractPayload(flatMapReceiver, null, constantResolver, context);
return extractPayload(transformReceiver, null, constantResolver, context);
}
private static Expression peelJustArgument(Expression expression) {