update tranistions

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

@@ -904,7 +904,17 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
     protected String[] extractMethodSuffix(String paramName, String currentSuffix) {
         if (paramName != null && !paramName.contains(" ") && !paramName.startsWith("new ")) {
             int bracketIndex = paramName.indexOf('[');
-            int dotIndex = paramName.indexOf('.');
+            int dotIndex = -1;
+            int parenDepth = 0;
+            for (int i = 0; i < paramName.length(); i++) {
+                char c = paramName.charAt(i);
+                if (c == '(') parenDepth++;
+                else if (c == ')') parenDepth--;
+                else if (c == '.' && parenDepth == 0) {
+                    dotIndex = i;
+                    break;
+                }
+            }
             
             if (bracketIndex > 0 && (dotIndex == -1 || bracketIndex < dotIndex)) {
                 return new String[] { paramName.substring(0, bracketIndex), paramName.substring(bracketIndex) + currentSuffix };

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

@@ -28,6 +28,16 @@ public class VariableTracer {
         return all.isEmpty() ? expr : all.get(all.size() - 1);
     }
 
+    private Expression peelExpression(Expression expr) {
+        while (expr instanceof ParenthesizedExpression pe) {
+            expr = pe.getExpression();
+        }
+        if (expr instanceof CastExpression ce) {
+            return peelExpression(ce.getExpression());
+        }
+        return expr;
+    }
+
     public List<Expression> traceVariableAll(Expression expr) {
         return dataFlowModel.getReachingDefinitions(expr);
     }
@@ -82,6 +92,30 @@ public class VariableTracer {
         return null;
     }
 
+    private String getFieldType(TypeDeclaration td, String fieldName, CodebaseContext context, java.util.Set<String> visited) {
+        if (td == null) return null;
+        String typeFqn = context.getFqn(td);
+        if (!visited.add(typeFqn)) return null;
+
+        for (FieldDeclaration fd : td.getFields()) {
+            for (Object fragObj : fd.fragments()) {
+                VariableDeclarationFragment frag = (VariableDeclarationFragment) fragObj;
+                if (frag.getName().getIdentifier().equals(fieldName)) {
+                    return fd.getType().toString();
+                }
+            }
+        }
+        
+        String superFqn = context.getSuperclassFqn(td);
+        if (superFqn != null) {
+            TypeDeclaration superTd = context.getTypeDeclaration(superFqn);
+            if (superTd != null) {
+                return getFieldType(superTd, fieldName, context, visited);
+            }
+        }
+        return null;
+    }
+
     public String getVariableDeclaredType(String methodFqn, String cleanFieldName) {
         if (methodFqn == null) return null;
         int lastDot = methodFqn.lastIndexOf('.');
@@ -200,15 +234,9 @@ public class VariableTracer {
                 if (foundType[0] != null) return foundType[0];
             }
             
-            // Fallback to fields
+            // Fallback to fields including superclasses
-            for (FieldDeclaration fd : td.getFields()) {
+            String fieldType = getFieldType(td, cleanFieldName, context, new java.util.HashSet<>());
-                for (Object fragObj : fd.fragments()) {
+            if (fieldType != null) return fieldType;
-                    VariableDeclarationFragment frag = (VariableDeclarationFragment) fragObj;
-                    if (frag.getName().getIdentifier().equals(cleanFieldName)) {
-                        return fd.getType().toString();
-                    }
-                }
-            }
         }
     }
         
@@ -240,14 +268,14 @@ public class VariableTracer {
                     @Override
                     public boolean visit(VariableDeclarationFragment node) {
                         if (node.getName().getIdentifier().equals(varName) && node.getInitializer() != null) {
-                            initializers.add(node.getInitializer());
+                            initializers.add(peelExpression(node.getInitializer()));
                         }
                         return super.visit(node);
                     }
                     @Override
                     public boolean visit(Assignment node) {
                         if (node.getLeftHandSide() instanceof SimpleName asn && asn.getIdentifier().equals(varName)) {
-                            initializers.add(node.getRightHandSide());
+                            initializers.add(peelExpression(node.getRightHandSide()));
                         }
                         return super.visit(node);
                     }
@@ -258,7 +286,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,8 +541,33 @@ public class VariableTracer {
         return (TypeDeclaration) parent;
     }
 
-    private Expression unwrapMethodInvocation(MethodInvocation mi, int depth) {
+    private Expression unwrapMethodInvocation(MethodInvocation mi, int depth, String methodFqn) {
         if (depth > 5) return mi;
+        
+        String targetMethodFqn = resolveTargetMethodFqn(mi, methodFqn);
+        if (targetMethodFqn != null && targetMethodFqn.contains(".")) {
+            String className = targetMethodFqn.substring(0, targetMethodFqn.lastIndexOf('.'));
+            String methodName = targetMethodFqn.substring(targetMethodFqn.lastIndexOf('.') + 1);
+            
+            TypeDeclaration currentTd = methodFqn != null && methodFqn.contains(".") ? context.getTypeDeclaration(methodFqn.substring(0, methodFqn.lastIndexOf('.'))) : null;
+            CompilationUnit cu = currentTd != null && currentTd.getRoot() instanceof CompilationUnit ? (CompilationUnit) currentTd.getRoot() : null;
+            
+            TypeDeclaration td = context.getTypeDeclaration(className, cu);
+            if (td != null) {
+                MethodDeclaration localMd = context.findMethodDeclaration(td, methodName, true);
+                if (localMd != null) {
+                    int paramIdx = getReturnedParameterIndex(localMd);
+                    if (paramIdx >= 0 && paramIdx < mi.arguments().size()) {
+                        Expression arg = peelExpression((Expression) mi.arguments().get(paramIdx));
+                        if (arg instanceof MethodInvocation innerMi) {
+                            return unwrapMethodInvocation(innerMi, depth + 1, methodFqn);
+                        }
+                        return arg;
+                    }
+                }
+            }
+        }
+        
         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") && !exprStr.equals("Mono") && !exprStr.equals("Flux")) {
@@ -538,12 +591,61 @@ public class VariableTracer {
                 return mi;
             }
             
-            Expression arg = (Expression) mi.arguments().get(0);
+            Expression arg = peelExpression((Expression) mi.arguments().get(0));
             if (arg instanceof MethodInvocation innerMi) {
-                return unwrapMethodInvocation(innerMi, depth + 1);
+                return unwrapMethodInvocation(innerMi, depth + 1, methodFqn);
             }
             return arg;
         }
         return mi;
     }
+    
+    private String resolveTargetMethodFqn(MethodInvocation mi, String currentMethodFqn) {
+        if (currentMethodFqn == null || !currentMethodFqn.contains(".")) return null;
+        String currentClass = currentMethodFqn.substring(0, currentMethodFqn.lastIndexOf('.'));
+        if (mi.getExpression() == null || mi.getExpression().toString().equals("this")) {
+            return currentClass + "." + mi.getName().getIdentifier();
+        }
+        if (mi.getExpression() instanceof SimpleName sn) {
+            String declaredType = getVariableDeclaredType(currentMethodFqn, sn.getIdentifier());
+            if (declaredType != null) {
+                return declaredType + "." + mi.getName().getIdentifier();
+            }
+        } else if (mi.getExpression() instanceof FieldAccess fa) {
+            String declaredType = getVariableDeclaredType(currentMethodFqn, fa.getName().getIdentifier());
+            if (declaredType != null) {
+                return declaredType + "." + mi.getName().getIdentifier();
+            }
+        }
+        return null;
+    }
+
+    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;
+    }
 }

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

@@ -173,4 +173,111 @@ class LocalVariableTest {
         assertThat(chain.getTriggerPoint().getPolymorphicEvents())
             .containsExactlyInAnyOrder("OrderEvents.CANCEL");
     }
+    @Test
+    void shouldUnwrapExplicitPassthroughMethod(@TempDir Path tempDir) throws IOException {
+        String source = """
+                package com.example;
+                public class OrderController {
+                    private OrderService service;
+                    private EventValidator validator;
+                    public void processOrderEvent(RichOrderEvent domainEvent) {
+                        OrderEvents resolved = validator.passthrough(domainEvent.getType());
+                        service.updateOrderState(resolved);
+                    }
+                }
+                
+                class EventValidator {
+                    public OrderEvents passthrough(OrderEvents e) {
+                        if (e == null) throw new IllegalArgumentException();
+                        return e;
+                    }
+                }
+                
+                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.PAY");
+    }
+
+    @Test
+    void shouldFallbackToPrefixMatchingForExternalMethods(@TempDir Path tempDir) throws IOException {
+        String source = """
+                package com.example;
+                public class OrderController {
+                    private OrderService service;
+                    public void processOrderEvent(RichOrderEvent domainEvent) {
+                        OrderEvents resolved = UnknownExternalClass.validate(domainEvent.getType());
+                        service.updateOrderState(resolved);
+                    }
+                }
+                
+                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())
+            .doesNotContain("OrderEvents.PAY");
+    }
 }