update tranistions

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

@@ -27,7 +27,13 @@ public class HeuristicEventMatchingEngine implements EventMatchingEngine {
                     hasPolyMatch = true;
                     break;
                 }
-                continue;
+                String classPe = pe.substring(0, pe.lastIndexOf('.'));
+                String classSm = smEventRaw.substring(0, smEventRaw.lastIndexOf('.'));
+                String simpleClassPe = classPe.contains(".") ? classPe.substring(classPe.lastIndexOf('.') + 1) : classPe;
+                String simpleClassSm = classSm.contains(".") ? classSm.substring(classSm.lastIndexOf('.') + 1) : classSm;
+                if (!simpleClassPe.equals(simpleClassSm)) {
+                    continue;
+                }
             }
 
             String simplePe = pe;
@@ -51,7 +57,13 @@ public class HeuristicEventMatchingEngine implements EventMatchingEngine {
             if (smEventRaw.equals(rawTriggerEvent) || smEventRaw.endsWith("." + rawTriggerEvent) || rawTriggerEvent.endsWith("." + smEventRaw)) {
                 return true;
             }
-            return false;
+            String classTrig = rawTriggerEvent.substring(0, rawTriggerEvent.lastIndexOf('.'));
+            String classSm = smEventRaw.substring(0, smEventRaw.lastIndexOf('.'));
+            String simpleClassTrig = classTrig.contains(".") ? classTrig.substring(classTrig.lastIndexOf('.') + 1) : classTrig;
+            String simpleClassSm = classSm.contains(".") ? classSm.substring(classSm.lastIndexOf('.') + 1) : classSm;
+            if (!simpleClassTrig.equals(simpleClassSm)) {
+                return false;
+            }
         }
 
         String triggerEvent = simplify(rawTriggerEvent);

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

@@ -929,6 +929,8 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
             if (!exprStr.equals("Optional") && !exprStr.equals("Stream") && !exprStr.equals("List") && !exprStr.equals("Set") && !exprStr.equals("Arrays") && !exprStr.equals("Objects")) {
                 return mi;
             }
