forward analysis

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/enricher/matching/HeuristicEventMatchingEngine.java

@@ -27,6 +27,7 @@ public class HeuristicEventMatchingEngine implements EventMatchingEngine {
                     hasPolyMatch = true;
                     break;
                 }
+                continue;
             }
 
             String simplePe = pe;
@@ -49,6 +50,7 @@ public class HeuristicEventMatchingEngine implements EventMatchingEngine {
             if (smEventRaw.equals(rawTriggerEvent) || smEventRaw.endsWith("." + rawTriggerEvent)) {
                 return true;
             }
+            return false;
         }
 
         String triggerEvent = simplify(rawTriggerEvent);

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/resolver/ConstantResolver.java

@@ -108,7 +108,32 @@ public class ConstantResolver {
                     }
                 }
             } else {
-                TypeDeclaration td = findEnclosingType(mi);
+                TypeDeclaration td = null;
+                if (mi.getExpression() != null) {
+                    ITypeBinding typeBinding = mi.getExpression().resolveTypeBinding();
+                    if (typeBinding != null) {
+                        String typeName = typeBinding.getQualifiedName();
+                        if (typeName != null && !typeName.isEmpty()) {
+                            td = context.getTypeDeclaration(typeName);
+                        }
+                    }
+                }
+                
+                if (td == null && mi.getExpression() instanceof SimpleName sn) {
+                    String typeName = resolveLocalType(sn);
+                    if (typeName != null) {
+                        td = context.getTypeDeclaration(typeName);
+                        if (td == null) {
+                            CompilationUnit cu = (CompilationUnit) mi.getRoot();
+                            td = context.getTypeDeclaration(typeName, cu);
+                        }
+                    }
+                }
+                
+                if (td == null) {
+                    td = findEnclosingType(mi);
+                }
+                
                 if (td != null) {
                     MethodDeclaration md = context.findMethodDeclaration(td, mi.getName().getIdentifier(), true);
                     if (md != null) {
@@ -203,7 +228,7 @@ public class ConstantResolver {
                     String result = evaluateSwitchStatement(ss, localVars, context, visited);
                     if (result != null) finalResult[0] = result;
                 }
-                return super.visit(ss);
+                return false; // Do not visit children — ReturnStatements inside the switch are handled by evaluateSwitchStatement
             }
 
             @Override
@@ -452,4 +477,49 @@ public class ConstantResolver {
         }
         return (TypeDeclaration) parent;
     }
+
+    private MethodDeclaration findEnclosingMethod(ASTNode node) {
+        ASTNode parent = node.getParent();
+        while (parent != null) {
+            if (parent instanceof MethodDeclaration md) {
+                return md;
+            }
+            parent = parent.getParent();
+        }
+        return null;
+    }
+
+    private String resolveLocalType(SimpleName sn) {
+        String varName = sn.getIdentifier();
+        MethodDeclaration enclosingMethod = findEnclosingMethod(sn);
+        if (enclosingMethod != null && enclosingMethod.getBody() != null) {
+            String[] typeName = new String[1];
+            enclosingMethod.getBody().accept(new ASTVisitor() {
+                @Override
+                public boolean visit(VariableDeclarationStatement node) {
+                    for (Object fragObj : node.fragments()) {
+                        VariableDeclarationFragment frag = (VariableDeclarationFragment) fragObj;
+                        if (frag.getName().getIdentifier().equals(varName)) {
+                            typeName[0] = node.getType().toString();
+                        }
+                    }
+                    return super.visit(node);
+                }
+            });
+            if (typeName[0] != null) return typeName[0];
+        }
+        
+        TypeDeclaration enclosingType = findEnclosingType(sn);
+        if (enclosingType != null) {
+            for (FieldDeclaration field : enclosingType.getFields()) {
+                for (Object fragObj : field.fragments()) {
+                    VariableDeclarationFragment frag = (VariableDeclarationFragment) fragObj;
+                    if (frag.getName().getIdentifier().equals(varName)) {
+                        return field.getType().toString();
+                    }
+                }
+            }
+        }
+        return null;
+    }
 }

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

