diff --git a/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/VariableTracer.java b/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/VariableTracer.java index eb21c15..288a682 100644 --- a/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/VariableTracer.java +++ b/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 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; }