diff --git a/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/resolver/ConstantResolver.java b/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/resolver/ConstantResolver.java index a0e25d9..45cf5c6 100644 --- a/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/resolver/ConstantResolver.java +++ b/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/resolver/ConstantResolver.java @@ -14,6 +14,14 @@ public class ConstantResolver { return resolveInternal(expr, context, new HashSet<>()); } + public String evaluateSwitchWithParams(org.eclipse.jdt.core.dom.SwitchExpression se, java.util.Map paramValues, CodebaseContext context) { + return evaluateSwitchExpression(se, paramValues, context, new HashSet<>()); + } + + public String evaluateSwitchWithParams(org.eclipse.jdt.core.dom.SwitchStatement ss, java.util.Map paramValues, CodebaseContext context) { + return evaluateSwitchStatement(ss, paramValues, context, new HashSet<>()); + } + private String resolveInternal(Expression expr, CodebaseContext context, Set visited) { if (expr == null) return null; 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 b38c3db..b0ebab1 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 @@ -78,7 +78,8 @@ import java.util.*; int paramIndex = getParameterIndex(target, currentParamName); if (paramIndex < 0) { // Not a parameter. Maybe it's a local variable initialized from a parameter or method call? - String tracedVar = traceLocalVariable(target, currentParamName); + Map parameterValues = buildParameterValuesMap(caller, target, callGraph, path, i); + String tracedVar = traceLocalVariable(target, currentParamName, parameterValues); if (tracedVar != null && !tracedVar.equals(currentParamName)) { // Extract method calls like .getType() from the traced variable String[] extractedTraced = extractMethodSuffix(tracedVar, methodSuffix); @@ -1511,6 +1512,149 @@ import java.util.*; return null; } + protected String traceLocalVariable(String methodFqn, String varName, java.util.Map parameterValues) { + String result = traceLocalVariable(methodFqn, varName); + if (result != null && parameterValues != null && !parameterValues.isEmpty()) { + TypeDeclaration td = context.getTypeDeclaration(methodFqn.substring(0, methodFqn.lastIndexOf('.'))); + if (td != null) { + String methodName = methodFqn.substring(methodFqn.lastIndexOf('.') + 1); + MethodDeclaration md = context.findMethodDeclaration(td, methodName, true); + if (md != null && md.getBody() != null) { + final org.eclipse.jdt.core.dom.ASTNode[] switchNode = new org.eclipse.jdt.core.dom.ASTNode[1]; + md.getBody().accept(new ASTVisitor() { + @Override + public boolean visit(VariableDeclarationFragment node) { + if (node.getName().getIdentifier().equals(varName) && node.getInitializer() != null) { + Expression init = node.getInitializer(); + if (init.getNodeType() == ASTNode.SWITCH_EXPRESSION) { + switchNode[0] = init; + return false; + } + if (init.getNodeType() == ASTNode.SWITCH_STATEMENT) { + switchNode[0] = init; + return false; + } + } + return super.visit(node); + } + @Override + public boolean visit(Assignment node) { + if (node.getLeftHandSide() instanceof SimpleName asn && asn.getIdentifier().equals(varName)) { + Expression rhs = node.getRightHandSide(); + if (rhs.getNodeType() == ASTNode.SWITCH_EXPRESSION) { + switchNode[0] = rhs; + return false; + } + if (rhs.getNodeType() == ASTNode.SWITCH_STATEMENT) { + switchNode[0] = rhs; + return false; + } + } + return super.visit(node); + } + }); + if (switchNode[0] != null) { + if (switchNode[0] instanceof SwitchExpression se) { + String evaluated = constantResolver.evaluateSwitchWithParams(se, parameterValues, context); + if (evaluated != null) return evaluated; + } else if (switchNode[0] instanceof SwitchStatement ss) { + String evaluated = constantResolver.evaluateSwitchWithParams(ss, parameterValues, context); + if (evaluated != null) return evaluated; + } + } + } + } + } + return result; + } + + protected java.util.Map buildParameterValuesMap(String caller, String target, Map> callGraph, List path, int pathIndex) { + java.util.Map paramValues = new java.util.HashMap<>(); + List edges = callGraph.get(caller); + if (edges == null) { + return paramValues; + } + + for (CallEdge edge : edges) { + if (edge.getTargetMethod().equals(target) || isHeuristicMatch(edge.getTargetMethod(), target)) { + List args = edge.getArguments(); + if (args == null || args.isEmpty()) break; + + int lastDot = target.lastIndexOf('.'); + if (lastDot < 0) break; + String className = target.substring(0, lastDot); + String methodName = target.substring(lastDot + 1); + + TypeDeclaration td = context.getTypeDeclaration(className); + if (td == null) break; + + MethodDeclaration md = context.findMethodDeclaration(td, methodName, true); + if (md == null) break; + + List paramNames = new java.util.ArrayList<>(); + for (Object paramObj : md.parameters()) { + org.eclipse.jdt.core.dom.SingleVariableDeclaration param = (org.eclipse.jdt.core.dom.SingleVariableDeclaration) paramObj; + paramNames.add(param.getName().getIdentifier()); + } + + int minSize = java.lang.Math.min(paramNames.size(), args.size()); + for (int j = 0; j < minSize; j++) { + String argValue = args.get(j); + String resolvedArgValue = resolveArgumentValue(argValue, caller, path, pathIndex, callGraph); + paramValues.put(paramNames.get(j), resolvedArgValue); + } + break; + } + } + return paramValues; + } + + protected String resolveArgumentValue(String argValue, String callerMethod, List path, int pathIndex, Map> callGraph) { + if (argValue == null) return null; + if (argValue.contains(".") || argValue.startsWith("new ") || argValue.matches(".*[0-9].*")) { + return argValue; + } + + int callerIndex = path.indexOf(callerMethod); + if (callerIndex <= 0) return argValue; + + String prevCaller = path.get(callerIndex - 1); + List prevEdges = callGraph.get(prevCaller); + if (prevEdges == null) return argValue; + + for (CallEdge edge : prevEdges) { + if (edge.getTargetMethod().equals(callerMethod) || isHeuristicMatch(edge.getTargetMethod(), callerMethod)) { + List prevArgs = edge.getArguments(); + if (prevArgs == null || prevArgs.isEmpty()) break; + + int lastDot = callerMethod.lastIndexOf('.'); + if (lastDot < 0) break; + String className = callerMethod.substring(0, lastDot); + String methodName = callerMethod.substring(lastDot + 1); + + TypeDeclaration td = context.getTypeDeclaration(className); + if (td == null) break; + + MethodDeclaration md = context.findMethodDeclaration(td, methodName, true); + if (md == null) break; + + List paramNames = new java.util.ArrayList<>(); + for (Object paramObj : md.parameters()) { + org.eclipse.jdt.core.dom.SingleVariableDeclaration param = (org.eclipse.jdt.core.dom.SingleVariableDeclaration) paramObj; + paramNames.add(param.getName().getIdentifier()); + } + + int paramIdx = paramNames.indexOf(argValue); + if (paramIdx >= 0 && paramIdx < prevArgs.size()) { + String prevArg = prevArgs.get(paramIdx); + return resolveArgumentValue(prevArg, prevCaller, path, pathIndex, callGraph); + } + break; + } + } + return argValue; + } + protected org.eclipse.jdt.core.dom.Expression traceLocalSetter(String methodFqn, String varName, String getterName) { if (methodFqn == null || !methodFqn.contains(".")) return null; String className = methodFqn.substring(0, methodFqn.lastIndexOf('.')); diff --git a/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/HeuristicCallGraphEngine.java b/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/HeuristicCallGraphEngine.java index c0d2db1..dc31169 100644 --- a/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/HeuristicCallGraphEngine.java +++ b/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/HeuristicCallGraphEngine.java @@ -562,7 +562,8 @@ public class HeuristicCallGraphEngine extends AbstractCallGraphEngine { String caller = path.get(i - 1); int paramIndex = getParameterIndex(target, currentParamName); if (paramIndex < 0) { - String tracedVar = traceLocalVariable(target, currentParamName); + java.util.Map parameterValues = buildParameterValuesMap(caller, target, callGraph, path, i); + String tracedVar = traceLocalVariable(target, currentParamName, parameterValues); if (tracedVar != null && !tracedVar.equals(currentParamName)) { String[] extractedTraced = extractMethodSuffix(tracedVar, methodSuffix); tracedVar = extractedTraced[0]; @@ -646,6 +647,21 @@ public class HeuristicCallGraphEngine extends AbstractCallGraphEngine { String entryMethod = path.get(0); int entryParamIndex = getParameterIndex(entryMethod, currentParamName); + + if (resolvedValue != null && resolvedValue.contains(".") && !resolvedValue.contains("(") && !resolvedValue.contains(" ")) { + TriggerPoint modifiedTp = TriggerPoint.builder() + .event(resolvedValue) + .className(tp.getClassName()) + .methodName(tp.getMethodName()) + .sourceFile(tp.getSourceFile()) + .sourceModule(tp.getSourceModule()) + .stateMachineId(tp.getStateMachineId()) + .sourceState(tp.getSourceState()) + .lineNumber(tp.getLineNumber()) + .build(); + return super.resolveTriggerPointParameters(modifiedTp, path, callGraph); + } + if (entryParamIndex < 0) { if (methodSuffix.startsWith(".get") || methodSuffix.equals(".type") || methodSuffix.equals(".event") || methodSuffix.equals(".type()") || methodSuffix.equals(".event()")) { String localMethodName = methodSuffix.substring(1).replace("()", ""); @@ -822,4 +838,5 @@ public class HeuristicCallGraphEngine extends AbstractCallGraphEngine { }); return initExpr[0]; } + } \ No newline at end of file diff --git a/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/HeuristicCallGraphEngineExtendedTest.java b/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/HeuristicCallGraphEngineExtendedTest.java index 92a9393..378b11f 100644 --- a/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/HeuristicCallGraphEngineExtendedTest.java +++ b/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/HeuristicCallGraphEngineExtendedTest.java @@ -130,4 +130,101 @@ class HeuristicCallGraphEngineExtendedTest { assertThat(chains.get(0).getTriggerPoint().getPolymorphicEvents()) .containsExactlyInAnyOrder("TransitionEnum.STATE_Y"); } + /** + * Diagnostic: Enum tracked through an overridden method and a super() call. + * + * Pattern: + * 1. EntryPoint calls an overridden method on a subclass (CustomHandler.handleEvent) + * with an initial enum (ActionType.BEGIN_CHECKOUT). + * 2. The subclass intercepts the call, does some localized logic, and delegates + * back to the parent using super.handleEvent(ActionType). + * 3. The parent class (BaseHandler) takes that enum, transforms it via a switch + * statement to a TransitionEnum, and fires the state machine. + * + * Root cause targeted: The engine must successfully resolve the `super` keyword + * to the parent class's AST node, maintain the argument mapping across the inheritance + * boundary, and evaluate the switch to get the final TransitionEnum. It must not + * get stuck in a recursive loop or eagerly return ActionType.BEGIN_CHECKOUT. + */ + @Test + void shouldResolveMappedEnumAcrossOverriddenMethodAndSuperDelegation() throws IOException { + String source = """ + package com.example; + + public class ApiEndpoint { + private CustomHandler handler; + + public void triggerCheckout() { + // Entry point passes the initial enum + handler.handleEvent(ActionType.BEGIN_CHECKOUT); + } + } + + abstract class BaseHandler { + protected StateMachine machine; + + // The method that gets overridden, but ultimately does the work + public void handleEvent(ActionType action) { + TransitionEnum transition = switch (action) { + case BEGIN_CHECKOUT -> TransitionEnum.CHECKOUT_STARTED; // Expected + case ABORT_CHECKOUT -> TransitionEnum.CHECKOUT_CANCELLED; + }; + machine.fire(transition); + } + } + + class CustomHandler extends BaseHandler { + + @Override + public void handleEvent(ActionType action) { + // Custom logic before delegation (simulated) + if (action == null) return; + + // The trap: The engine must follow 'super' to BaseHandler + // and carry the 'action' argument with it. + super.handleEvent(action); + } + } + + class StateMachine { + public void fire(TransitionEnum event) {} + } + + enum ActionType { BEGIN_CHECKOUT, ABORT_CHECKOUT } + enum TransitionEnum { CHECKOUT_STARTED, CHECKOUT_CANCELLED } + """; + + java.nio.file.Path tempDir = java.nio.file.Files.createTempDirectory("callgraph_test_super_override"); + java.nio.file.Files.writeString(tempDir.resolve("OrderConfig.java"), source); + + CodebaseContext context = new CodebaseContext(); + context.scan(tempDir); + + HeuristicCallGraphEngine engine = new HeuristicCallGraphEngine(context); + + EntryPoint entryPoint = EntryPoint.builder() + .className("com.example.ApiEndpoint") + .methodName("triggerCheckout") + .build(); + + TriggerPoint trigger = TriggerPoint.builder() + .className("com.example.StateMachine") + .methodName("fire") + .event("event") + .build(); + + List chains = engine.findChains(List.of(entryPoint), List.of(trigger)); + + assertThat(chains).hasSize(1); + + System.out.println("SUPER DELEGATION polymorphicEvents: " + + chains.get(0).getTriggerPoint().getPolymorphicEvents()); + + // The assertion to prove the bug is fixed. + // It MUST NOT contain "ActionType.BEGIN_CHECKOUT". + // It MUST successfully trace through super.handleEvent() and resolve the switch. + assertThat(chains.get(0).getTriggerPoint().getPolymorphicEvents()) + .containsExactlyInAnyOrder("TransitionEnum.CHECKOUT_STARTED"); + } + }