+        } else {
+            return mi;
         }
         
         if (!mi.arguments().isEmpty()) {

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

@@ -190,6 +190,13 @@ public class HeuristicCallGraphEngine extends AbstractCallGraphEngine {
         Expression receiver = node.getExpression();
         String methodName = node.getName().getIdentifier();
 
+        if (receiver instanceof MethodInvocation miReceiver) {
+            String returnType = resolveMethodInvocationReturnType(miReceiver);
+            if (returnType != null) {
+                return returnType + "." + methodName;
+            }
+        }
+
         if (receiver == null) {
             TypeDeclaration td = findEnclosingType(node);
             if (td != null) {
@@ -572,4 +579,44 @@ public class HeuristicCallGraphEngine extends AbstractCallGraphEngine {
         });
         return foundReceiver[0];
     }
+
+    protected String resolveMethodInvocationReturnType(MethodInvocation mi) {
+        Expression receiver = mi.getExpression();
+        String receiverType = null;
+        if (receiver == null) {
+            TypeDeclaration td = findEnclosingType(mi);
+            if (td != null) {
+                receiverType = context.getFqn(td);
+            }
+        } else if (receiver instanceof SimpleName sn) {
+            receiverType = resolveReceiverTypeFallback(sn);
+        } else if (receiver instanceof FieldAccess fa) {
+            receiverType = resolveReceiverTypeFallback(fa.getName());
+        } else if (receiver instanceof MethodInvocation innerMi) {
+            receiverType = resolveMethodInvocationReturnType(innerMi);
+        }
+        
+        if (receiverType == null) {
+            ITypeBinding binding = receiver != null ? receiver.resolveTypeBinding() : null;
+            if (binding != null) {
+                receiverType = binding.getErasure().getQualifiedName();
+            }
+        }
+        
+        if (receiverType != null) {
+            String methodName = mi.getName().getIdentifier();
+            TypeDeclaration td = context.getTypeDeclaration(receiverType);
+            if (td == null && receiverType.contains(".")) {
+                String simpleClass = receiverType.substring(receiverType.lastIndexOf('.') + 1);
+                td = context.getTypeDeclaration(simpleClass);
+            }
+            if (td != null) {
+                MethodDeclaration md = context.findMethodDeclaration(td, methodName, true);
+                if (md != null && md.getReturnType2() != null) {
+                    return resolveTypeToFqn(md.getReturnType2(), mi);
+                }
+            }
+        }
+        return null;
+    }
 }

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

@@ -170,6 +170,13 @@ public class JdtCallGraphEngine extends AbstractCallGraphEngine {
         Expression receiver = node.getExpression();
         String methodName = node.getName().getIdentifier();
 
+        if (receiver instanceof MethodInvocation miReceiver) {
+            String returnType = resolveMethodInvocationReturnType(miReceiver);
+            if (returnType != null) {
+                return returnType + "." + methodName;
+            }
+        }
+
         if (receiver == null) {
             TypeDeclaration td = findEnclosingType(node);
             if (td != null) {
@@ -344,4 +351,44 @@ public class JdtCallGraphEngine extends AbstractCallGraphEngine {
         }
         return expr;
     }
+
+    protected String resolveMethodInvocationReturnType(MethodInvocation mi) {
+        Expression receiver = mi.getExpression();
+        String receiverType = null;
+        if (receiver == null) {
+            TypeDeclaration td = findEnclosingType(mi);
+            if (td != null) {
+                receiverType = context.getFqn(td);
+            }
+        } else if (receiver instanceof SimpleName sn) {
+            receiverType = resolveReceiverTypeFallback(sn);
+        } else if (receiver instanceof FieldAccess fa) {
+            receiverType = resolveReceiverTypeFallback(fa.getName());
+        } else if (receiver instanceof MethodInvocation innerMi) {
+            receiverType = resolveMethodInvocationReturnType(innerMi);
+        }
+        
+        if (receiverType == null) {
+            ITypeBinding binding = receiver != null ? receiver.resolveTypeBinding() : null;
+            if (binding != null) {
+                receiverType = binding.getErasure().getQualifiedName();
+            }
+        }
+        
+        if (receiverType != null) {
+            String methodName = mi.getName().getIdentifier();
+            TypeDeclaration td = context.getTypeDeclaration(receiverType);
+            if (td == null && receiverType.contains(".")) {
+                String simpleClass = receiverType.substring(receiverType.lastIndexOf('.') + 1);
+                td = context.getTypeDeclaration(simpleClass);
+            }
+            if (td != null) {
+                MethodDeclaration md = context.findMethodDeclaration(td, methodName, true);
+                if (md != null && md.getReturnType2() != null) {
+                    return resolveTypeToFqn(md.getReturnType2(), mi);
+                }
+            }
+        }
+        return null;
+    }
 }

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

@@ -258,7 +258,7 @@ public class VariableTracer {
                     for (Expression expr : initializers) {
                         Expression traced = traceVariable(expr);
                         if (traced instanceof MethodInvocation mi) {
-                            Expression innerMost = unwrapMethodInvocation(mi, 0);
+                            Expression innerMost = unwrapMethodInvocation(mi, 0, methodFqn);
                             stringified.add(innerMost.toString());
                         } else if (traced instanceof ArrayInitializer) {
                             stringified.add("new Object[]" + traced.toString());
@@ -513,13 +513,61 @@ public class VariableTracer {
         return (TypeDeclaration) parent;
     }
 
-    private Expression unwrapMethodInvocation(MethodInvocation mi, int depth) {
+    private int getReturnedParameterIndex(MethodDeclaration md) {
+        if (md.getBody() == null) return -1;
+        List<?> parameters = md.parameters();
+        if (parameters.isEmpty()) return -1;
+        
+        final List<String> returnedExprs = new ArrayList<>();
+        md.getBody().accept(new ASTVisitor() {
+            @Override
+            public boolean visit(ReturnStatement node) {
+                if (node.getExpression() != null) {
+                    returnedExprs.add(node.getExpression().toString());
+                }
+                return super.visit(node);
+            }
+        });
+        
+        if (returnedExprs.size() == 1) {
+            String retStr = returnedExprs.get(0);
+            for (int i = 0; i < parameters.size(); i++) {
+                if (parameters.get(i) instanceof SingleVariableDeclaration svd) {
+                    if (svd.getName().getIdentifier().equals(retStr)) {
+                        return i;
+                    }
+                }
+            }
+        }
+        return -1;
+    }
+
+    private Expression unwrapMethodInvocation(MethodInvocation mi, int depth, String methodFqn) {
         if (depth > 5) return mi;
         if (mi.getExpression() != null) {
             String exprStr = mi.getExpression().toString();
             if (!exprStr.equals("Optional") && !exprStr.equals("Stream") && !exprStr.equals("List") && !exprStr.equals("Set") && !exprStr.equals("Arrays") && !exprStr.equals("Objects")) {
                 return mi;
             }
+        } else {
+            if (methodFqn != null && methodFqn.contains(".")) {
+                String className = methodFqn.substring(0, methodFqn.lastIndexOf('.'));
+                TypeDeclaration td = context.getTypeDeclaration(className);
+                if (td != null) {
+                    MethodDeclaration localMd = context.findMethodDeclaration(td, mi.getName().getIdentifier(), true);
+                    if (localMd != null) {
+                        int paramIdx = getReturnedParameterIndex(localMd);
+                        if (paramIdx >= 0 && paramIdx < mi.arguments().size()) {
+                            Expression arg = (Expression) mi.arguments().get(paramIdx);
+                            if (arg instanceof MethodInvocation innerMi) {
+                                return unwrapMethodInvocation(innerMi, depth + 1, methodFqn);
+                            }
+                            return arg;
+                        }
+                    }
+                }
+            }
+            return mi;
         }
         
         if (!mi.arguments().isEmpty()) {
@@ -540,7 +588,7 @@ public class VariableTracer {
             
             Expression arg = (Expression) mi.arguments().get(0);
             if (arg instanceof MethodInvocation innerMi) {
-                return unwrapMethodInvocation(innerMi, depth + 1);
+                return unwrapMethodInvocation(innerMi, depth + 1, methodFqn);
             }
             return arg;
         }

state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/enricher/matching/HeuristicEventMatchingEngineTest.java

@@ -106,4 +106,23 @@ class HeuristicEventMatchingEngineTest {
         TriggerPoint triggerPoint2 = TriggerPoint.builder().event("pay").build();
         assertThat(engine.matches(smEvent2, triggerPoint2)).isTrue();
     }
+
+    @Test
+    void shouldMatchDottedEventPackageMismatch() {
+        Event smEvent = Event.of("CREATE", "com.example.other.OrderEvent.CREATE");
+        TriggerPoint triggerPoint = TriggerPoint.builder().event("com.example.some.OrderEvent.CREATE").build();
+
+        assertThat(engine.matches(smEvent, triggerPoint)).isTrue();
+    }
+
+    @Test
+    void shouldMatchPolymorphicDottedEventPackageMismatch() {
+        Event smEvent = Event.of("CREATE", "com.example.other.OrderEvent.CREATE");
+        TriggerPoint triggerPoint = TriggerPoint.builder()
+                .event("getType()")
+                .polymorphicEvents(List.of("com.example.some.OrderEvent.CREATE"))
+                .build();
+
+        assertThat(engine.matches(smEvent, triggerPoint)).isTrue();
+    }
 }

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

@@ -119,4 +119,58 @@ class LocalVariableTest {
         assertThat(chain.getTriggerPoint().getPolymorphicEvents())
             .containsExactlyInAnyOrder("OrderEvents.PAY", "OrderEvents.CANCEL");
     }
+
+    @Test
+    void shouldNotUnwrapLocalDomainMethod(@TempDir Path tempDir) throws IOException {
+        String source = """
+                package com.example;
+                public class OrderController {
+                    private OrderService service;
+                    public void processOrderEvent(RichOrderEvent domainEvent) {
+                        OrderEvents resolved = validateEvent(domainEvent);
+                        service.updateOrderState(resolved);
+                    }
+                    
+                    private OrderEvents validateEvent(RichOrderEvent event) {
+                        return OrderEvents.CANCEL;
+                    }
+                }
+                
+                class OrderService {
+                    public void updateOrderState(OrderEvents event) {}
+                }
+                
+                class RichOrderEvent {
+                    public OrderEvents getType() { 
+                        return OrderEvents.PAY; 
+                    }
+                }
+                
+                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.CANCEL");
+    }
 }

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

@@ -0,0 +1,95 @@
+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 PolymorphicDispatchCallGraphTest {
+
+    @Test
+    void shouldResolvePolymorphicDispatchCalls(@TempDir Path tempDir) throws IOException {
+        String source = """
+                package com.example;
+                
+                public class OrderController {
+                    private OrderServiceRegistry registry;
+                    
+                    public void handleRequest(String order, String event) {
+                        registry.getOrderService(order).processOrderEvent(event);
+                    }
+                }
+                
+                class OrderServiceRegistry {
+                    public OrderService getOrderService(String order) {
+                        return null; // dynamic lookup
+                    }
+                }
+                
+                interface OrderService {
+                    void processOrderEvent(String event);
+                }
+                
+                class PlatformOrderService implements OrderService {
+                    public void processOrderEvent(String event) {
+                        // trigger state machine
+                    }
+                }
+                
+                class TransportOrderService implements OrderService {
+                    public void processOrderEvent(String event) {
+                        // trigger state machine
+                    }
+                }
+                """;
+        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("handleRequest")
+                .build();
+
+        TriggerPoint triggerPlatform = TriggerPoint.builder()
+                .className("com.example.PlatformOrderService")
+                .methodName("processOrderEvent")
+                .event("event")
+                .build();
+
+        TriggerPoint triggerTransport = TriggerPoint.builder()
+                .className("com.example.TransportOrderService")
+                .methodName("processOrderEvent")
+                .event("event")
+                .build();
+
+        List<CallChain> chainsPlatform = builder.findChains(List.of(entryPoint), List.of(triggerPlatform));
+        List<CallChain> chainsTransport = builder.findChains(List.of(entryPoint), List.of(triggerTransport));
+
+        assertThat(chainsPlatform).hasSize(1);
+        assertThat(chainsPlatform.get(0).getMethodChain())
+                .containsExactly(
+                        "com.example.OrderController.handleRequest",
+                        "com.example.PlatformOrderService.processOrderEvent"
+                );
+
+        assertThat(chainsTransport).hasSize(1);
+        assertThat(chainsTransport.get(0).getMethodChain())
+                .containsExactly(
+                        "com.example.OrderController.handleRequest",
+                        "com.example.TransportOrderService.processOrderEvent"
+                );
+    }
+}