From c7e826c0fbdc8ee81064d8e3fd3fac9590beea88 Mon Sep 17 00:00:00 2001 From: Kamil Patryk Kozakowski Date: Mon, 13 Jul 2026 20:16:58 +0200 Subject: [PATCH] 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 --- .../service/AbstractCallGraphEngine.java | 134 ++++++++++++++++-- .../ComplexMultiModuleCancelJmsTest.java | 68 +++++++++ .../service/CrossClassSupplierLambdaTest.java | 107 ++++++++++++++ .../StateMachineConfig.json | 8 +- 4 files changed, 302 insertions(+), 15 deletions(-) create mode 100644 state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/ComplexMultiModuleCancelJmsTest.java create mode 100644 state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/CrossClassSupplierLambdaTest.java diff --git 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 index c12df08..cf9ffd3 100644 --- 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 constants, List path, Map> callGraph, String scopeMethod) { + if (expr instanceof MethodInvocation mi + && "get".equals(mi.getName().getIdentifier()) + && mi.arguments().isEmpty() + && !expressionAccessClassifier.isKeyedLookup(mi, scopeMethod)) { + List callerFrameConstants = resolveSupplierConstantsFromCallerFrame(mi, path, callGraph, scopeMethod); + if (!callerFrameConstants.isEmpty()) { + for (String constant : callerFrameConstants) { + if (!constants.contains(constant)) { + constants.add(constant); + } + } + return; + } + } List 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 resolveSupplierConstantsFromCallerFrame( + MethodInvocation getCall, + List path, + Map> 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 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 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 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 path, diff --git 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 new file mode 100644 index 0000000..0cf3a93 --- /dev/null +++ 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 chains = engine.findChains(List.of(entryPoint), List.of(trigger)); + assertThat(chains).hasSize(1); + assertThat(chains.get(0).getTriggerPoint().getPolymorphicEvents()) + .contains("OrderEvent.CANCEL"); + } +} diff --git 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 new file mode 100644 index 0000000..9e325ba --- /dev/null +++ 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 stateMachine; + + @Override + public void sendMessageWithProvider(Supplier 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 { + void sendMessageWithProvider(Supplier 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 chains = engine.findChains(List.of(entryPoint), List.of(trigger)); + assertThat(chains).hasSize(1); + assertThat(chains.get(0).getTriggerPoint().getPolymorphicEvents()) + .containsExactly("OrderEvent.CANCEL"); + } +} diff --git a/state_machine_exporter/src/test/resources/golden/StateMachineConfig/StateMachineConfig.json b/state_machine_exporter/src/test/resources/golden/StateMachineConfig/StateMachineConfig.json index e887182..5847d8f 100644 --- 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" : { }