update tranistions 3

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,8 +513,82 @@ 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 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 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 = (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")) {
@@ -540,7 +614,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;
         }