Extend reactive operator support for map and switchMap lambda remapping.

Co-authored-by: Cursor <cursoragent@cursor.com>

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<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) {

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<CallChain> 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<CallChain> chains = engine.findChains(List.of(entry), List.of(trigger));
+            assertThat(chains).hasSize(1);
+            assertThat(chains.get(0).getTriggerPoint().getPolymorphicEvents())
+                    .containsExactly("OrderEvent.SHIP");
+        }
     }
 }

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()");
+    }
 }