event maching heuristic update 3

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/HeuristicCallGraphEngine.java

@@ -63,14 +63,9 @@ public class HeuristicCallGraphEngine implements CallGraphEngine {
         String methodSuffix = "";
 
         // Extract method calls like .getType() so we can trace the base parameter
-        int dotIndex = currentParamName.indexOf('.');
+        String[] extractedEntry = extractMethodSuffix(currentParamName, methodSuffix);
-        if (dotIndex > 0 && dotIndex + 1 < currentParamName.length()) {
+        currentParamName = extractedEntry[0];
-            char nextChar = currentParamName.charAt(dotIndex + 1);
+        methodSuffix = extractedEntry[1];
-            if (Character.isLowerCase(nextChar)) {
-                methodSuffix = currentParamName.substring(dotIndex);
-                currentParamName = currentParamName.substring(0, dotIndex);
-            }
-        }
         
         // Walk backwards up the call chain
         for (int i = path.size() - 1; i > 0; i--) {
@@ -84,11 +79,9 @@ public class HeuristicCallGraphEngine implements CallGraphEngine {
                 String tracedVar = traceLocalVariable(target, currentParamName);
                 if (tracedVar != null && !tracedVar.equals(currentParamName)) {
                     // Extract method calls like .getType() from the traced variable
-                    int dotIdx = tracedVar.indexOf('.');
+                    String[] extractedTraced = extractMethodSuffix(tracedVar, methodSuffix);
-                    if (dotIdx > 0 && dotIdx + 1 < tracedVar.length() && Character.isLowerCase(tracedVar.charAt(dotIdx + 1))) {
+                    tracedVar = extractedTraced[0];
-                        methodSuffix = tracedVar.substring(dotIdx) + methodSuffix;
+                    methodSuffix = extractedTraced[1];
-                        tracedVar = tracedVar.substring(0, dotIdx);
-                    }
                     currentParamName = tracedVar;
                     resolvedValue = tracedVar + methodSuffix;
                     paramIndex = getParameterIndex(target, currentParamName);
@@ -108,11 +101,9 @@ public class HeuristicCallGraphEngine implements CallGraphEngine {
                             String arg = edge.getArguments().get(paramIndex);
                             if (arg != null) {
                                 // If the argument passed has a method call, extract it
-                                int dotIdx = arg.indexOf('.');
+                                String[] extractedArg = extractMethodSuffix(arg, methodSuffix);
-                                if (dotIdx > 0 && dotIdx + 1 < arg.length() && Character.isLowerCase(arg.charAt(dotIdx + 1))) {
+                                arg = extractedArg[0];
-                                    methodSuffix = arg.substring(dotIdx) + methodSuffix;
+                                methodSuffix = extractedArg[1];
-                                    arg = arg.substring(0, dotIdx);
-                                }
                                 currentParamName = arg;
                                 resolvedValue = arg + methodSuffix;
                                 found = true;
@@ -154,11 +145,9 @@ public class HeuristicCallGraphEngine implements CallGraphEngine {
 
             String tracedVar = traceLocalVariable(entryMethod, currentParamName);
             if (tracedVar != null && !tracedVar.equals(currentParamName)) {
-                int dotIdx = tracedVar.indexOf('.');
+                String[] extractedFinalTraced = extractMethodSuffix(tracedVar, methodSuffix);
-                if (dotIdx > 0 && dotIdx + 1 < tracedVar.length() && Character.isLowerCase(tracedVar.charAt(dotIdx + 1))) {
+                tracedVar = extractedFinalTraced[0];
-                    methodSuffix = tracedVar.substring(dotIdx) + methodSuffix;
+                methodSuffix = extractedFinalTraced[1];
-                    tracedVar = tracedVar.substring(0, dotIdx);
-                }
                 currentParamName = tracedVar;
                 resolvedValue = tracedVar + methodSuffix;
             }
@@ -166,6 +155,24 @@ public class HeuristicCallGraphEngine implements CallGraphEngine {
         
         List<String> polymorphicEvents = new ArrayList<>();
         
+        if (resolvedValue.startsWith("ENUM_SET:")) {
+            for (String eVal : resolvedValue.substring(9).split(",")) {
+                String parsed = parseEnumSetElement(eVal);
+                if (!polymorphicEvents.contains(parsed)) polymorphicEvents.add(parsed);
+            }
+            return TriggerPoint.builder()
+                    .event(tp.getEvent())
+                    .className(tp.getClassName())
+                    .methodName(tp.getMethodName())
+                    .sourceFile(tp.getSourceFile())
+                    .sourceModule(tp.getSourceModule())
+                    .stateMachineId(tp.getStateMachineId())
+                    .sourceState(tp.getSourceState())
+                    .lineNumber(tp.getLineNumber())
+                    .polymorphicEvents(polymorphicEvents)
+                    .build();
+        }
+        
         // Parse resolvedValue using JDT to robustly handle complex expressions
         org.eclipse.jdt.core.dom.ASTParser exprParser = org.eclipse.jdt.core.dom.ASTParser.newParser(org.eclipse.jdt.core.dom.AST.getJLSLatest());
         exprParser.setSource(resolvedValue.toCharArray());
@@ -494,7 +501,7 @@ public class HeuristicCallGraphEngine implements CallGraphEngine {
                                 if (val != null) {
                                     if (val.startsWith("ENUM_SET:")) {
                                         for (String eVal : val.substring(9).split(",")) {
-                                            String parsed = eVal.substring(eVal.lastIndexOf('.') + 1);
+                                            String parsed = parseEnumSetElement(eVal);
                                             if (!constants.contains(parsed)) constants.add(parsed);
                                         }
                                     } else {
@@ -777,6 +784,25 @@ public class HeuristicCallGraphEngine implements CallGraphEngine {
         return (MethodDeclaration) parent;
     }
 
+    private ASTNode findStatement(ASTNode node) {
+        while (node != null && !(node instanceof Statement)) {
+            node = node.getParent();
+        }
+        return node;
+    }
+
+    private String[] extractMethodSuffix(String paramName, String currentSuffix) {
+        int dotIndex = paramName.indexOf('.');
+        if (dotIndex > 0 && dotIndex + 1 < paramName.length() && Character.isLowerCase(paramName.charAt(dotIndex + 1))) {
+            return new String[] { paramName.substring(0, dotIndex), paramName.substring(dotIndex) + currentSuffix };
+        }
+        return new String[] { paramName, currentSuffix };
+    }
+
+    private String parseEnumSetElement(String eVal) {
+        return eVal.contains(".") ? eVal.substring(eVal.lastIndexOf('.', eVal.lastIndexOf('.') - 1) + 1) : eVal;
+    }
+
     private Map<String, List<CallEdge>> buildCallGraph() {
         graph = new HashMap<>();
         for (CompilationUnit cu : context.getCompilationUnits()) {
@@ -1256,7 +1282,7 @@ public class HeuristicCallGraphEngine implements CallGraphEngine {
                     if (val != null) {
                         if (val.startsWith("ENUM_SET:")) {
                             for (String eVal : val.substring(9).split(",")) {
-                                results.add(eVal.substring(eVal.lastIndexOf('.') + 1));
+                                results.add(parseEnumSetElement(eVal));
                             }
                         } else {
                             results.add(val);

state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/HeuristicCallGraphEngineTest.java

@@ -79,6 +79,124 @@ class HeuristicCallGraphEngineTest {
                 .containsExactlyInAnyOrder("OrderEvents.A8", "OrderEvents.A133");
     }
 
+    @Test
+    void shouldResolveOldStyleSwitchStatementPolymorphism(@TempDir Path tempDir) throws IOException {
+        String source = """
+                package com.example;
+                public class OrderService {
+                    public void updateOrderState(String vv) {
+                        StateMachine sm = new StateMachine();
+                        sm.sendEvent(getEventFromOldSwitch(vv));
+                    }
+                    
+                    private OrderEvents getEventFromOldSwitch(String q) {
+                        switch (q) {
+                            case "a": return OrderEvents.A8;
+                            case "b": 
+                            case "c": return OrderEvents.A133;
+                            default: return OrderEvents.PAY;
+                        }
+                    }
+                }
+                
+                enum OrderEvents { A8, A133, PAY }
+                
+                class StateMachine {
+                    public void sendEvent(OrderEvents event) {}
+                }
+                """;
+        Files.writeString(tempDir.resolve("OrderService.java"), source);
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+        HeuristicCallGraphEngine builder = new HeuristicCallGraphEngine(context);
+
+        EntryPoint entryPoint = EntryPoint.builder().className("com.example.OrderService").methodName("updateOrderState").build();
+        TriggerPoint trigger = TriggerPoint.builder().className("com.example.StateMachine").methodName("sendEvent").event("event").build();
+        List<CallChain> chains = builder.findChains(List.of(entryPoint), List.of(trigger));
+        assertThat(chains).hasSize(1);
+        assertThat(chains.get(0).getTriggerPoint().getPolymorphicEvents())
+                .containsExactlyInAnyOrder("OrderEvents.A8", "OrderEvents.A133", "OrderEvents.PAY");
+    }
+
+    @Test
+    void shouldResolveMultiValueSwitchExpressionPolymorphism(@TempDir Path tempDir) throws IOException {
+        String source = """
+                package com.example;
+                public class OrderService {
+                    public void updateOrderState(String vv) {
+                        StateMachine sm = new StateMachine();
+                        sm.sendEvent(getEventMultiValue(vv));
+                    }
+                    
+                    private OrderEvents getEventMultiValue(String q) {
+                        return switch (q) {
+                            case "a", "b" -> OrderEvents.A8;
+                            case "c" -> OrderEvents.A133;
+                            default -> OrderEvents.PAY;
+                        };
+                    }
+                }
+                
+                enum OrderEvents { A8, A133, PAY }
+                
+                class StateMachine {
+                    public void sendEvent(OrderEvents event) {}
+                }
+                """;
+        Files.writeString(tempDir.resolve("OrderService.java"), source);
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+        HeuristicCallGraphEngine builder = new HeuristicCallGraphEngine(context);
+
+        EntryPoint entryPoint = EntryPoint.builder().className("com.example.OrderService").methodName("updateOrderState").build();
+        TriggerPoint trigger = TriggerPoint.builder().className("com.example.StateMachine").methodName("sendEvent").event("event").build();
+        List<CallChain> chains = builder.findChains(List.of(entryPoint), List.of(trigger));
+        assertThat(chains).hasSize(1);
+        assertThat(chains.get(0).getTriggerPoint().getPolymorphicEvents())
+                .containsExactlyInAnyOrder("OrderEvents.A8", "OrderEvents.A133", "OrderEvents.PAY");
+    }
+
+    @Test
+    void shouldResolveSwitchExpressionWithBlockAndYieldPolymorphism(@TempDir Path tempDir) throws IOException {
+        String source = """
+                package com.example;
+                public class OrderService {
+                    public void updateOrderState(String vv) {
+                        StateMachine sm = new StateMachine();
+                        sm.sendEvent(getEventWithYield(vv));
+                    }
+                    
+                    private OrderEvents getEventWithYield(String q) {
+                        return switch (q) {
+                            case "a" -> {
+                                System.out.println("Processing a");
+                                yield OrderEvents.A8;
+                            }
+                            case "b" -> { yield OrderEvents.A133; }
+                            default -> OrderEvents.PAY;
+                        };
+                    }
+                }
+                
+                enum OrderEvents { A8, A133, PAY }
+                
+                class StateMachine {
+                    public void sendEvent(OrderEvents event) {}
+                }
+                """;
+        Files.writeString(tempDir.resolve("OrderService.java"), source);
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+        HeuristicCallGraphEngine builder = new HeuristicCallGraphEngine(context);
+
+        EntryPoint entryPoint = EntryPoint.builder().className("com.example.OrderService").methodName("updateOrderState").build();
+        TriggerPoint trigger = TriggerPoint.builder().className("com.example.StateMachine").methodName("sendEvent").event("event").build();
+        List<CallChain> chains = builder.findChains(List.of(entryPoint), List.of(trigger));
+        assertThat(chains).hasSize(1);
+        assertThat(chains.get(0).getTriggerPoint().getPolymorphicEvents())
+                .containsExactlyInAnyOrder("OrderEvents.A8", "OrderEvents.A133", "OrderEvents.PAY");
+    }
+
     @Test
     void shouldFindCallChainWithLambdaMethodReference(@TempDir Path tempDir) throws IOException {
         String source = """