Trace Supplier lambda arguments across call frames so JMS cancel chains resolve the concrete event.

When callee code calls eventProvider.get(), walk back to the caller's Supplier argument instead of stopping at the parameter name; also skip Supplier-vs-enum type checks for already-resolved lambda/enum call-site arguments under JDT bindings. Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-13 20:16:58 +02:00
parent 024458d638
commit c7e826c0fb
4 changed files with 302 additions and 15 deletions
--- a/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/AbstractCallGraphEngine.java
+++ b/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/AbstractCallGraphEngine.java
@@ -334,18 +334,20 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
                        }
                        if (expectedType != null && paramIndex < edge.getArguments().size()) {
                            String argValue = edge.getArguments().get(paramIndex);
-                            String actualType = null;
-                            if (argValue.contains(".") && !argValue.contains("(")) {
-                                String prefix = argValue.substring(0, argValue.lastIndexOf('.'));
-                                if (!prefix.equals("this") && !prefix.equals("super")) {
-                                    actualType = prefix;
+                            if (!isProvablyResolvedCallSiteArgument(argValue, expectedType)) {
+                                String actualType = null;
+                                if (argValue.contains(".") && !argValue.contains("(")) {
+                                    String prefix = argValue.substring(0, argValue.lastIndexOf('.'));
+                                    if (!prefix.equals("this") && !prefix.equals("super")) {
+                                        actualType = prefix;
+                                    }
+                                }
+                                if (actualType == null && !argValue.contains(".")) {
+                                    actualType = variableTracer.getVariableDeclaredType(caller, argValue);
+                                }
+                                if (actualType != null && !typeResolver.isTypeCompatible(actualType, expectedType)) {
+                                    continue;
                                }
-                            }
-                            if (actualType == null && !argValue.contains(".")) {
-                                actualType = variableTracer.getVariableDeclaredType(caller, argValue);
-                            }
-                            if (actualType != null && !typeResolver.isTypeCompatible(actualType, expectedType)) {
-                                continue;
                            }
                        }
                        if (paramIndex < edge.getArguments().size()) {
@@ -3047,6 +3049,20 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {

    private void addConstantsFromTracedExpression(
            Expression expr, List<String> constants, List<String> path, Map<String, List<CallEdge>> callGraph, String scopeMethod) {
+        if (expr instanceof MethodInvocation mi
+                && "get".equals(mi.getName().getIdentifier())
+                && mi.arguments().isEmpty()
+                && !expressionAccessClassifier.isKeyedLookup(mi, scopeMethod)) {
+            List<String> callerFrameConstants = resolveSupplierConstantsFromCallerFrame(mi, path, callGraph, scopeMethod);
+            if (!callerFrameConstants.isEmpty()) {
+                for (String constant : callerFrameConstants) {
+                    if (!constants.contains(constant)) {
+                        constants.add(constant);
+                    }
+                }
+                return;
+            }
+        }
        List<Expression> traced = variableTracer.traceVariableAll(expr);
        for (Expression tracedExpr : traced) {
            addConstantsFromSingleExpression(tracedExpr, constants, path, callGraph, scopeMethod);
@@ -3470,6 +3486,102 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
        return name.matches("[A-Z_][A-Z0-9_]*");
    }

+    private boolean isProvablyResolvedCallSiteArgument(String argValue, String expectedType) {
+        if (argValue == null || argValue.isBlank() || expectedType == null) {
+            return false;
+        }
+        if (!isFunctionalInterfaceType(expectedType)) {
+            return false;
+        }
+        if (looksLikeEnumConstant(argValue)) {
+            return true;
+        }
+        return isLambdaArgument(argValue);
+    }
+
+    private boolean isLambdaArgument(String argValue) {
+        String trimmed = argValue.trim();
+        return trimmed.contains("->")
+                && (trimmed.startsWith("(") || trimmed.matches("^\\w+\\s*->.*"));
+    }
+
+    private boolean isFunctionalInterfaceType(String typeName) {
+        if (typeName == null || typeName.isBlank()) {
+            return false;
+        }
+        String simple = typeName.contains(".") ? typeName.substring(typeName.lastIndexOf('.') + 1) : typeName;
+        if (simple.contains("<")) {
+            simple = simple.substring(0, simple.indexOf('<'));
+        }
+        return "Supplier".equals(simple) || "Function".equals(simple) || "Callable".equals(simple);
+    }
+
+    private List<String> resolveSupplierConstantsFromCallerFrame(
+            MethodInvocation getCall,
+            List<String> path,
+            Map<String, List<CallEdge>> callGraph,
+            String scopeMethod) {
+        if (path == null || path.size() < 2 || scopeMethod == null || getCall.getExpression() == null) {
+            return List.of();
+        }
+        if (!(getCall.getExpression() instanceof SimpleName supplierParam)) {
+            return List.of();
+        }
+        int scopeIndex = -1;
+        for (int i = path.size() - 1; i >= 0; i--) {
+            if (scopeMethod.equals(path.get(i))) {
+                scopeIndex = i;
+                break;
+            }
+        }
+        if (scopeIndex <= 0) {
+            return List.of();
+        }
+        String callee = path.get(scopeIndex);
+        String caller = path.get(scopeIndex - 1);
+        int paramIndex = typeResolver.getParameterIndex(callee, supplierParam.getIdentifier());
+        if (paramIndex < 0) {
+            return List.of();
+        }
+        List<CallEdge> edges = callGraph.get(caller);
+        if (edges == null) {
+            return List.of();
+        }
+        for (CallEdge edge : edges) {
+            if (!edge.getTargetMethod().equals(callee) && !pathFinder.isHeuristicMatch(edge.getTargetMethod(), callee)) {
+                continue;
+            }
+            if (paramIndex >= edge.getArguments().size()) {
+                continue;
+            }
+            return extractConstantsFromResolvedCallSiteArgument(edge.getArguments().get(paramIndex));
+        }
+        return List.of();
+    }
+
+    private List<String> extractConstantsFromResolvedCallSiteArgument(String argValue) {
+        if (argValue == null || argValue.isBlank()) {
+            return List.of();
+        }
+        ASTNode parsed = parseExpressionString(argValue);
+        if (!(parsed instanceof Expression expr)) {
+            if (looksLikeEnumConstant(argValue)) {
+                return List.of(argValue);
+            }
+            return List.of();
+        }
+        String resolved = resolveArgument(expr);
+        ASTNode resolvedNode = parseExpressionString(resolved);
+        List<String> constants = new ArrayList<>();
+        if (resolvedNode instanceof Expression resolvedExpr) {
+            constantExtractor.extractConstantsFromExpression(resolvedExpr, constants);
+        }
+        if (constants.isEmpty() && looksLikeEnumConstant(resolved)) {
+            constants.add(resolved);
+        }
+        return constants;
+    }
+
    private void trackKeyedMapLookupFlags(
            Expression expr,
            List<String> path,
--- a/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/ComplexMultiModuleCancelJmsTest.java
+++ b/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/ComplexMultiModuleCancelJmsTest.java
@@ -0,0 +1,68 @@
+package click.kamil.springstatemachineexporter.analysis.service;
+
+import click.kamil.springstatemachineexporter.analysis.model.CallChain;
+import click.kamil.springstatemachineexporter.analysis.model.EntryPoint;
+import click.kamil.springstatemachineexporter.analysis.model.TriggerPoint;
+import click.kamil.springstatemachineexporter.analysis.spring.InjectionPointAnalyzer;
+import click.kamil.springstatemachineexporter.analysis.spring.SpringBeanRegistry;
+import click.kamil.springstatemachineexporter.analysis.spring.SpringContextScanner;
+import click.kamil.springstatemachineexporter.analysis.spring.SpringDependencyResolver;
+import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.file.Path;
+import java.util.Collections;
+import java.util.List;
+import java.util.Set;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+class ComplexMultiModuleCancelJmsTest {
+
+    private CodebaseContext context;
+    private JdtCallGraphEngine engine;
+
+    @BeforeEach
+    void setUp() throws IOException {
+        Path projectRoot = Path.of("../state_machines/complex_multi_module_sm").toAbsolutePath().normalize();
+
+        context = new CodebaseContext();
+        context.setProjectRoot(projectRoot);
+        context.setSourcepath(List.of(
+                projectRoot.resolve("domain/src/main/java").toString(),
+                projectRoot.resolve("service/src/main/java").toString(),
+                projectRoot.resolve("web/src/main/java").toString()));
+        context.setClasspath(ToolingClasspath.currentJvmJarEntries());
+        context.setResolveBindings(true);
+        context.scan(Set.of(projectRoot), Collections.emptySet());
+
+        SpringBeanRegistry registry = new SpringBeanRegistry();
+        SpringContextScanner scanner = new SpringContextScanner(registry);
+        for (var cu : context.getCompilationUnits()) {
+            cu.accept(scanner);
+        }
+        InjectionPointAnalyzer injectionAnalyzer = new InjectionPointAnalyzer(new SpringDependencyResolver(registry));
+        engine = new JdtCallGraphEngine(context, injectionAnalyzer);
+    }
+
+    @Test
+    void shouldResolveCancelEventFromJmsSupplierLambda() {
+        EntryPoint entryPoint = EntryPoint.builder()
+                .className("click.kamil.web.JmsOrderListener")
+                .methodName("receiveCancelCommand")
+                .build();
+
+        TriggerPoint trigger = TriggerPoint.builder()
+                .className("click.kamil.service.StateMachineServiceImpl")
+                .methodName("sendMessageWithProvider")
+                .event("eventProvider")
+                .build();
+
+        List<CallChain> chains = engine.findChains(List.of(entryPoint), List.of(trigger));
+        assertThat(chains).hasSize(1);
+        assertThat(chains.get(0).getTriggerPoint().getPolymorphicEvents())
+                .contains("OrderEvent.CANCEL");
+    }
+}
--- a/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/CrossClassSupplierLambdaTest.java
+++ b/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/CrossClassSupplierLambdaTest.java
@@ -0,0 +1,107 @@
+package click.kamil.springstatemachineexporter.analysis.service;
+
+import click.kamil.springstatemachineexporter.analysis.model.CallChain;
+import click.kamil.springstatemachineexporter.analysis.model.EntryPoint;
+import click.kamil.springstatemachineexporter.analysis.model.TriggerPoint;
+import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+class CrossClassSupplierLambdaTest {
+
+    @Test
+    void shouldResolveEventThroughCrossClassSupplierLambda() throws IOException {
+        String webSource = """
+                package com.example.web;
+
+                import com.example.service.StateMachineService;
+                import com.example.domain.OrderEvent;
+                import org.springframework.stereotype.Component;
+
+                @Component
+                public class JmsOrderListener {
+                    private final StateMachineService stateMachineService;
+
+                    public JmsOrderListener(StateMachineService stateMachineService) {
+                        this.stateMachineService = stateMachineService;
+                    }
+
+                    public void receiveCancelCommand(String message) {
+                        stateMachineService.sendMessageWithProvider(() -> OrderEvent.CANCEL);
+                    }
+                }
+                """;
+
+        String serviceSource = """
+                package com.example.service;
+
+                import com.example.domain.OrderEvent;
+                import org.springframework.statemachine.StateMachine;
+                import org.springframework.messaging.support.MessageBuilder;
+                import reactor.core.publisher.Mono;
+                import java.util.function.Supplier;
+
+                public class StateMachineServiceImpl implements StateMachineService {
+                    private StateMachine<OrderEvent, OrderEvent> stateMachine;
+
+                    @Override
+                    public <T extends OrderEvent> void sendMessageWithProvider(Supplier<T> eventProvider) {
+                        T event = eventProvider.get();
+                        if (event != null) {
+                            stateMachine.sendEvent(Mono.just(MessageBuilder.withPayload(event).build())).subscribe();
+                        }
+                    }
+                }
+                """;
+
+        String ifaceSource = """
+                package com.example.service;
+
+                import com.example.domain.OrderEvent;
+                import java.util.function.Supplier;
+
+                public interface StateMachineService {
+                    <T extends OrderEvent> void sendMessageWithProvider(Supplier<T> eventProvider);
+                }
+                """;
+
+        String domainSource = """
+                package com.example.domain;
+
+                public enum OrderEvent { CANCEL, PROCESS }
+                """;
+
+        Path tempDir = Files.createTempDirectory("callgraph_cross_class_lambda");
+        Files.writeString(tempDir.resolve("JmsOrderListener.java"), webSource);
+        Files.writeString(tempDir.resolve("StateMachineServiceImpl.java"), serviceSource);
+        Files.writeString(tempDir.resolve("StateMachineService.java"), ifaceSource);
+        Files.writeString(tempDir.resolve("OrderEvent.java"), domainSource);
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        HeuristicCallGraphEngine engine = new HeuristicCallGraphEngine(context);
+
+        EntryPoint entryPoint = EntryPoint.builder()
+                .className("com.example.web.JmsOrderListener")
+                .methodName("receiveCancelCommand")
+                .build();
+
+        TriggerPoint trigger = TriggerPoint.builder()
+                .className("com.example.service.StateMachineServiceImpl")
+                .methodName("sendMessageWithProvider")
+                .event("eventProvider")
+                .build();
+
+        List<CallChain> chains = engine.findChains(List.of(entryPoint), List.of(trigger));
+        assertThat(chains).hasSize(1);
+        assertThat(chains.get(0).getTriggerPoint().getPolymorphicEvents())
+                .containsExactly("OrderEvent.CANCEL");
+    }
+}
--- a/state_machine_exporter/src/test/resources/golden/StateMachineConfig/StateMachineConfig.json
+++ b/state_machine_exporter/src/test/resources/golden/StateMachineConfig/StateMachineConfig.json
@@ -452,7 +452,7 @@
      },
      "methodChain" : [ "click.kamil.web.JmsOrderListener.receiveCancelCommand", "click.kamil.service.StateMachineServiceImpl.sendMessageWithProvider" ],
      "triggerPoint" : {
-        "event" : "eventProvider",
+        "event" : "click.kamil.domain.OrderEvent.CANCEL",
        "className" : "click.kamil.service.StateMachineServiceImpl",
        "methodName" : "sendMessageWithProvider",
        "sourceFile" : "service/src/main/java/click/kamil/service/StateMachineServiceImpl.java",
@@ -460,14 +460,14 @@
        "stateMachineId" : null,
        "sourceState" : null,
        "lineNumber" : 50,
-        "polymorphicEvents" : null,
+        "polymorphicEvents" : [ "click.kamil.domain.OrderEvent.CANCEL" ],
        "external" : false,
        "constraint" : "event != null",
-        "ambiguous" : false
+        "ambiguous" : true
      },
      "contextMachineId" : null,
      "matchedTransitions" : null,
-      "linkResolution" : "NO_MATCH"
+      "linkResolution" : "AMBIGUOUS_WIDEN"
    } ],
    "properties" : {
      "default" : { }