From 73754cf6ea4be4feca66159eb6533cf3a6452b66 Mon Sep 17 00:00:00 2001 From: Kamil Patryk Kozakowski Date: Mon, 13 Jul 2026 20:56:14 +0200 Subject: [PATCH] Extend reactive operator support for map and switchMap lambda remapping. Co-authored-by: Cursor --- .../pipeline/ReactiveExpressionSupport.java | 58 ++++++------ .../InheritanceAndNestedResolutionTest.java | 90 +++++++++++++++++++ .../ReactiveExpressionSupportTest.java | 42 +++++++++ 3 files changed, 162 insertions(+), 28 deletions(-) diff --git a/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/pipeline/ReactiveExpressionSupport.java b/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/pipeline/ReactiveExpressionSupport.java index 822ca0d..fa1a96e 100644 --- a/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/pipeline/ReactiveExpressionSupport.java +++ b/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/pipeline/ReactiveExpressionSupport.java @@ -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 REACTIVE_FACTORY_METHODS = Set.of("just", "withPayload", "success"); + private static final Set 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) { diff --git a/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/pipeline/InheritanceAndNestedResolutionTest.java b/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/pipeline/InheritanceAndNestedResolutionTest.java index 4c5c5f7..71cc534 100644 --- a/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/pipeline/InheritanceAndNestedResolutionTest.java +++ b/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/pipeline/InheritanceAndNestedResolutionTest.java @@ -603,5 +603,95 @@ class InheritanceAndNestedResolutionTest { assertThat(chains.get(0).getTriggerPoint().getPolymorphicEvents()) .containsExactly("OrderEvent.SHIP"); } + + @Test + void callGraphShouldResolveGetterInsideMapLambda() throws IOException { + writeJava("com/example/Api.java", """ + package com.example; + public class ApiController { + StateMachine sm; + Payload payload; + public void handle() { + Mono.just(payload).map(p -> sm.sendEvent(p.getEvent())); + } + } + class Payload { + OrderEvent event = OrderEvent.PAY; + OrderEvent getEvent() { return event; } + } + enum OrderEvent { PAY, SHIP } + class StateMachine { void sendEvent(OrderEvent event) {} } + class Mono { + static Mono just(Object o) { return new Mono(); } + Mono map(Object fn) { return this; } + } + class OrderStateMachineConfig {} + """); + scan(); + + HeuristicCallGraphEngine engine = new HeuristicCallGraphEngine(context); + EntryPoint entry = EntryPoint.builder() + .className("com.example.ApiController") + .methodName("handle") + .build(); + TriggerPoint trigger = TriggerPoint.builder() + .className("com.example.StateMachine") + .methodName("sendEvent") + .event("event") + .build(); + + List chains = engine.findChains(List.of(entry), List.of(trigger)); + assertThat(chains).hasSize(1); + assertThat(chains.get(0).getTriggerPoint().getPolymorphicEvents()) + .containsExactly("OrderEvent.PAY"); + } + + @Test + void callGraphShouldResolveGetterThroughNestedSwitchMap() throws IOException { + writeJava("com/example/Api.java", """ + package com.example; + public class ApiController { + StateMachine sm; + Payload payload; + public void handle() { + Mono.just(payload) + .switchMap(x -> Mono.just(x.getInner())) + .switchMap(p -> sm.sendEvent(p.getEvent())); + } + } + class Payload { + Inner inner = new Inner(); + Inner getInner() { return inner; } + } + class Inner { + OrderEvent event = OrderEvent.SHIP; + OrderEvent getEvent() { return event; } + } + enum OrderEvent { PAY, SHIP } + class StateMachine { void sendEvent(OrderEvent event) {} } + class Mono { + static Mono just(Object o) { return new Mono(); } + Mono switchMap(Object fn) { return this; } + } + class OrderStateMachineConfig {} + """); + scan(); + + HeuristicCallGraphEngine engine = new HeuristicCallGraphEngine(context); + EntryPoint entry = EntryPoint.builder() + .className("com.example.ApiController") + .methodName("handle") + .build(); + TriggerPoint trigger = TriggerPoint.builder() + .className("com.example.StateMachine") + .methodName("sendEvent") + .event("event") + .build(); + + List chains = engine.findChains(List.of(entry), List.of(trigger)); + assertThat(chains).hasSize(1); + assertThat(chains.get(0).getTriggerPoint().getPolymorphicEvents()) + .containsExactly("OrderEvent.SHIP"); + } } } diff --git a/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/pipeline/ReactiveExpressionSupportTest.java b/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/pipeline/ReactiveExpressionSupportTest.java index aabf930..d4ad208 100644 --- a/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/pipeline/ReactiveExpressionSupportTest.java +++ b/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/pipeline/ReactiveExpressionSupportTest.java @@ -96,4 +96,46 @@ class ReactiveExpressionSupportTest { eventExpr, context.getConstantResolver(), context); assertThat(remapped).isEqualTo("payload.getEvent()"); } + + @Test + void shouldRemapLambdaParameterGetterInsideMap(@TempDir Path tempDir) throws Exception { + String source = """ + package com.example; + public class Runner { + Payload payload; + StateMachine sm; + void run() { + Mono.just(payload).map(p -> p.getEvent()); + } + } + class Payload { String getEvent() { return "X"; } } + class StateMachine { void sendEvent(String e) {} } + class Mono { + static Mono just(Object o) { return new Mono(); } + Mono map(Object fn) { return this; } + } + """; + Path file = tempDir.resolve("Runner.java"); + Files.writeString(file, source); + + CodebaseContext context = new CodebaseContext(); + context.scan(tempDir); + + org.eclipse.jdt.core.dom.TypeDeclaration runner = context.getTypeDeclaration("com.example.Runner"); + CompilationUnit cu = (CompilationUnit) runner.getRoot(); + final MethodInvocation[] getEventCall = new MethodInvocation[1]; + cu.accept(new org.eclipse.jdt.core.dom.ASTVisitor() { + @Override + public boolean visit(MethodInvocation node) { + if ("getEvent".equals(node.getName().getIdentifier()) && getEventCall[0] == null) { + getEventCall[0] = node; + } + return super.visit(node); + } + }); + + String remapped = ReactiveExpressionSupport.remapLambdaParameterGetter( + getEventCall[0], context.getConstantResolver(), context); + assertThat(remapped).isEqualTo("payload.getEvent()"); + } }