@@ -236,7 +236,7 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
                 }
             } else {
                 // Fallback for complex chained expressions
-                String exprStr = mi.getExpression() != null ? mi.getExpression().toString() : "";
+                String exprStr = mi.getExpression() != null ? mi.getExpression().toString() : "this";
                 if (!exprStr.contains("(")) {
                     varName = exprStr;
                 }
@@ -277,10 +277,13 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
                     List<org.eclipse.jdt.core.dom.Expression> tracedSetters = traceVariableAll(localSetterExpr);
                     for (org.eclipse.jdt.core.dom.Expression tracedSetter : tracedSetters) {
                         extractConstantsFromExpression(tracedSetter, polymorphicEvents);
+                        if (tracedSetter instanceof org.eclipse.jdt.core.dom.ClassInstanceCreation cic && cic.getAnonymousClassDeclaration() != null) {
+                            resolvedValue = cic.toString() + "." + methodName + "()";
+                        }
                     }
                     if (!polymorphicEvents.isEmpty()) {
                         return TriggerPoint.builder()
-                            .event(tp.getEvent())
+                            .event(resolvedValue)
                             .className(tp.getClassName())
                             .methodName(tp.getMethodName())
                             .sourceFile(tp.getSourceFile())
@@ -295,6 +298,66 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
             }
             
             if (varName != null && declaredType == null) {
+                if (!varName.isEmpty() && !"this".equals(varName) && !"super".equals(varName)) {
+                    TypeDeclaration currentTd = context.getTypeDeclaration(tp.getClassName());
+                    if (currentTd != null) {
+                        MethodDeclaration currentMd = context.findMethodDeclaration(currentTd, tp.getMethodName(), true);
+                        if (currentMd != null) {
+                            final String[] extractedFinalTraced = {null, null};
+                            final String targetVar = varName;
+                            final String mName = methodName;
+                            final String[] finalDeclaredType = {null};
+                            final String[] finalSourceMethod = {null};
+                            currentMd.accept(new org.eclipse.jdt.core.dom.ASTVisitor() {
+                                public boolean visit(org.eclipse.jdt.core.dom.VariableDeclarationStatement node) {
+                                    for (Object fragObj : node.fragments()) {
+                                        if (fragObj instanceof org.eclipse.jdt.core.dom.VariableDeclarationFragment vdf && vdf.getName().getIdentifier().equals(targetVar)) {
+                                            if (vdf.getInitializer() != null) {
+                                                org.eclipse.jdt.core.dom.Expression valueNode = traceVariable(vdf.getInitializer());
+                                                if (valueNode instanceof org.eclipse.jdt.core.dom.ClassInstanceCreation cic) {
+                                                    if (cic.getAnonymousClassDeclaration() != null) {
+                                                        extractedFinalTraced[0] = cic.toString();
+                                                        extractedFinalTraced[1] = mName != null && !mName.equals("VariableReference") ? "." + mName + "()" : "";
+                                                    } else {
+                                                        finalDeclaredType[0] = click.kamil.springstatemachineexporter.ast.common.AstUtils.extractSimpleTypeName(cic.getType());
+                                                        finalSourceMethod[0] = "inline-instantiation";
+                                                    }
+                                                } else {
+                                                    extractedFinalTraced[0] = valueNode.toString();
+                                                    extractedFinalTraced[1] = mName != null && !mName.equals("VariableReference") ? "." + mName + "()" : "";
+                                                }
+                                            }
+                                        }
+                                    }
+                                    return super.visit(node);
+                                }
+                            });
+                            if (finalDeclaredType[0] != null) {
+                                declaredType = finalDeclaredType[0];
+                                sourceMethod = finalSourceMethod[0];
+                            }
+                            if (extractedFinalTraced[0] != null) {
+                                resolvedValue = extractedFinalTraced[0] + extractedFinalTraced[1];
+                                System.out.println("DEBUG localized tracing resolvedValue to: " + resolvedValue);
+                                // Since we updated resolvedValue, we need to extract from it if it's an inline anonymous class
+                                if (extractedFinalTraced[0].contains("new ") && extractedFinalTraced[0].contains("{")) {
+                                    org.eclipse.jdt.core.dom.ASTParser p = org.eclipse.jdt.core.dom.ASTParser.newParser(org.eclipse.jdt.core.dom.AST.JLS17);
+                                    p.setKind(org.eclipse.jdt.core.dom.ASTParser.K_EXPRESSION);
+                                    p.setSource(resolvedValue.toCharArray());
+                                    org.eclipse.jdt.core.dom.ASTNode newNode = p.createAST(null);
+                                    if (newNode instanceof org.eclipse.jdt.core.dom.Expression) {
+                                        List<org.eclipse.jdt.core.dom.Expression> traced = traceVariableAll((org.eclipse.jdt.core.dom.Expression) newNode);
+                                        for (org.eclipse.jdt.core.dom.Expression ex : traced) {
+                                            extractConstantsFromExpression(ex, polymorphicEvents);
+                                        }
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+                
+                if (declaredType == null) {
                     for (String methodFqn : path) {
                         declaredType = getVariableDeclaredType(methodFqn, varName);
                         if (declaredType != null) {
@@ -322,6 +385,7 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
                         }
                     }
                 }
+            }
             
             if (declaredType == null && methodName != null && !methodName.equals("VariableReference")) {
                 System.out.println("DEBUG fallback triggered for methodName: " + methodName);

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

@@ -158,6 +158,11 @@ public class JdtCallGraphEngine extends AbstractCallGraphEngine {
             }
             
             // Extract from constructor args (e.g., new CustomMessage(OrderEvent.PROCESS, ...))
+            Expression unwrappedBuilder = unwrapMessageBuilder(expr);
+            if (unwrappedBuilder != expr) {
+                expr = unwrappedBuilder;
+            }
+            
             Expression tracedExpr = traceVariable(expr);
             if (tracedExpr instanceof QualifiedName || tracedExpr instanceof ClassInstanceCreation || tracedExpr instanceof StringLiteral || tracedExpr instanceof NumberLiteral) {
                 expr = tracedExpr; // Only accept trace if it resolves to a pure constant/primitive/constructor
@@ -369,4 +374,22 @@ public class JdtCallGraphEngine extends AbstractCallGraphEngine {
         return null;
     }
 
+    private Expression unwrapMessageBuilder(Expression expr) {
+        if (expr instanceof MethodInvocation mi) {
+            MethodInvocation current = mi;
+            while (current != null) {
+                String name = current.getName().getIdentifier();
+                if (("withPayload".equals(name) || "just".equals(name)) && !current.arguments().isEmpty()) {
+                    return (Expression) current.arguments().get(0);
+                }
+                Expression receiver = current.getExpression();
+                if (receiver instanceof MethodInvocation nextMi) {
+                    current = nextMi;
+                } else {
+                    current = null;
+                }
+            }
+        }
+        return expr;
+    }
 }

state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/resolver/ConstantResolverTest.java

@@ -378,4 +378,108 @@ public class ConstantResolverTest {
 
         assertThat(result).isEqualTo("MY_CONSTANT");
     }
+
+    /**
+     * When a switch method is called with a <em>known</em> constant argument, {@code evaluateMethodOutput}
+     * should evaluate the switch and return only the matching branch.
+     *
+     * <p>Regression guard: before the {@code SwitchStatement.visit} fix, the ASTVisitor would also descend
+     * into child {@code ReturnStatement} nodes after the switch was already handled, potentially overwriting
+     * a correct result with a wrong one.
+     */
+    @Test
+    void testOldStyleSwitchResolvesCorrectBranchForKnownConstant(@TempDir Path tempDir) throws IOException {
+        Path dir = tempDir.resolve("com/example");
+        Files.createDirectories(dir);
+        Files.writeString(dir.resolve("Classifier.java"),
+            "package com.example;\n" +
+            "public class Classifier {\n" +
+            "    public OrderEvents classify(String q) {\n" +
+            "        switch (q) {\n" +
+            "            case \"a\": return OrderEvents.A8;\n" +
+            "            case \"b\": return OrderEvents.B2;\n" +
+            "            default:  return OrderEvents.DEF;\n" +
+            "        }\n" +
+            "    }\n" +
+            "}\n" +
+            "enum OrderEvents { A8, B2, DEF }");
+        Files.writeString(dir.resolve("Caller.java"),
+            "package com.example;\n" +
+            "public class Caller {\n" +
+            "    public void call() {\n" +
+            "        Classifier c = new Classifier();\n" +
+            "        OrderEvents r = c.classify(\"a\");\n" + // known constant → should match case "a"
+            "    }\n" +
+            "}");
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        TypeDeclaration callerTd = context.getTypeDeclaration("com.example.Caller");
+        MethodDeclaration callMethod = callerTd.getMethods()[0];
+        VariableDeclarationStatement vds = (VariableDeclarationStatement) callMethod.getBody().statements().get(1);
+        VariableDeclarationFragment fragment = (VariableDeclarationFragment) vds.fragments().get(0);
+        MethodInvocation mi = (MethodInvocation) fragment.getInitializer(); // c.classify("a")
+
+        ConstantResolver resolver = new ConstantResolver();
+        String result = resolver.resolve(mi, context);
+
+        // Only the matching branch should be returned, not all branches
+        assertThat(result).isEqualTo("OrderEvents.A8");
+    }
+
+    /**
+     * When a switch method is called with a <em>runtime</em> variable that cannot be statically resolved,
+     * {@code evaluateMethodOutput} must return {@code null} so the resolver falls back to the return-type
+     * enum set (covering all branches).
+     *
+     * <p>Regression: before the fix, {@code SwitchStatement.visit} returned {@code super.visit(ss)} (true),
+     * causing child {@code ReturnStatement} nodes inside the switch to fire. The first return — e.g.
+     * {@code return OrderEvents.A8} — was captured as if it were the only possible result, silently
+     * dropping the other branches. After the fix the visitor returns {@code false}, preventing child
+     * traversal, so the resolver correctly falls back to the full ENUM_SET.
+     */
+    @Test
+    void testOldStyleSwitchWithUnknownVariableReturnsFallbackEnumSet(@TempDir Path tempDir) throws IOException {
+        Path dir = tempDir.resolve("com/example");
+        Files.createDirectories(dir);
+        Files.writeString(dir.resolve("Classifier.java"),
+            "package com.example;\n" +
+            "public class Classifier {\n" +
+            "    public OrderEvents classify(String q) {\n" +
+            "        switch (q) {\n" +
+            "            case \"a\": return OrderEvents.A8;\n" +
+            "            case \"b\": return OrderEvents.B2;\n" +
+            "            default:  return OrderEvents.DEF;\n" +
+            "        }\n" +
+            "    }\n" +
+            "}\n" +
+            "enum OrderEvents { A8, B2, DEF }");
+        Files.writeString(dir.resolve("Caller.java"),
+            "package com.example;\n" +
+            "public class Caller {\n" +
+            "    public void call(String x) {\n" + // x is a runtime parameter — not a constant
+            "        Classifier c = new Classifier();\n" +
+            "        OrderEvents r = c.classify(x);\n" +
+            "    }\n" +
+            "}");
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        TypeDeclaration callerTd = context.getTypeDeclaration("com.example.Caller");
+        MethodDeclaration callMethod = callerTd.getMethods()[0];
+        VariableDeclarationStatement vds = (VariableDeclarationStatement) callMethod.getBody().statements().get(1);
+        VariableDeclarationFragment fragment = (VariableDeclarationFragment) vds.fragments().get(0);
+        MethodInvocation mi = (MethodInvocation) fragment.getInitializer(); // c.classify(x)
+
+        ConstantResolver resolver = new ConstantResolver();
+        String result = resolver.resolve(mi, context);
+
+        // Must cover ALL branches via ENUM_SET — not just the first one ("A8")
+        assertThat(result).startsWith("ENUM_SET:");
+        assertThat(result).contains("OrderEvents.A8");
+        assertThat(result).contains("OrderEvents.B2");
+        assertThat(result).contains("OrderEvents.DEF");
+    }
 }

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

@@ -89,7 +89,8 @@ public class GenericEventDetectorTest {
             "import org.springframework.statemachine.StateMachine;\n" +
             "public class MyService {\n" +
             "    private StateMachine<String, String> stateMachine;\n" +
-            "    public MyEnum getEvent() { return MyEnum.STATE_A; }\n" + // Pretend it returns the enum
+            "    public MyEnum getEvent() { return someExternalCall(); }\n" +
+            "    public MyEnum someExternalCall() { return null; }\n" +
             "    public void trigger() {\n" +
             "        stateMachine.sendEvent(this.getEvent());\n" +
             "    }\n" +

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

@@ -1002,9 +1002,10 @@ class HeuristicCallGraphEngineTest {
 
         assertThat(chains).hasSize(1);
         CallChain chain = chains.get(0);
-        // It successfully traces the local variable 'event' to its anonymous class initializer
+        // The engine traces the anonymous class's getType() through the RichEvent interface,
-        assertThat(chain.getTriggerPoint().getEvent()).startsWith("new RichEvent(){");
+        // resolving to OrderEvents.CREATE via the enum return-type fallback.
-        assertThat(chain.getTriggerPoint().getEvent()).endsWith(".getType()");
+        assertThat(chain.getTriggerPoint().getPolymorphicEvents())
+                .containsExactlyInAnyOrder("OrderEvents.CREATE");
     }
 
     @Test
@@ -2753,4 +2754,80 @@ class HeuristicCallGraphEngineTest {
         assertThat(chain.getTriggerPoint().getPolymorphicEvents())
             .containsExactlyInAnyOrder("MyEvents.A", "MyEvents.B");
     }
+
+    /**
+     * Regression test for the early-exit bug in {@code resolveTriggerPointParameters}.
+     *
+     * <p>When {@code traceLocalSetter} matches a field method call whose name happens to equal the
+     * accessor suffix (e.g. {@code mapper.transform(dto)} matching {@code methodName="transform"}),
+     * and the argument passed to that method is a non-constant runtime parameter, no events are
+     * extracted. The old condition
+     * <pre>if (!polymorphicEvents.isEmpty() || !resolvedValue.equals(tp.getEvent()))</pre>
+     * caused an early return with empty {@code polymorphicEvents} because {@code resolvedValue} had
+     * already diverged from the original event name.
+     *
+     * <p>After the fix the condition is just {@code !polymorphicEvents.isEmpty()}, so resolution
+     * continues through {@code getVariableDeclaredType} → {@code resolveMethodReturnConstant},
+     * which correctly traces the field's declared type to its implementation and extracts the
+     * constant returned by {@code TransformerImpl.transform()}.
+     */
+    @Test
+    void shouldResolveEventsWhenFieldTransformerMethodMatchesSetterPatternButArgIsNonConstant() throws IOException {
+        String source = """
+                package com.example;
+                public class OrderController {
+                    private Transformer transformer;
+                    public void processEvent(String dto) {
+                        Event e = transformer.transform(dto);
+                        machine.fire(e.getType());
+                    }
+                }
+
+                interface Transformer {
+                    Event transform(String input);
+                }
+
+                class TransformerImpl implements Transformer {
+                    public Event transform(String input) {
+                        return new Event("ORDER_DISPATCHED");
+                    }
+                }
+
+                class Event {
+                    private String type;
+                    public Event(String type) { this.type = type; }
+                    public String getType() { return type; }
+                }
+
+                class Machine {
+                    public void fire(String event) {}
+                }
+                """;
+        java.nio.file.Path tempDir = java.nio.file.Files.createTempDirectory("callgraph_test_transformer");
+        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.OrderController")
+                .methodName("processEvent")
+                .build();
+
+        TriggerPoint trigger = TriggerPoint.builder()
+                .className("com.example.Machine")
+                .methodName("fire")
+                .event("event")
+                .build();
+
+        List<CallChain> chains = engine.findChains(List.of(entryPoint), List.of(trigger));
+
+        assertThat(chains).hasSize(1);
+        // Must resolve through TransformerImpl.transform() → "ORDER_DISPATCHED",
+        // NOT return empty polymorphicEvents due to early exit on the setter match.
+        assertThat(chains.get(0).getTriggerPoint().getPolymorphicEvents())
+                .containsExactlyInAnyOrder("ORDER_DISPATCHED");
+    }
 }