event maching heuristic update 4

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

@@ -17,6 +17,19 @@ public class ConstantResolver {
     private String resolveInternal(Expression expr, CodebaseContext context, Set<String> visited) {
         if (expr == null) return null;
 
+        // Unwrap common wrappers
+        if (expr instanceof CastExpression ce) {
+            return resolveInternal(ce.getExpression(), context, visited);
+        }
+        if (expr instanceof ParenthesizedExpression pe) {
+            return resolveInternal(pe.getExpression(), context, visited);
+        }
+        if (expr instanceof MethodInvocation mi && (mi.getName().getIdentifier().equals("requireNonNull") || mi.getName().getIdentifier().equals("notNull"))) {
+            if (!mi.arguments().isEmpty()) {
+                return resolveInternal((Expression) mi.arguments().get(0), context, visited);
+            }
+        }
+
         // 1. Literal?
         if (expr instanceof StringLiteral sl) {
             return sl.getLiteralValue();

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

@@ -126,7 +126,10 @@ public class HeuristicCallGraphEngine implements CallGraphEngine {
                 org.eclipse.jdt.core.dom.Expression localSetterExpr = traceLocalSetter(entryMethod, currentParamName, localMethodName);
                 if (localSetterExpr != null) {
                     List<String> setterEvents = new ArrayList<>();
-                    extractConstantsFromExpression(localSetterExpr, setterEvents);
+                    List<org.eclipse.jdt.core.dom.Expression> tracedSetters = traceVariableAll(localSetterExpr);
+                    for (org.eclipse.jdt.core.dom.Expression tracedSetter : tracedSetters) {
+                        extractConstantsFromExpression(tracedSetter, setterEvents);
+                    }
                     if (!setterEvents.isEmpty()) {
                         return TriggerPoint.builder()
                             .event(tp.getEvent())
@@ -240,7 +243,10 @@ public class HeuristicCallGraphEngine implements CallGraphEngine {
                 org.eclipse.jdt.core.dom.Expression localSetterExpr = traceLocalSetter(entryMethod, varName, methodName);
                 System.out.println("localSetterExpr = " + (localSetterExpr != null ? localSetterExpr.toString() : "null"));
                 if (localSetterExpr != null) {
-                    extractConstantsFromExpression(localSetterExpr, polymorphicEvents);
+                    List<org.eclipse.jdt.core.dom.Expression> tracedSetters = traceVariableAll(localSetterExpr);
+                    for (org.eclipse.jdt.core.dom.Expression tracedSetter : tracedSetters) {
+                        extractConstantsFromExpression(tracedSetter, polymorphicEvents);
+                    }
                     System.out.println("Extracted from setter: " + polymorphicEvents);
                     if (!polymorphicEvents.isEmpty()) {
                         return TriggerPoint.builder()
@@ -496,30 +502,33 @@ public class HeuristicCallGraphEngine implements CallGraphEngine {
                                 }
                             }
                             if (!handled) {
-                                extractConstantsFromExpression(retExpr, constants);
+                                List<Expression> tracedRetExprs = traceVariableAll(retExpr);
-                                String val = constantResolver.resolve(retExpr, context);
+                                for (Expression tracedRetExpr : tracedRetExprs) {
-                                if (val != null) {
+                                    extractConstantsFromExpression(tracedRetExpr, constants);
-                                    if (val.startsWith("ENUM_SET:")) {
+                                    String val = constantResolver.resolve(tracedRetExpr, context);
-                                        for (String eVal : val.substring(9).split(",")) {
+                                    if (val != null) {
-                                            String parsed = parseEnumSetElement(eVal);
+                                        if (val.startsWith("ENUM_SET:")) {
-                                            if (!constants.contains(parsed)) constants.add(parsed);
+                                            for (String eVal : val.substring(9).split(",")) {
+                                                String parsed = parseEnumSetElement(eVal);
+                                                if (!constants.contains(parsed)) constants.add(parsed);
+                                            }
+                                        } else {
+                                            if (!constants.contains(val)) constants.add(val);
+                                        }
+                                    } else if (tracedRetExpr instanceof org.eclipse.jdt.core.dom.SimpleName sn) {
+                                        List<String> consts = traceFieldInConstructors(td, sn.getIdentifier(), context, visited);
+                                        if (!consts.isEmpty()) {
+                                            constants.addAll(consts);
+                                        } else {
+                                            constants.add(sn.toString());
+                                        }
+                                    } else if (tracedRetExpr instanceof org.eclipse.jdt.core.dom.FieldAccess fa) {
+                                        List<String> consts = traceFieldInConstructors(td, fa.getName().getIdentifier(), context, visited);
+                                        if (!consts.isEmpty()) {
+                                            constants.addAll(consts);
+                                        } else {
+                                            constants.add(fa.getName().getIdentifier());
                                         }
-                                    } else {
-                                        if (!constants.contains(val)) constants.add(val);
-                                    }
-                                } else if (retExpr instanceof org.eclipse.jdt.core.dom.SimpleName sn) {
-                                    List<String> consts = traceFieldInConstructors(td, sn.getIdentifier(), context, visited);
-                                    if (!consts.isEmpty()) {
-                                        constants.addAll(consts);
-                                    } else {
-                                        constants.add(sn.toString());
-                                    }
-                                } else if (retExpr instanceof org.eclipse.jdt.core.dom.FieldAccess fa) {
-                                    List<String> consts = traceFieldInConstructors(td, fa.getName().getIdentifier(), context, visited);
-                                    if (!consts.isEmpty()) {
-                                        constants.addAll(consts);
-                                    } else {
-                                        constants.add(fa.getName().getIdentifier());
                                     }
                                 }
                             }
@@ -542,9 +551,11 @@ public class HeuristicCallGraphEngine implements CallGraphEngine {
         for (MethodDeclaration md : td.getMethods()) {
             if (md.getBody() != null) {
                 md.getBody().accept(new ASTVisitor() {
-                    @Override
                     public boolean visit(ReturnStatement node) {
-                        extractConstantsFromExpression(node.getExpression(), resolvedConstants);
+                        List<org.eclipse.jdt.core.dom.Expression> tracedReturns = traceVariableAll(node.getExpression());
+                        for (org.eclipse.jdt.core.dom.Expression tracedRet : tracedReturns) {
+                            extractConstantsFromExpression(tracedRet, resolvedConstants);
+                        }
                         return super.visit(node);
                     }
                 });
@@ -567,6 +578,8 @@ public class HeuristicCallGraphEngine implements CallGraphEngine {
             extractConstantsFromExpression(pe.getExpression(), constants);
         } else if (expr instanceof org.eclipse.jdt.core.dom.CastExpression ce) {
             extractConstantsFromExpression(ce.getExpression(), constants);
+        } else if (expr instanceof org.eclipse.jdt.core.dom.Assignment assignment) {
+            extractConstantsFromExpression(assignment.getRightHandSide(), constants);
         } else if (expr instanceof org.eclipse.jdt.core.dom.SwitchExpression se) {
             se.accept(new org.eclipse.jdt.core.dom.ASTVisitor() {
                 @Override
@@ -1135,33 +1148,43 @@ public class HeuristicCallGraphEngine implements CallGraphEngine {
         return (TypeDeclaration) parent;
     }
     private Expression traceVariable(Expression expr) {
+        List<Expression> all = traceVariableAll(expr);
+        return all.isEmpty() ? expr : all.get(all.size() - 1);
+    }
+
+    private List<Expression> traceVariableAll(Expression expr) {
+        List<Expression> results = new ArrayList<>();
+        System.out.println("traceVariableAll called with: " + expr);
         if (expr instanceof SimpleName sn) {
             String varName = sn.getIdentifier();
             MethodDeclaration enclosingMethod = findEnclosingMethod(expr);
+            System.out.println("traceVariableAll SimpleName=" + varName + " enclosingMethod=" + (enclosingMethod != null ? enclosingMethod.getName() : "null"));
             if (enclosingMethod != null) {
-                final Expression[] initializer = new Expression[1];
                 enclosingMethod.accept(new ASTVisitor() {
                     @Override
                     public boolean visit(VariableDeclarationFragment node) {
                         if (node.getName().getIdentifier().equals(varName) && node.getInitializer() != null) {
-                            initializer[0] = node.getInitializer();
+                            System.out.println("traceVariableAll found init: " + node.getInitializer());
+                            results.addAll(traceVariableAll(node.getInitializer()));
                         }
                         return super.visit(node);
                     }
                     @Override
                     public boolean visit(Assignment node) {
                         if (node.getLeftHandSide() instanceof SimpleName asn && asn.getIdentifier().equals(varName)) {
-                            initializer[0] = node.getRightHandSide();
+                            System.out.println("traceVariableAll found assign: " + node.getRightHandSide());
+                            results.addAll(traceVariableAll(node.getRightHandSide()));
                         }
                         return super.visit(node);
                     }
                 });
-                if (initializer[0] != null) {
+                if (!results.isEmpty()) {
-                    return traceVariable(initializer[0]);
+                    return results;
                 }
             }
         }
-        return expr;
+        results.add(expr);
+        return results;
     }
 
     private List<String> traceFieldInConstructors(TypeDeclaration td, String fieldName, CodebaseContext context, Set<String> visited) {
@@ -1173,14 +1196,17 @@ public class HeuristicCallGraphEngine implements CallGraphEngine {
                 org.eclipse.jdt.core.dom.VariableDeclarationFragment frag = (org.eclipse.jdt.core.dom.VariableDeclarationFragment) fragObj;
                 if (frag.getName().getIdentifier().equals(fieldName) && frag.getInitializer() != null) {
                     Expression right = frag.getInitializer();
-                    if (right instanceof MethodInvocation mi) {
+                    List<Expression> tracedRights = traceVariableAll(right);
-                        String calledName = mi.getName().getIdentifier();
+                    for (Expression tracedRight : tracedRights) {
-                        String targetClass = context.getFqn(td);
+                        if (tracedRight instanceof MethodInvocation mi) {
-                        CompilationUnit cu = td.getRoot() instanceof CompilationUnit ? (CompilationUnit) td.getRoot() : null;
+                            String calledName = mi.getName().getIdentifier();
-                        results.addAll(resolveMethodReturnConstant(targetClass, calledName, 0, new HashSet<>(visited), cu));
+                            String targetClass = context.getFqn(td);
-                    } else {
+                            CompilationUnit cu = td.getRoot() instanceof CompilationUnit ? (CompilationUnit) td.getRoot() : null;
-                        String val = constantResolver.resolve(right, context);
+                            results.addAll(resolveMethodReturnConstant(targetClass, calledName, 0, new HashSet<>(visited), cu));
-                        if (val != null) results.add(val);
+                        } else {
+                            String val = constantResolver.resolve(tracedRight, context);
+                            if (val != null) results.add(val);
+                        }
                     }
                 }
             }
@@ -1194,26 +1220,29 @@ public class HeuristicCallGraphEngine implements CallGraphEngine {
                     (left instanceof FieldAccess && ((FieldAccess) left).getName().getIdentifier().equals(fieldName))) {
                     
                     Expression right = node.getRightHandSide();
-                    if (right instanceof MethodInvocation mi) {
+                    List<Expression> tracedRights = traceVariableAll(right);
-                        String calledName = mi.getName().getIdentifier();
+                    for (Expression tracedRight : tracedRights) {
-                        String targetClass = context.getFqn(td);
+                        if (tracedRight instanceof MethodInvocation mi) {
-                        
+                            String calledName = mi.getName().getIdentifier();
-                        if (mi.getExpression() != null && mi.resolveMethodBinding() != null && mi.resolveMethodBinding().getDeclaringClass() != null) {
+                            String targetClass = context.getFqn(td);
-                            targetClass = mi.resolveMethodBinding().getDeclaringClass().getQualifiedName();
+                            
-                        } else if (mi.getExpression() instanceof SimpleName) {
+                            if (mi.getExpression() != null && mi.resolveMethodBinding() != null && mi.resolveMethodBinding().getDeclaringClass() != null) {
-                            String exprName = ((SimpleName) mi.getExpression()).getIdentifier();
+                                targetClass = mi.resolveMethodBinding().getDeclaringClass().getQualifiedName();
-                            CompilationUnit cu = td.getRoot() instanceof CompilationUnit ? (CompilationUnit) td.getRoot() : null;
+                            } else if (mi.getExpression() instanceof SimpleName) {
-                            TypeDeclaration targetTd = context.resolveStaticImport(exprName, cu);
+                                String exprName = ((SimpleName) mi.getExpression()).getIdentifier();
-                            if (targetTd != null) {
+                                CompilationUnit cu = td.getRoot() instanceof CompilationUnit ? (CompilationUnit) td.getRoot() : null;
-                                targetClass = context.getFqn(targetTd);
+                                TypeDeclaration targetTd = context.resolveStaticImport(exprName, cu);
+                                if (targetTd != null) {
+                                    targetClass = context.getFqn(targetTd);
+                                }
                             }
-                        }
 
-                        CompilationUnit cu = td.getRoot() instanceof CompilationUnit ? (CompilationUnit) td.getRoot() : null;
+                            CompilationUnit cu = td.getRoot() instanceof CompilationUnit ? (CompilationUnit) td.getRoot() : null;
-                        results.addAll(resolveMethodReturnConstant(targetClass, calledName, 0, new HashSet<>(visited), cu));
+                            results.addAll(resolveMethodReturnConstant(targetClass, calledName, 0, new HashSet<>(visited), cu));
-                    } else {
+                        } else {
-                        String val = constantResolver.resolve(right, context);
+                            String val = constantResolver.resolve(tracedRight, context);
-                        if (val != null) results.add(val);
+                            if (val != null) results.add(val);
+                        }
                     }
                 }
                 return super.visit(node);
@@ -1303,10 +1332,11 @@ public class HeuristicCallGraphEngine implements CallGraphEngine {
             public boolean visit(Assignment asn) {
                 Expression left = asn.getLeftHandSide();
                 if ((left instanceof SimpleName sn && sn.getIdentifier().equals(fieldName)) ||
-                    (left instanceof FieldAccess fa && fa.getName().getIdentifier().equals(fieldName))) {
+                    (left instanceof FieldAccess fa && fa.getName().getIdentifier().equals(fieldName)) ||
+                    (left instanceof SuperFieldAccess sfa && sfa.getName().getIdentifier().equals(fieldName))) {
                     Expression right = asn.getRightHandSide();
-                    if (right instanceof SimpleName snRight) {
+                    String rightName = extractVariableName(right);
-                        String rightName = snRight.getIdentifier();
+                    if (rightName != null) {
                         for (int i = 0; i < constructorMd.parameters().size(); i++) {
                             SingleVariableDeclaration svd = (SingleVariableDeclaration) constructorMd.parameters().get(i);
                             if (svd.getName().getIdentifier().equals(rightName)) {
@@ -1319,15 +1349,22 @@ public class HeuristicCallGraphEngine implements CallGraphEngine {
             }
             @Override
             public boolean visit(ConstructorInvocation node) {
+                org.eclipse.jdt.core.dom.IMethodBinding resolvedBinding = node.resolveConstructorBinding();
                 TypeDeclaration enclosingTd = findEnclosingType(node);
                 if (enclosingTd != null) {
                     for (MethodDeclaration otherMd : enclosingTd.getMethods()) {
-                        if (otherMd.isConstructor() && otherMd != constructorMd && otherMd.parameters().size() == node.arguments().size()) {
+                        boolean matches = false;
+                        if (resolvedBinding != null && otherMd.resolveBinding() != null) {
+                            matches = resolvedBinding.isEqualTo(otherMd.resolveBinding()) || resolvedBinding.getKey().equals(otherMd.resolveBinding().getKey());
+                        } else {
+                            matches = otherMd.isConstructor() && otherMd != constructorMd && otherMd.parameters().size() == node.arguments().size();
+                        }
+                        if (matches) {
                             int targetIdx = findAssignedParameterIndex(otherMd, fieldName, context, visited);
                             if (targetIdx >= 0 && targetIdx < node.arguments().size()) {
                                 Expression arg = (Expression) node.arguments().get(targetIdx);
-                                if (arg instanceof SimpleName snArg) {
+                                String argName = extractVariableName(arg);
-                                    String argName = snArg.getIdentifier();
+                                if (argName != null) {
                                     for (int i = 0; i < constructorMd.parameters().size(); i++) {
                                         SingleVariableDeclaration svd = (SingleVariableDeclaration) constructorMd.parameters().get(i);
                                         if (svd.getName().getIdentifier().equals(argName)) {
@@ -1343,6 +1380,7 @@ public class HeuristicCallGraphEngine implements CallGraphEngine {
             }
             @Override
             public boolean visit(SuperConstructorInvocation node) {
+                org.eclipse.jdt.core.dom.IMethodBinding resolvedBinding = node.resolveConstructorBinding();
                 TypeDeclaration enclosingTd = findEnclosingType(node);
                 if (enclosingTd != null) {
                     String superFqn = context.getSuperclassFqn(enclosingTd);
@@ -1350,12 +1388,18 @@ public class HeuristicCallGraphEngine implements CallGraphEngine {
                         TypeDeclaration superTd = context.getTypeDeclaration(superFqn);
                         if (superTd != null) {
                             for (MethodDeclaration otherMd : superTd.getMethods()) {
-                                if (otherMd.isConstructor() && otherMd.parameters().size() == node.arguments().size()) {
+                                boolean matches = false;
+                                if (resolvedBinding != null && otherMd.resolveBinding() != null) {
+                                    matches = resolvedBinding.isEqualTo(otherMd.resolveBinding()) || resolvedBinding.getKey().equals(otherMd.resolveBinding().getKey());
+                                } else {
+                                    matches = otherMd.isConstructor() && otherMd.parameters().size() == node.arguments().size();
+                                }
+                                if (matches) {
                                     int targetIdx = findAssignedParameterIndex(otherMd, fieldName, context, visited);
                                     if (targetIdx >= 0 && targetIdx < node.arguments().size()) {
                                         Expression arg = (Expression) node.arguments().get(targetIdx);
-                                        if (arg instanceof SimpleName snArg) {
+                                        String argName = extractVariableName(arg);
-                                            String argName = snArg.getIdentifier();
+                                        if (argName != null) {
                                             for (int i = 0; i < constructorMd.parameters().size(); i++) {
                                                 SingleVariableDeclaration svd = (SingleVariableDeclaration) constructorMd.parameters().get(i);
                                                 if (svd.getName().getIdentifier().equals(argName)) {
@@ -1374,4 +1418,16 @@ public class HeuristicCallGraphEngine implements CallGraphEngine {
         });
         return foundIdx[0];
     }
+
+    private String extractVariableName(Expression expr) {
+        if (expr instanceof SimpleName sn) return sn.getIdentifier();
+        if (expr instanceof FieldAccess fa) return fa.getName().getIdentifier();
+        if (expr instanceof SuperFieldAccess sfa) return sfa.getName().getIdentifier();
+        if (expr instanceof CastExpression ce) return extractVariableName(ce.getExpression());
+        if (expr instanceof ParenthesizedExpression pe) return extractVariableName(pe.getExpression());
+        if (expr instanceof MethodInvocation mi && (mi.getName().getIdentifier().equals("requireNonNull") || mi.getName().getIdentifier().equals("notNull"))) {
+            if (!mi.arguments().isEmpty()) return extractVariableName((Expression)mi.arguments().get(0));
+        }
+        return null;
+    }
 }

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

@@ -0,0 +1,69 @@
+package click.kamil.springstatemachineexporter.analysis.service;
+
+import click.kamil.springstatemachineexporter.analysis.model.CallChain;
+import click.kamil.springstatemachineexporter.analysis.model.EntryPoint;
+import click.kamil.springstatemachineexporter.analysis.model.TriggerPoint;
+import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+import static org.assertj.core.api.Assertions.assertThat;
+
+class BuilderLocalVariableTest {
+
+    @Test
+    void shouldTraceLocalVariableInBuilderPattern(@TempDir Path tempDir) throws IOException {
+        String source = """
+                package com.example;
+                public class OrderController {
+                    private OrderService service;
+                    public void processOrderEvent() {
+                        OrderEvents e = OrderEvents.PAY;
+                        RichOrderEvent event = new RichOrderEvent();
+                        event.setType(e);
+                        service.updateOrderState(event.getType());
+                    }
+                }
+                
+                class OrderService {
+                    public void updateOrderState(OrderEvents event) {}
+                }
+                
+                class RichOrderEvent {
+                    private OrderEvents type;
+                    public void setType(OrderEvents type) { this.type = type; }
+                    public OrderEvents getType() { return type; }
+                }
+                
+                enum OrderEvents { PAY }
+                """;
+        Files.writeString(tempDir.resolve("OrderConfig.java"), source);
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        HeuristicCallGraphEngine builder = new HeuristicCallGraphEngine(context);
+
+        EntryPoint entryPoint = EntryPoint.builder()
+                .className("com.example.OrderController")
+                .methodName("processOrderEvent")
+                .build();
+
+        TriggerPoint trigger = TriggerPoint.builder()
+                .className("com.example.OrderService")
+                .methodName("updateOrderState")
+                .event("event")
+                .build();
+
+        List<CallChain> chains = builder.findChains(List.of(entryPoint), List.of(trigger));
+
+        assertThat(chains).hasSize(1);
+        CallChain chain = chains.get(0);
+        System.out.println("POLY EVENTS FOUND: " + chain.getTriggerPoint().getPolymorphicEvents());
+        assertThat(chain.getTriggerPoint().getPolymorphicEvents())
+            .containsExactlyInAnyOrder("OrderEvents.PAY");
+    }
+}

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

@@ -266,6 +266,143 @@ class ConstructorInvocationTest {
 
         List<CallChain> chains = builder.findChains(List.of(entryPoint), List.of(trigger));
 
+        assertThat(chains).hasSize(1);
+        CallChain chain = chains.get(0);
+        assertThat(chain.getTriggerPoint().getPolymorphicEvents())
+            .contains("OrderEvents.PAY");
+    }
+    @Test
+    void shouldTraceThroughWrappedConstructorArguments(@TempDir Path tempDir) throws IOException {
+        String source = """
+                package com.example;
+                import java.util.Objects;
+                
+                public class OrderController {
+                    private OrderService service;
+                    public void processOrderEvent(RichOrderEvent domainEvent) {
+                        service.updateOrderState(domainEvent.getType());
+                    }
+                }
+                
+                class OrderService {
+                    public void updateOrderState(OrderEvents event) {}
+                }
+                
+                class RichOrderEvent {
+                    private OrderEvents type;
+                    
+                    public RichOrderEvent() {
+                        this(Objects.requireNonNull(OrderEvents.PAY));
+                    }
+                    
+                    public RichOrderEvent(OrderEvents type) {
+                        this((OrderEvents) type, "DEFAULT_INFO");
+                    }
+                    
+                    public RichOrderEvent(OrderEvents type, String info) {
+                        this.type = Objects.requireNonNull(type);
+                    }
+                    
+                    public OrderEvents getType() { return type; }
+                }
+                
+                enum OrderEvents { PAY }
+                """;
+        Files.writeString(tempDir.resolve("OrderConfig.java"), source);
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        HeuristicCallGraphEngine builder = new HeuristicCallGraphEngine(context);
+
+        EntryPoint entryPoint = EntryPoint.builder()
+                .className("com.example.OrderController")
+                .methodName("processOrderEvent")
+                .build();
+
+        TriggerPoint trigger = TriggerPoint.builder()
+                .className("com.example.OrderService")
+                .methodName("updateOrderState")
+                .event("event")
+                .build();
+
+        List<CallChain> chains = builder.findChains(List.of(entryPoint), List.of(trigger));
+
+        assertThat(chains).hasSize(1);
+        CallChain chain = chains.get(0);
+        assertThat(chain.getTriggerPoint().getPolymorphicEvents())
+            .contains("OrderEvents.PAY");
+    }
+
+    @Test
+    void shouldTraceThroughComplexOverloadedConstructors(@TempDir Path tempDir) throws IOException {
+        String source = """
+                package com.example;
+                
+                public class OrderController {
+                    private OrderService service;
+                    public void processOrderEvent(EnterpriseEvent domainEvent) {
+                        service.updateOrderState(domainEvent.getType());
+                    }
+                }
+                
+                class OrderService {
+                    public void updateOrderState(OrderEvents event) {}
+                }
+                
+                class EnterpriseEvent {
+                    private OrderEvents type;
+                    private String info;
+                    
+                    public EnterpriseEvent() {
+                        this(OrderEvents.PAY);
+                    }
+                    
+                    // Matches by parameter count but wrong type!
+                    public EnterpriseEvent(String info) {
+                        this.info = info;
+                    }
+                    
+                    public EnterpriseEvent(OrderEvents type) {
+                        this("info", type);
+                    }
+                    
+                    // Target of previous constructor
+                    public EnterpriseEvent(String info, OrderEvents type) {
+                        this(type, info, null);
+                    }
+                    
+                    // Final destination
+                    public EnterpriseEvent(OrderEvents type, String info, Object extra) {
+                        this.type = type;
+                        this.info = info;
+                    }
+                    
+                    public OrderEvents getType() { return type; }
+                }
+                
+                enum OrderEvents { PAY }
+                """;
+        Files.writeString(tempDir.resolve("OrderConfig.java"), source);
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        HeuristicCallGraphEngine builder = new HeuristicCallGraphEngine(context);
+
+        EntryPoint entryPoint = EntryPoint.builder()
+                .className("com.example.OrderController")
+                .methodName("processOrderEvent")
+                .build();
+
+        TriggerPoint trigger = TriggerPoint.builder()
+                .className("com.example.OrderService")
+                .methodName("updateOrderState")
+                .event("event")
+                .build();
+
+        List<CallChain> chains = builder.findChains(List.of(entryPoint), List.of(trigger));
+
         assertThat(chains).hasSize(1);
         CallChain chain = chains.get(0);
         assertThat(chain.getTriggerPoint().getPolymorphicEvents())

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

@@ -0,0 +1,72 @@
+package click.kamil.springstatemachineexporter.analysis.service;
+
+import click.kamil.springstatemachineexporter.analysis.model.CallChain;
+import click.kamil.springstatemachineexporter.analysis.model.EntryPoint;
+import click.kamil.springstatemachineexporter.analysis.model.TriggerPoint;
+import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+import static org.assertj.core.api.Assertions.assertThat;
+
+class ConstructorLocalVariableTest {
+
+    @Test
+    void shouldTraceLocalVariableInConstructor(@TempDir Path tempDir) throws IOException {
+        String source = """
+                package com.example;
+                public class OrderController {
+                    private OrderService service;
+                    public void processOrderEvent(RichOrderEvent domainEvent) {
+                        service.updateOrderState(domainEvent.getType());
+                    }
+                }
+                
+                class OrderService {
+                    public void updateOrderState(OrderEvents event) {}
+                }
+                
+                class RichOrderEvent {
+                    private OrderEvents type;
+                    
+                    public RichOrderEvent() {
+                        OrderEvents e = OrderEvents.PAY;
+                        this.type = e;
+                    }
+                    
+                    public OrderEvents getType() { 
+                        return type; 
+                    }
+                }
+                
+                enum OrderEvents { PAY }
+                """;
+        Files.writeString(tempDir.resolve("OrderConfig.java"), source);
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        HeuristicCallGraphEngine builder = new HeuristicCallGraphEngine(context);
+
+        EntryPoint entryPoint = EntryPoint.builder()
+                .className("com.example.OrderController")
+                .methodName("processOrderEvent")
+                .build();
+
+        TriggerPoint trigger = TriggerPoint.builder()
+                .className("com.example.OrderService")
+                .methodName("updateOrderState")
+                .event("event")
+                .build();
+
+        List<CallChain> chains = builder.findChains(List.of(entryPoint), List.of(trigger));
+
+        assertThat(chains).hasSize(1);
+        CallChain chain = chains.get(0);
+        assertThat(chain.getTriggerPoint().getPolymorphicEvents())
+            .containsExactlyInAnyOrder("OrderEvents.PAY");
+    }
+}

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

@@ -0,0 +1,122 @@
+package click.kamil.springstatemachineexporter.analysis.service;
+
+import click.kamil.springstatemachineexporter.analysis.model.CallChain;
+import click.kamil.springstatemachineexporter.analysis.model.EntryPoint;
+import click.kamil.springstatemachineexporter.analysis.model.TriggerPoint;
+import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+import static org.assertj.core.api.Assertions.assertThat;
+
+class LocalVariableTest {
+
+    @Test
+    void shouldTraceLocalVariableInGetter(@TempDir Path tempDir) throws IOException {
+        String source = """
+                package com.example;
+                public class OrderController {
+                    private OrderService service;
+                    public void processOrderEvent(RichOrderEvent domainEvent) {
+                        service.updateOrderState(domainEvent.getType());
+                    }
+                }
+                
+                class OrderService {
+                    public void updateOrderState(OrderEvents event) {}
+                }
+                
+                class RichOrderEvent {
+                    public OrderEvents getType() { 
+                        OrderEvents e = OrderEvents.PAY;
+                        return e; 
+                    }
+                }
+                
+                enum OrderEvents { PAY }
+                """;
+        Files.writeString(tempDir.resolve("OrderConfig.java"), source);
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        HeuristicCallGraphEngine builder = new HeuristicCallGraphEngine(context);
+
+        EntryPoint entryPoint = EntryPoint.builder()
+                .className("com.example.OrderController")
+                .methodName("processOrderEvent")
+                .build();
+
+        TriggerPoint trigger = TriggerPoint.builder()
+                .className("com.example.OrderService")
+                .methodName("updateOrderState")
+                .event("event")
+                .build();
+
+        List<CallChain> chains = builder.findChains(List.of(entryPoint), List.of(trigger));
+
+        assertThat(chains).hasSize(1);
+        CallChain chain = chains.get(0);
+        assertThat(chain.getTriggerPoint().getPolymorphicEvents())
+            .contains("OrderEvents.PAY")
+            .doesNotContain("e"); // Should not blindly add the variable name
+    }
+
+    @Test
+    void shouldTraceMultipleAssignments(@TempDir Path tempDir) throws IOException {
+        String source = """
+                package com.example;
+                public class OrderController {
+                    private OrderService service;
+                    public void processOrderEvent(RichOrderEvent domainEvent) {
+                        service.updateOrderState(domainEvent.getType());
+                    }
+                }
+                
+                class OrderService {
+                    public void updateOrderState(OrderEvents event) {}
+                }
+                
+                class RichOrderEvent {
+                    public OrderEvents getType() { 
+                        OrderEvents e = null;
+                        if (System.currentTimeMillis() % 2 == 0) {
+                            e = OrderEvents.PAY;
+                        } else {
+                            e = OrderEvents.CANCEL;
+                        }
+                        return e; 
+                    }
+                }
+                
+                enum OrderEvents { PAY, CANCEL }
+                """;
+        Files.writeString(tempDir.resolve("OrderConfig.java"), source);
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        HeuristicCallGraphEngine builder = new HeuristicCallGraphEngine(context);
+
+        EntryPoint entryPoint = EntryPoint.builder()
+                .className("com.example.OrderController")
+                .methodName("processOrderEvent")
+                .build();
+
+        TriggerPoint trigger = TriggerPoint.builder()
+                .className("com.example.OrderService")
+                .methodName("updateOrderState")
+                .event("event")
+                .build();
+
+        List<CallChain> chains = builder.findChains(List.of(entryPoint), List.of(trigger));
+
+        assertThat(chains).hasSize(1);
+        CallChain chain = chains.get(0);
+        assertThat(chain.getTriggerPoint().getPolymorphicEvents())
+            .containsExactlyInAnyOrder("OrderEvents.PAY", "OrderEvents.CANCEL");
+    }
+}

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

@@ -0,0 +1,66 @@
+package click.kamil.springstatemachineexporter.analysis.service;
+
+import click.kamil.springstatemachineexporter.analysis.model.CallChain;
+import click.kamil.springstatemachineexporter.analysis.model.EntryPoint;
+import click.kamil.springstatemachineexporter.analysis.model.TriggerPoint;
+import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+import static org.assertj.core.api.Assertions.assertThat;
+
+class ReturnStatementTest {
+
+    @Test
+    void shouldHandleAssignmentInReturn(@TempDir Path tempDir) throws IOException {
+        String source = """
+                package com.example;
+                public class OrderController {
+                    private OrderService service;
+                    public void processOrderEvent(RichOrderEvent domainEvent) {
+                        service.updateOrderState(domainEvent.getType());
+                    }
+                }
+                
+                class OrderService {
+                    public void updateOrderState(OrderEvents event) {}
+                }
+                
+                class RichOrderEvent {
+                    private OrderEvents type;
+                    public OrderEvents getType() { 
+                        return this.type = OrderEvents.PAY; 
+                    }
+                }
+                
+                enum OrderEvents { PAY }
+                """;
+        Files.writeString(tempDir.resolve("OrderConfig.java"), source);
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        HeuristicCallGraphEngine builder = new HeuristicCallGraphEngine(context);
+
+        EntryPoint entryPoint = EntryPoint.builder()
+                .className("com.example.OrderController")
+                .methodName("processOrderEvent")
+                .build();
+
+        TriggerPoint trigger = TriggerPoint.builder()
+                .className("com.example.OrderService")
+                .methodName("updateOrderState")
+                .event("event")
+                .build();
+
+        List<CallChain> chains = builder.findChains(List.of(entryPoint), List.of(trigger));
+
+        assertThat(chains).hasSize(1);
+        CallChain chain = chains.get(0);
+        assertThat(chain.getTriggerPoint().getPolymorphicEvents())
+            .containsExactlyInAnyOrder("OrderEvents.PAY");
+    }
+}