heuristics updated

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<String, String> paramValues, CodebaseContext context) {
+        return evaluateSwitchExpression(se, paramValues, context, new HashSet<>());
+    }
+
+    public String evaluateSwitchWithParams(org.eclipse.jdt.core.dom.SwitchStatement ss, java.util.Map<String, String> paramValues, CodebaseContext context) {
+        return evaluateSwitchStatement(ss, paramValues, context, new HashSet<>());
+    }
+
     private String resolveInternal(Expression expr, CodebaseContext context, Set<String> visited) {
         if (expr == null) return null;
 

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<String, String> 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<String, String> 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<String, String> buildParameterValuesMap(String caller, String target, Map<String, List<CallEdge>> callGraph, List<String> path, int pathIndex) {
+        java.util.Map<String, String> paramValues = new java.util.HashMap<>();
+        List<CallEdge> edges = callGraph.get(caller);
+        if (edges == null) {
+            return paramValues;
+        }
+
+        for (CallEdge edge : edges) {
+            if (edge.getTargetMethod().equals(target) || isHeuristicMatch(edge.getTargetMethod(), target)) {
+                List<String> 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<String> 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<String> path, int pathIndex, Map<String, List<CallEdge>> 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<CallEdge> prevEdges = callGraph.get(prevCaller);
+        if (prevEdges == null) return argValue;
+
+        for (CallEdge edge : prevEdges) {
+            if (edge.getTargetMethod().equals(callerMethod) || isHeuristicMatch(edge.getTargetMethod(), callerMethod)) {
+                List<String> 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<String> 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('.'));

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<String, String> 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];
     }
+
 }

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