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 ec5fb4f..e954337 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 @@ -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('.'); - if (dotIndex > 0 && dotIndex + 1 < currentParamName.length()) { - char nextChar = currentParamName.charAt(dotIndex + 1); - if (Character.isLowerCase(nextChar)) { - methodSuffix = currentParamName.substring(dotIndex); - currentParamName = currentParamName.substring(0, dotIndex); - } - } + String[] extractedEntry = extractMethodSuffix(currentParamName, methodSuffix); + currentParamName = extractedEntry[0]; + methodSuffix = extractedEntry[1]; // 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('.'); - if (dotIdx > 0 && dotIdx + 1 < tracedVar.length() && Character.isLowerCase(tracedVar.charAt(dotIdx + 1))) { - methodSuffix = tracedVar.substring(dotIdx) + methodSuffix; - tracedVar = tracedVar.substring(0, dotIdx); - } + String[] extractedTraced = extractMethodSuffix(tracedVar, methodSuffix); + tracedVar = extractedTraced[0]; + methodSuffix = extractedTraced[1]; 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('.'); - if (dotIdx > 0 && dotIdx + 1 < arg.length() && Character.isLowerCase(arg.charAt(dotIdx + 1))) { - methodSuffix = arg.substring(dotIdx) + methodSuffix; - arg = arg.substring(0, dotIdx); - } + String[] extractedArg = extractMethodSuffix(arg, methodSuffix); + arg = extractedArg[0]; + methodSuffix = extractedArg[1]; 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('.'); - if (dotIdx > 0 && dotIdx + 1 < tracedVar.length() && Character.isLowerCase(tracedVar.charAt(dotIdx + 1))) { - methodSuffix = tracedVar.substring(dotIdx) + methodSuffix; - tracedVar = tracedVar.substring(0, dotIdx); - } + String[] extractedFinalTraced = extractMethodSuffix(tracedVar, methodSuffix); + tracedVar = extractedFinalTraced[0]; + methodSuffix = extractedFinalTraced[1]; currentParamName = tracedVar; resolvedValue = tracedVar + methodSuffix; } @@ -166,6 +155,24 @@ public class HeuristicCallGraphEngine implements CallGraphEngine { List 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> 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); diff --git a/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/HeuristicCallGraphEngineTest.java b/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/HeuristicCallGraphEngineTest.java index 1c421a9..4376024 100644 --- a/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/HeuristicCallGraphEngineTest.java +++ b/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 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 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 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 = """