fixes 2

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

@@ -1323,8 +1323,14 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
                                 public boolean visit(org.eclipse.jdt.core.dom.ExpressionMethodReference node) {
                                     if (node.getName().getIdentifier().equals(mName)) {
                                         if (node.getParent() instanceof org.eclipse.jdt.core.dom.MethodInvocation parentMi) {
-                                            if (parentMi.getExpression() != null) {
+                                            if (click.kamil.springstatemachineexporter.ast.common.AstUtils.isFunctionalCollectionMethod(parentMi)) {
                                                 callerExprs.addAll(traceVariableAll(parentMi.getExpression(), visitedVariables));
+                                            } else {
+                                                for (Object arg : parentMi.arguments()) {
+                                                    if (arg != node) {
+                                                        callerExprs.addAll(traceVariableAll((Expression) arg, visitedVariables));
+                                                    }
+                                                }
                                             }
                                         }
                                     }

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

@@ -46,8 +46,7 @@ public class HeuristicCallGraphEngine extends AbstractCallGraphEngine {
                                             String refMethod = resolveMethodInType(td2, emr.getName().getIdentifier());
                                             if (refMethod != null) {
                                                 List<String> implicitArgs = new ArrayList<>();
-                                                String invokedMethod = node.getName().getIdentifier();
+                                                if (click.kamil.springstatemachineexporter.ast.common.AstUtils.isFunctionalCollectionMethod(node)) {
-                                                if (node.getExpression() != null && (invokedMethod.equals("forEach") || invokedMethod.equals("map") || invokedMethod.equals("filter") || invokedMethod.equals("flatMap"))) {
                                                     implicitArgs.add(node.getExpression().toString());
                                                 } else {
                                                     implicitArgs.addAll(args);
@@ -66,8 +65,7 @@ public class HeuristicCallGraphEngine extends AbstractCallGraphEngine {
                                         if (fallbackTypeFqn != null) {
                                             String calledMethod = fallbackTypeFqn + "." + emr.getName().getIdentifier();
                                             List<String> implicitArgs = new ArrayList<>();
-                                            String invokedMethod = node.getName().getIdentifier();
+                                            if (click.kamil.springstatemachineexporter.ast.common.AstUtils.isFunctionalCollectionMethod(node)) {
-                                            if (node.getExpression() != null && (invokedMethod.equals("forEach") || invokedMethod.equals("map") || invokedMethod.equals("filter") || invokedMethod.equals("flatMap"))) {
                                                 implicitArgs.add(node.getExpression().toString());
                                             } else {
                                                 implicitArgs.addAll(args);

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

@@ -50,8 +50,7 @@ public class JdtCallGraphEngine extends AbstractCallGraphEngine {
                                             String refMethod = resolveMethodInType(td2, emr.getName().getIdentifier());
                                             if (refMethod != null) {
                                                 List<String> implicitArgs = new ArrayList<>();
-                                                String invokedMethod = node.getName().getIdentifier();
+                                                if (click.kamil.springstatemachineexporter.ast.common.AstUtils.isFunctionalCollectionMethod(node)) {
-                                                if (node.getExpression() != null && (invokedMethod.equals("forEach") || invokedMethod.equals("map") || invokedMethod.equals("filter") || invokedMethod.equals("flatMap"))) {
                                                     implicitArgs.add(node.getExpression().toString());
                                                 } else {
                                                     implicitArgs.addAll(args);
@@ -70,8 +69,7 @@ public class JdtCallGraphEngine extends AbstractCallGraphEngine {
                                         if (fallbackTypeFqn != null) {
                                             String calledMethod = fallbackTypeFqn + "." + emr.getName().getIdentifier();
                                             List<String> implicitArgs = new ArrayList<>();
-                                            String invokedMethod = node.getName().getIdentifier();
+                                            if (click.kamil.springstatemachineexporter.ast.common.AstUtils.isFunctionalCollectionMethod(node)) {
-                                            if (node.getExpression() != null && (invokedMethod.equals("forEach") || invokedMethod.equals("map") || invokedMethod.equals("filter") || invokedMethod.equals("flatMap"))) {
                                                 implicitArgs.add(node.getExpression().toString());
                                             } else {
                                                 implicitArgs.addAll(args);

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/common/AstUtils.java

@@ -36,4 +36,62 @@ public final class AstUtils {
             return type.toString(); // fallback
         }
     }
+
+    public static boolean isFunctionalCollectionMethod(org.eclipse.jdt.core.dom.MethodInvocation node) {
+        if (node.getExpression() == null) {
+            return false;
+        }
+
+        org.eclipse.jdt.core.dom.ITypeBinding typeBinding = node.getExpression().resolveTypeBinding();
+        if (typeBinding != null) {
+            if (isCollectionOrStreamType(typeBinding)) {
+                return true;
+            }
+            return false;
+        }
+        
+        String methodName = node.getName().getIdentifier();
+        if (methodName.equals("forEach") || methodName.equals("map") || 
+            methodName.equals("filter") || methodName.equals("flatMap") || 
+            methodName.equals("anyMatch") || methodName.equals("allMatch") || 
+            methodName.equals("noneMatch") || methodName.equals("reduce") ||
+            methodName.equals("collect") || methodName.equals("removeIf") ||
+            methodName.equals("ifPresent") || methodName.equals("ifPresentOrElse") ||
+            methodName.equals("compute") || methodName.equals("computeIfAbsent") ||
+            methodName.equals("computeIfPresent") || methodName.equals("replaceAll") ||
+            methodName.equals("merge") || methodName.equals("peek") ||
+            methodName.equals("doOnNext") || methodName.equals("doOnSuccess") ||
+            methodName.equals("doOnError") || methodName.equals("concatMap") ||
+            methodName.equals("switchMap")) {
+            return true;
+        }
+        
+        return false; 
+    }
+
+    private static boolean isCollectionOrStreamType(org.eclipse.jdt.core.dom.ITypeBinding typeBinding) {
+        String fqn = typeBinding.getErasure().getQualifiedName();
+        if (fqn.startsWith("java.util.stream.") || 
+            fqn.equals("java.lang.Iterable") || 
+            fqn.equals("java.util.Collection") || 
+            fqn.equals("java.util.List") || 
+            fqn.equals("java.util.Set") || 
+            fqn.equals("java.util.Map") || 
+            fqn.equals("java.util.Iterator") || 
+            fqn.equals("java.util.Optional") ||
+            fqn.equals("reactor.core.publisher.Flux") ||
+            fqn.equals("reactor.core.publisher.Mono")) {
+            return true;
+        }
+        
+        for (org.eclipse.jdt.core.dom.ITypeBinding intf : typeBinding.getInterfaces()) {
+            if (isCollectionOrStreamType(intf)) return true;
+        }
+        
+        if (typeBinding.getSuperclass() != null) {
+            return isCollectionOrStreamType(typeBinding.getSuperclass());
+        }
+        
+        return false;
+    }
 }