Remove CallSiteArgumentResolver; route functional args through dataflow.

Delete the parallel Supplier/Runnable bridging layer and resolve functional call-site arguments via VariableTracer, PathBindingEvaluator, and JdtDataFlowModel instead.

Co-authored-by: Cursor <cursoragent@cursor.com>

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

@@ -86,6 +86,7 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
 
         // Wire up cross dependencies
         this.variableTracer.setConstantExtractor(this.constantExtractor);
+        this.variableTracer.setTypeResolver(this.typeResolver);
         this.constantExtractor.setVariableTracer(this.variableTracer);
         this.constantExtractor.setConstructorAnalyzer(this.constructorAnalyzer);
         this.constructorAnalyzer.setVariableTracer(this.variableTracer);
@@ -3101,35 +3102,6 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
 
     private void addConstantsFromTracedExpression(
             Expression expr, List<String> constants, List<String> path, Map<String, List<CallEdge>> callGraph, String scopeMethod) {
-        if (expr instanceof MethodInvocation mi
-                && "get".equals(mi.getName().getIdentifier())
-                && mi.arguments().isEmpty()
-                && !expressionAccessClassifier.isKeyedLookup(mi, scopeMethod)) {
-            List<String> callerFrameConstants = CallSiteArgumentResolver.resolveFunctionalGetFromCallerFrame(
-                    mi, path, callGraph, scopeMethod, typeResolver, pathFinder, constantExtractor, callSiteHooks());
-            if (!callerFrameConstants.isEmpty()) {
-                for (String constant : callerFrameConstants) {
-                    if (!constants.contains(constant)) {
-                        constants.add(constant);
-                    }
-                }
-                return;
-            }
-        }
-        if (expr instanceof MethodInvocation mi
-                && "run".equals(mi.getName().getIdentifier())
-                && mi.arguments().isEmpty()) {
-            List<String> callerFrameConstants = CallSiteArgumentResolver.resolveFunctionalRunFromCallerFrame(
-                    mi, path, callGraph, scopeMethod, typeResolver, pathFinder, constantExtractor, callSiteHooks());
-            if (!callerFrameConstants.isEmpty()) {
-                for (String constant : callerFrameConstants) {
-                    if (!constants.contains(constant)) {
-                        constants.add(constant);
-                    }
-                }
-                return;
-            }
-        }
         List<Expression> traced = variableTracer.traceVariableAll(expr);
         for (Expression tracedExpr : traced) {
             addConstantsFromSingleExpression(tracedExpr, constants, path, callGraph, scopeMethod);
@@ -3169,11 +3141,6 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
         }
         Map<String, String> bindings = pathBindingEvaluator.traceBindings(path, callGraph, pathFinder);
         if (bindings.isEmpty()) {
-            List<String> runnableResolved = CallSiteArgumentResolver.resolveRunnableLambdaFromPath(
-                    path, callGraph, typeResolver, context, pathFinder, constantExtractor, callSiteHooks());
-            if (!runnableResolved.isEmpty()) {
-                return runnableResolved;
-            }
             if (hasMultiClassPolymorphicPath(path)
                     && polymorphicEvents.stream().noneMatch(this::looksLikeEnumConstant)) {
                 List<String> callerArgumentResolved = resolvePolymorphicEventsFromCallerArgument(path);
@@ -3558,21 +3525,6 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
         return name.matches("[A-Z_][A-Z0-9_]*");
     }
 
-    private CallSiteArgumentResolver.ExpressionHooks callSiteHooks() {
-        return new CallSiteArgumentResolver.ExpressionHooks() {
-            @Override
-            public String resolveArgument(Expression expr) {
-                return AbstractCallGraphEngine.this.resolveArgument(expr);
-            }
-
-            @Override
-            public Expression parseExpression(String expr) {
-                ASTNode node = parseExpressionString(expr);
-                return node instanceof Expression e ? e : null;
-            }
-        };
-    }
-
     private void trackKeyedMapLookupFlags(
             Expression expr,
             List<String> path,
@@ -3612,10 +3564,4 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
             results.add(parsed);
         }
     }
-
-    List<String> resolveRunnableLambdaConstantsForPath(
-            List<String> path, Map<String, List<CallEdge>> callGraph) {
-        return CallSiteArgumentResolver.resolveRunnableLambdaFromPath(
-                path, callGraph, typeResolver, context, pathFinder, constantExtractor, callSiteHooks());
-    }
 }

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

@@ -1,416 +0,0 @@
-package click.kamil.springstatemachineexporter.analysis.service;
-
-import click.kamil.springstatemachineexporter.analysis.model.CallEdge;
-import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
-import click.kamil.springstatemachineexporter.ast.common.JdtDataFlowModel;
-import org.eclipse.jdt.core.dom.ASTNode;
-import org.eclipse.jdt.core.dom.ASTVisitor;
-import org.eclipse.jdt.core.dom.Expression;
-import org.eclipse.jdt.core.dom.ExpressionStatement;
-import org.eclipse.jdt.core.dom.IMethodBinding;
-import org.eclipse.jdt.core.dom.LambdaExpression;
-import org.eclipse.jdt.core.dom.MethodDeclaration;
-import org.eclipse.jdt.core.dom.MethodInvocation;
-import org.eclipse.jdt.core.dom.SimpleName;
-import org.eclipse.jdt.core.dom.TypeDeclaration;
-import org.eclipse.jdt.core.dom.Block;
-
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Map;
-
-/**
- * Resolves call-site arguments across stack frames (e.g. {@code Supplier.get()} → caller lambda).
- */
-public final class CallSiteArgumentResolver {
-
-    public interface ExpressionHooks {
-        /** Unwrap lambda / enum literal argument expressions to a stable string form. */
-        String resolveArgument(Expression expr);
-
-        /** Parse a stringified expression back to AST for constant extraction. */
-        Expression parseExpression(String expr);
-    }
-
-    private CallSiteArgumentResolver() {
-    }
-
-    /**
-     * When {@code supplierParam.get()} is seen inside {@code scopeMethod}, walk to the caller frame
-     * and extract enum/constants from the argument bound to {@code supplierParam}.
-     */
-    public static List<String> resolveFunctionalGetFromCallerFrame(
-            MethodInvocation getCall,
-            List<String> path,
-            Map<String, List<CallEdge>> callGraph,
-            String scopeMethod,
-            TypeResolver typeResolver,
-            CallGraphPathFinder pathFinder,
-            ConstantExtractor constantExtractor,
-            ExpressionHooks hooks) {
-        if (path == null || path.size() < 2 || scopeMethod == null || getCall.getExpression() == null) {
-            return List.of();
-        }
-        if (!(getCall.getExpression() instanceof SimpleName supplierParam)) {
-            return List.of();
-        }
-        int scopeIndex = indexOfMethod(path, scopeMethod);
-        if (scopeIndex <= 0) {
-            return List.of();
-        }
-        String callee = path.get(scopeIndex);
-        String caller = path.get(scopeIndex - 1);
-        int paramIndex = typeResolver.getParameterIndex(callee, supplierParam.getIdentifier());
-        if (paramIndex < 0) {
-            return List.of();
-        }
-        String argValue = findCallSiteArgument(caller, callee, paramIndex, callGraph, pathFinder);
-        if (argValue == null) {
-            return List.of();
-        }
-        return extractConstantsFromCallSiteArgument(argValue, constantExtractor, hooks);
-    }
-
-    /**
-     * When {@code task.run()} is seen inside {@code scopeMethod}, walk to the caller frame and extract
-     * enum/constants from the {@code Runnable} argument bound at the {@code execute(Runnable)} call site.
-     */
-    public static List<String> resolveFunctionalRunFromCallerFrame(
-            MethodInvocation runCall,
-            List<String> path,
-            Map<String, List<CallEdge>> callGraph,
-            String scopeMethod,
-            TypeResolver typeResolver,
-            CallGraphPathFinder pathFinder,
-            ConstantExtractor constantExtractor,
-            ExpressionHooks hooks) {
-        if (path == null || path.size() < 2 || scopeMethod == null || runCall.getExpression() == null) {
-            return List.of();
-        }
-        if (!(runCall.getExpression() instanceof SimpleName runnableParam)) {
-            return List.of();
-        }
-        if (!"run".equals(runCall.getName().getIdentifier()) || !runCall.arguments().isEmpty()) {
-            return List.of();
-        }
-        int scopeIndex = indexOfMethod(path, scopeMethod);
-        if (scopeIndex <= 0) {
-            return List.of();
-        }
-        String callee = path.get(scopeIndex);
-        String caller = path.get(scopeIndex - 1);
-        int paramIndex = typeResolver.getParameterIndex(callee, runnableParam.getIdentifier());
-        if (paramIndex < 0) {
-            paramIndex = findRunnableParameterIndex(callee, typeResolver);
-        }
-        if (paramIndex < 0) {
-            return List.of();
-        }
-        String argValue = findCallSiteArgument(caller, callee, paramIndex, callGraph, pathFinder);
-        if (argValue == null) {
-            return List.of();
-        }
-        return extractConstantsFromCallSiteArgument(argValue, constantExtractor, hooks);
-    }
-
-    /**
-     * When the call path contains {@code java.lang.Runnable.run}, resolve the {@code Runnable} lambda
-     * passed from the frame two hops earlier (e.g. controller → executor.execute → Runnable.run).
-     */
-    public static List<String> resolveRunnableLambdaFromPath(
-            List<String> path,
-            Map<String, List<CallEdge>> callGraph,
-            TypeResolver typeResolver,
-            CodebaseContext context,
-            CallGraphPathFinder pathFinder,
-            ConstantExtractor constantExtractor,
-            ExpressionHooks hooks) {
-        if (path == null || path.size() < 3) {
-            return List.of();
-        }
-        for (int i = 1; i < path.size(); i++) {
-            if (!isRunnableRunMethod(path.get(i))) {
-                continue;
-            }
-            String executeMethod = path.get(i - 1);
-            String callerMethod = path.get(i - 2);
-            int paramIndex = findRunnableParameterIndex(executeMethod, typeResolver);
-            if (paramIndex < 0) {
-                paramIndex = 0;
-            }
-            String argValue = findCallSiteArgument(callerMethod, executeMethod, paramIndex, callGraph, pathFinder);
-            List<String> constants = extractRunnableArgumentConstants(
-                    callerMethod, executeMethod, paramIndex, argValue, context, typeResolver, constantExtractor, hooks);
-            if (!constants.isEmpty()) {
-                return constants;
-            }
-        }
-        return List.of();
-    }
-
-    private static boolean isRunnableRunMethod(String methodFqn) {
-        return "java.lang.Runnable.run".equals(methodFqn) || "Runnable.run".equals(methodFqn);
-    }
-
-    private static int findRunnableParameterIndex(String methodFqn, TypeResolver typeResolver) {
-        if (methodFqn == null || typeResolver == null) {
-            return -1;
-        }
-        for (int i = 0; i < 8; i++) {
-            String paramType = typeResolver.getParameterType(methodFqn, i);
-            if (paramType == null) {
-                break;
-            }
-            if (FunctionalInterfaceTypes.isFunctionalInterface(paramType)) {
-                return i;
-            }
-        }
-        return -1;
-    }
-
-    public static String findCallSiteArgument(
-            String caller,
-            String callee,
-            int paramIndex,
-            Map<String, List<CallEdge>> callGraph,
-            CallGraphPathFinder pathFinder) {
-        List<CallEdge> edges = callGraph.get(caller);
-        if (edges == null) {
-            return null;
-        }
-        for (CallEdge edge : edges) {
-            if (!edge.getTargetMethod().equals(callee) && !pathFinder.isHeuristicMatch(edge.getTargetMethod(), callee)) {
-                continue;
-            }
-            if (paramIndex < edge.getArguments().size()) {
-                return edge.getArguments().get(paramIndex);
-            }
-        }
-        return null;
-    }
-
-    public static int indexOfMethod(List<String> path, String methodFqn) {
-        for (int i = path.size() - 1; i >= 0; i--) {
-            if (methodFqn.equals(path.get(i))) {
-                return i;
-            }
-        }
-        return -1;
-    }
-
-    public static List<String> extractConstantsFromCallSiteArgument(
-            String argValue,
-            ConstantExtractor constantExtractor,
-            ExpressionHooks hooks) {
-        if (argValue == null || argValue.isBlank()) {
-            return List.of();
-        }
-        Expression parsed = hooks.parseExpression(argValue);
-        if (parsed == null) {
-            if (FunctionalInterfaceTypes.looksLikeEnumConstant(argValue)) {
-                return List.of(argValue);
-            }
-            return List.of();
-        }
-        String resolved = hooks.resolveArgument(parsed);
-        Expression resolvedExpr = hooks.parseExpression(resolved);
-        List<String> constants = new ArrayList<>();
-        if (resolvedExpr != null) {
-            extractSideEffectConstantsFromLambda(resolvedExpr, constantExtractor, constants);
-            if (constants.isEmpty()) {
-                constantExtractor.extractConstantsFromExpression(resolvedExpr, constants);
-            }
-        }
-        if (constants.isEmpty() && parsed instanceof LambdaExpression lambda) {
-            extractSideEffectConstantsFromLambda(lambda, constantExtractor, constants);
-        }
-        if (constants.isEmpty() && FunctionalInterfaceTypes.looksLikeEnumConstant(resolved)) {
-            constants.add(resolved);
-        }
-        return constants;
-    }
-
-    private static void extractSideEffectConstantsFromLambda(
-            Expression expression,
-            ConstantExtractor constantExtractor,
-            List<String> constants) {
-        if (!(expression instanceof LambdaExpression lambda)) {
-            return;
-        }
-        ASTNode body = lambda.getBody();
-        if (body instanceof Expression bodyExpr) {
-            constantExtractor.extractConstantsFromExpression(bodyExpr, constants);
-            return;
-        }
-        if (!(body instanceof Block block)) {
-            return;
-        }
-        for (Object statement : block.statements()) {
-            if (statement instanceof ExpressionStatement expressionStatement) {
-                constantExtractor.extractConstantsFromExpression(expressionStatement.getExpression(), constants);
-            } else if (statement instanceof org.eclipse.jdt.core.dom.ReturnStatement returnStatement
-                    && returnStatement.getExpression() != null) {
-                constantExtractor.extractConstantsFromExpression(returnStatement.getExpression(), constants);
-            }
-        }
-    }
-
-    private static List<String> extractRunnableArgumentConstants(
-            String callerMethod,
-            String executeMethod,
-            int paramIndex,
-            String argValue,
-            CodebaseContext context,
-            TypeResolver typeResolver,
-            ConstantExtractor constantExtractor,
-            ExpressionHooks hooks) {
-        List<String> constants = new ArrayList<>();
-        Expression sourceArg = findCallSiteArgumentExpression(callerMethod, executeMethod, paramIndex, context);
-        if (sourceArg != null) {
-            extractConstantsViaDataflow(sourceArg, context, constantExtractor, constants);
-        }
-        if (constants.isEmpty() && argValue != null) {
-            Expression parsed = hooks.parseExpression(argValue);
-            if (parsed != null) {
-                extractConstantsViaDataflow(parsed, context, constantExtractor, constants);
-                if (constants.isEmpty() && parsed instanceof MethodInvocation sideEffect) {
-                    extractSideEffectConstantsFromParsedInvocation(
-                            sideEffect, callerMethod, constantExtractor, constants, context, typeResolver);
-                }
-            }
-            if (constants.isEmpty()) {
-                constants.addAll(extractConstantsFromCallSiteArgument(argValue, constantExtractor, hooks));
-            }
-        }
-        return constants;
-    }
-
-    private static void extractSideEffectConstantsFromParsedInvocation(
-            MethodInvocation invocation,
-            String callerMethodFqn,
-            ConstantExtractor constantExtractor,
-            List<String> constants,
-            CodebaseContext context,
-            TypeResolver typeResolver) {
-        if (context == null || typeResolver == null || callerMethodFqn == null || !callerMethodFqn.contains(".")) {
-            return;
-        }
-        String callerClass = callerMethodFqn.substring(0, callerMethodFqn.lastIndexOf('.'));
-        TypeDeclaration callerType = context.getTypeDeclaration(callerClass);
-        if (callerType == null || !(invocation.getExpression() instanceof SimpleName receiver)) {
-            return;
-        }
-        String receiverFqn = null;
-        for (Object field : callerType.getFields()) {
-            if (field instanceof org.eclipse.jdt.core.dom.FieldDeclaration fieldDeclaration) {
-                for (Object fragmentObj : fieldDeclaration.fragments()) {
-                    if (fragmentObj instanceof org.eclipse.jdt.core.dom.VariableDeclarationFragment fragment
-                            && receiver.getIdentifier().equals(fragment.getName().getIdentifier())) {
-                        receiverFqn = typeResolver.resolveTypeToFqn(fieldDeclaration.getType(), callerType.getRoot());
-                        break;
-                    }
-                }
-            }
-            if (receiverFqn != null) {
-                break;
-            }
-        }
-        if (receiverFqn == null) {
-            return;
-        }
-        MethodDeclaration method = context.findMethodDeclaration(
-                context.getTypeDeclaration(receiverFqn),
-                invocation.getName().getIdentifier(),
-                true);
-        if (method == null || method.getBody() == null) {
-            return;
-        }
-        for (Object statement : method.getBody().statements()) {
-            if (statement instanceof ExpressionStatement expressionStatement
-                    && expressionStatement.getExpression() instanceof MethodInvocation nested) {
-                for (Object arg : nested.arguments()) {
-                    constantExtractor.extractConstantsFromExpression((Expression) arg, constants);
-                }
-            }
-        }
-    }
-
-    private static void extractConstantsViaDataflow(
-            Expression expression,
-            CodebaseContext context,
-            ConstantExtractor constantExtractor,
-            List<String> constants) {
-        Expression effective = unwrapLambdaBody(expression);
-        if (effective != expression) {
-            extractConstantsViaDataflow(effective, context, constantExtractor, constants);
-            if (!constants.isEmpty()) {
-                return;
-            }
-        }
-        JdtDataFlowModel dataFlow = new JdtDataFlowModel(context);
-        for (Expression def : dataFlow.getReachingDefinitions(effective)) {
-            constantExtractor.extractConstantsFromExpression(def, constants);
-        }
-        if (constants.isEmpty()) {
-            constantExtractor.extractConstantsFromExpression(effective, constants);
-        }
-    }
-
-    private static Expression unwrapLambdaBody(Expression expression) {
-        if (!(expression instanceof LambdaExpression lambda)) {
-            return expression;
-        }
-        ASTNode body = lambda.getBody();
-        if (body instanceof Expression bodyExpr) {
-            return bodyExpr;
-        }
-        return expression;
-    }
-
-    static Expression findCallSiteArgumentExpression(
-            String caller,
-            String callee,
-            int paramIndex,
-            CodebaseContext context) {
-        if (caller == null || callee == null || paramIndex < 0 || context == null
-                || !caller.contains(".") || !callee.contains(".")) {
-            return null;
-        }
-        String callerClass = caller.substring(0, caller.lastIndexOf('.'));
-        String callerMethod = caller.substring(caller.lastIndexOf('.') + 1);
-        String calleeClass = callee.substring(0, callee.lastIndexOf('.'));
-        String calleeMethod = callee.substring(callee.lastIndexOf('.') + 1);
-        TypeDeclaration callerType = context.getTypeDeclaration(callerClass);
-        if (callerType == null) {
-            return null;
-        }
-        MethodDeclaration method = context.findMethodDeclaration(callerType, callerMethod, true);
-        if (method == null || method.getBody() == null) {
-            return null;
-        }
-        final Expression[] found = new Expression[1];
-        method.getBody().accept(new ASTVisitor() {
-            @Override
-            public boolean visit(MethodInvocation node) {
-                if (found[0] != null || !calleeMethod.equals(node.getName().getIdentifier())) {
-                    return super.visit(node);
-                }
-                IMethodBinding binding = node.resolveMethodBinding();
-                if (binding != null && binding.getDeclaringClass() != null) {
-                    String declaringFqn = binding.getDeclaringClass().getErasure().getQualifiedName();
-                    if (!declaringFqn.equals(calleeClass) && !declaringFqn.endsWith("." + calleeClass)) {
-                        String calleeSimple = calleeClass.substring(calleeClass.lastIndexOf('.') + 1);
-                        if (!declaringFqn.endsWith("." + calleeSimple)) {
-                            return super.visit(node);
-                        }
-                    }
-                }
-                if (paramIndex < node.arguments().size()) {
-                    found[0] = (Expression) node.arguments().get(paramIndex);
-                }
-                return super.visit(node);
-            }
-        });
-        return found[0];
-    }
-}

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

@@ -143,6 +143,19 @@ public class PathBindingEvaluator {
             Map<String, List<CallEdge>> callGraph,
             CallGraphPathFinder pathFinder,
             Map<String, String> bindings) {
+        if (VariableTracer.isRunnableRunMethod(target) && pathIndex >= 2) {
+            int paramIndex = variableTracer.findFunctionalParameterIndex(caller);
+            if (paramIndex < 0) {
+                paramIndex = 0;
+            }
+            List<String> constants = variableTracer.resolveConstantsFromCallSiteArgument(
+                    path.get(pathIndex - 2), caller, paramIndex);
+            for (String constant : constants) {
+                if (constant != null && !constant.isBlank()) {
+                    bindings.putIfAbsent(constant, constant);
+                }
+            }
+        }
         Map<String, String> paramValues = variableTracer.buildParameterValuesMap(caller, target, callGraph, path, pathIndex);
         mergeResolvedBindings(bindings, paramValues, caller);
 
@@ -153,6 +166,14 @@ public class PathBindingEvaluator {
         for (int i = 0; i < edge.getArguments().size() && i < targetParams.size(); i++) {
             String arg = edge.getArguments().get(i);
             String paramName = targetParams.get(i);
+            if (typeResolver != null
+                    && FunctionalInterfaceTypes.isFunctionalInterface(typeResolver.getParameterType(target, i))) {
+                List<String> constants = variableTracer.resolveConstantsFromCallSiteArgument(caller, target, i);
+                if (constants.size() == 1) {
+                    bindings.put(paramName, constants.get(0));
+                    continue;
+                }
+            }
             String resolved = resolveBindingValue(caller, arg, bindings);
             if (shouldStoreBinding(paramName, resolved)) {
                 bindings.put(paramName, resolved);

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

@@ -18,6 +18,7 @@ public class VariableTracer {
     private final CodebaseContext context;
     private final ConstantResolver constantResolver;
     private ConstantExtractor constantExtractor;
+    private TypeResolver typeResolver;
 
     private final click.kamil.springstatemachineexporter.ast.common.DataFlowModel dataFlowModel;
     private final Map<String, String> variableDeclaredTypeCache = new HashMap<>();
@@ -34,6 +35,10 @@ public class VariableTracer {
         this.constantExtractor = constantExtractor;
     }
 
+    public void setTypeResolver(TypeResolver typeResolver) {
+        this.typeResolver = typeResolver;
+    }
+
     public void clearAnalysisCaches() {
         variableDeclaredTypeCache.clear();
         if (dataFlowModel instanceof JdtDataFlowModel jdtDataFlowModel) {
@@ -537,6 +542,13 @@ public class VariableTracer {
                 for (int j = 0; j < minSize; j++) {
                     String argValue = args.get(j);
                     String resolvedArgValue = resolveArgumentValue(argValue, caller, path, pathIndex, callGraph);
+                    if (typeResolver != null
+                            && FunctionalInterfaceTypes.isFunctionalInterface(typeResolver.getParameterType(target, j))) {
+                        List<String> constants = resolveConstantsFromCallSiteArgument(caller, target, j);
+                        if (constants.size() == 1) {
+                            resolvedArgValue = constants.get(0);
+                        }
+                    }
                     paramValues.put(paramNames.get(j), resolvedArgValue);
                 }
                 break;
@@ -698,4 +710,110 @@ public class VariableTracer {
         }
         return null;
     }
+
+    /**
+     * Resolves enum/constants passed at a call site using source AST + {@link JdtDataFlowModel}.
+     * Used for {@code Supplier}, {@code Runnable}, and similar functional parameters.
+     */
+    public List<String> resolveConstantsFromCallSiteArgument(String caller, String callee, int paramIndex) {
+        if (constantExtractor == null) {
+            return List.of();
+        }
+        List<String> constants = new ArrayList<>();
+        Expression sourceArg = findCallSiteArgumentExpression(caller, callee, paramIndex);
+        if (sourceArg != null) {
+            resolveConstantsFromExpressionViaDataflow(sourceArg, constants);
+        }
+        return constants;
+    }
+
+    public int findFunctionalParameterIndex(String methodFqn) {
+        if (methodFqn == null || typeResolver == null) {
+            return -1;
+        }
+        for (int i = 0; i < 8; i++) {
+            String paramType = typeResolver.getParameterType(methodFqn, i);
+            if (paramType == null) {
+                break;
+            }
+            if (FunctionalInterfaceTypes.isFunctionalInterface(paramType)) {
+                return i;
+            }
+        }
+        return -1;
+    }
+
+    public static boolean isRunnableRunMethod(String methodFqn) {
+        return "java.lang.Runnable.run".equals(methodFqn) || "Runnable.run".equals(methodFqn);
+    }
+
+    private void resolveConstantsFromExpressionViaDataflow(Expression expression, List<String> constants) {
+        Expression effective = unwrapLambdaBody(expression);
+        if (effective != expression) {
+            resolveConstantsFromExpressionViaDataflow(effective, constants);
+            if (!constants.isEmpty()) {
+                return;
+            }
+        }
+        for (Expression def : dataFlowModel.getReachingDefinitions(effective)) {
+            constantExtractor.extractConstantsFromExpression(def, constants);
+        }
+        if (constants.isEmpty()) {
+            constantExtractor.extractConstantsFromExpression(effective, constants);
+        }
+    }
+
+    private static Expression unwrapLambdaBody(Expression expression) {
+        if (!(expression instanceof LambdaExpression lambda)) {
+            return expression;
+        }
+        ASTNode body = lambda.getBody();
+        if (body instanceof Expression bodyExpr) {
+            return bodyExpr;
+        }
+        return expression;
+    }
+
+    Expression findCallSiteArgumentExpression(String caller, String callee, int paramIndex) {
+        if (caller == null || callee == null || paramIndex < 0
+                || !caller.contains(".") || !callee.contains(".")) {
+            return null;
+        }
+        String callerClass = caller.substring(0, caller.lastIndexOf('.'));
+        String callerMethod = caller.substring(caller.lastIndexOf('.') + 1);
+        String calleeMethod = callee.substring(callee.lastIndexOf('.') + 1);
+        TypeDeclaration callerType = context.getTypeDeclaration(callerClass);
+        if (callerType == null) {
+            return null;
+        }
+        MethodDeclaration method = context.findMethodDeclaration(callerType, callerMethod, true);
+        if (method == null || method.getBody() == null) {
+            return null;
+        }
+        final Expression[] found = new Expression[1];
+        method.getBody().accept(new ASTVisitor() {
+            @Override
+            public boolean visit(MethodInvocation node) {
+                if (found[0] != null || !calleeMethod.equals(node.getName().getIdentifier())) {
+                    return super.visit(node);
+                }
+                IMethodBinding binding = node.resolveMethodBinding();
+                if (binding != null && binding.getDeclaringClass() != null) {
+                    String declaringFqn = binding.getDeclaringClass().getErasure().getQualifiedName();
+                    String calleeClass = callee.substring(0, callee.lastIndexOf('.'));
+                    if (!declaringFqn.equals(calleeClass) && !declaringFqn.endsWith("." + calleeClass)) {
+                        String calleeSimple = calleeClass.substring(calleeClass.lastIndexOf('.') + 1);
+                        if (!declaringFqn.endsWith("." + calleeSimple)) {
+                            return super.visit(node);
+                        }
+                    }
+                }
+                if (paramIndex < node.arguments().size()) {
+                    found[0] = (Expression) node.arguments().get(paramIndex);
+                }
+                return super.visit(node);
+            }
+        });
+        return found[0];
+    }
 }

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

@@ -200,6 +200,23 @@ public class JdtDataFlowModel implements DataFlowModel {
                 return resolved;
             }
 
+            if ("run".equals(mName) && mi.arguments().isEmpty() && isRunnableReceiver(mi.getExpression())) {
+                List<Expression> functional = resolveFunctionalInterfaceInvocation(
+                        mi, visited, paramBindings, instanceFieldBindings, depth);
+                if (!functional.isEmpty()) {
+                    visited.remove(expr);
+                    return functional;
+                }
+            }
+            if ("get".equals(mName) && mi.arguments().isEmpty() && isSupplierLikeReceiver(mi.getExpression())) {
+                List<Expression> functional = resolveFunctionalInterfaceInvocation(
+                        mi, visited, paramBindings, instanceFieldBindings, depth);
+                if (!functional.isEmpty()) {
+                    visited.remove(expr);
+                    return functional;
+                }
+            }
+
             List<Expression> inlined = tryInlineAccessorInvocation(
                     mi, new HashSet<>(visited), paramBindings, instanceFieldBindings, depth);
             if (!inlined.isEmpty()) {
@@ -1466,6 +1483,79 @@ public class JdtDataFlowModel implements DataFlowModel {
         return findMethodDeclarationInType(mb.getDeclaringClass(), mb);
     }
 
+    private static boolean isSupplierLikeReceiver(Expression receiver) {
+        if (receiver == null) {
+            return false;
+        }
+        ITypeBinding typeBinding = receiver.resolveTypeBinding();
+        if (typeBinding == null) {
+            return false;
+        }
+        String fqn = typeBinding.getErasure().getQualifiedName();
+        if (fqn == null) {
+            return false;
+        }
+        return fqn.contains("Supplier") || fqn.contains("Function") || fqn.contains("Callable");
+    }
+
+    private static boolean isRunnableReceiver(Expression receiver) {
+        if (receiver == null) {
+            return false;
+        }
+        ITypeBinding typeBinding = receiver.resolveTypeBinding();
+        if (typeBinding == null) {
+            return false;
+        }
+        String fqn = typeBinding.getErasure().getQualifiedName();
+        return "java.lang.Runnable".equals(fqn) || fqn.endsWith(".Runnable");
+    }
+
+    private List<Expression> resolveFunctionalInterfaceInvocation(
+            MethodInvocation mi,
+            Set<ASTNode> visited,
+            Map<IVariableBinding, Expression> paramBindings,
+            Map<IVariableBinding, Expression> instanceFieldBindings,
+            int depth) {
+        Expression receiver = mi.getExpression();
+        if (receiver instanceof SimpleName sn) {
+            IBinding binding = sn.resolveBinding();
+            if (binding instanceof IVariableBinding variableBinding && paramBindings.containsKey(variableBinding)) {
+                return resolveFunctionalBoundExpression(
+                        paramBindings.get(variableBinding), visited, paramBindings, instanceFieldBindings, depth + 1);
+            }
+        }
+        List<Expression> receiverDefs = getReachingDefinitions(receiver, visited, paramBindings, instanceFieldBindings, depth + 1);
+        List<Expression> results = new ArrayList<>();
+        for (Expression receiverDef : receiverDefs) {
+            results.addAll(resolveFunctionalBoundExpression(
+                    receiverDef, visited, paramBindings, instanceFieldBindings, depth + 1));
+        }
+        return results;
+    }
+
+    private List<Expression> resolveFunctionalBoundExpression(
+            Expression bound,
+            Set<ASTNode> visited,
+            Map<IVariableBinding, Expression> paramBindings,
+            Map<IVariableBinding, Expression> instanceFieldBindings,
+            int depth) {
+        if (bound == null) {
+            return List.of();
+        }
+        if (bound instanceof LambdaExpression lambda) {
+            ASTNode body = lambda.getBody();
+            if (body instanceof Expression bodyExpr) {
+                return getReachingDefinitions(bodyExpr, visited, paramBindings, instanceFieldBindings, depth);
+            }
+            if (body instanceof Block block) {
+                List<Expression> results = new ArrayList<>();
+                collectSideEffectExpressions(block, results, visited, paramBindings, instanceFieldBindings, depth);
+                return results;
+            }
+        }
+        return getReachingDefinitions(bound, visited, paramBindings, instanceFieldBindings, depth);
+    }
+
     private void collectSideEffectExpressions(
             Block body,
             List<Expression> results,

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

@@ -103,83 +103,4 @@ class RunnableLambdaDispatchTest {
         assertThat(chains.get(0).getTriggerPoint().getPolymorphicEvents())
                 .containsExactly("OrderCommand.PAY");
     }
-
-    @Test
-    void shouldResolveRunnableLambdaFromExecuteCallSite() throws IOException {
-        String controllerSource = """
-                package com.example.web;
-
-                import com.example.service.AbstractOrderService;
-                import com.example.executor.TaskExecutor;
-
-                public class OrderController {
-                    private final TaskExecutor executor;
-                    private final AbstractOrderService orderService;
-
-                    public OrderController(TaskExecutor executor, AbstractOrderService orderService) {
-                        this.executor = executor;
-                        this.orderService = orderService;
-                    }
-
-                    public void process() {
-                        executor.execute(() -> orderService.fireEvent());
-                    }
-                }
-                """;
-
-        String executorSource = """
-                package com.example.executor;
-
-                public class TaskExecutor {
-                    public void execute(Runnable task) {
-                        task.run();
-                    }
-                }
-                """;
-
-        String serviceSource = """
-                package com.example.service;
-
-                import com.example.domain.OrderCommand;
-
-                public abstract class AbstractOrderService {
-                    public void fireEvent() {
-                        sendEvent(OrderCommand.PAY);
-                    }
-
-                    protected void sendEvent(OrderCommand event) {}
-                }
-                """;
-
-        String enumSource = """
-                package com.example.domain;
-
-                public enum OrderCommand { PAY, SHIP, META }
-                """;
-
-        Path tempDir = Files.createTempDirectory("callgraph_runnable_resolver");
-        Files.writeString(tempDir.resolve("OrderController.java"), controllerSource);
-        Files.writeString(tempDir.resolve("TaskExecutor.java"), executorSource);
-        Files.writeString(tempDir.resolve("AbstractOrderService.java"), serviceSource);
-        Files.writeString(tempDir.resolve("OrderCommand.java"), enumSource);
-
-        CodebaseContext context = new CodebaseContext();
-        context.scan(tempDir);
-
-        HeuristicCallGraphEngine engine = new HeuristicCallGraphEngine(context);
-        var graph = engine.buildCallGraph();
-
-        assertThat(graph.get("com.example.executor.TaskExecutor.execute"))
-                .extracting(click.kamil.springstatemachineexporter.analysis.model.CallEdge::getTargetMethod)
-                .anyMatch(target -> target.endsWith("Runnable.run"));
-
-        List<String> path = List.of(
-                "com.example.web.OrderController.process",
-                "com.example.executor.TaskExecutor.execute",
-                graph.get("com.example.executor.TaskExecutor.execute").get(0).getTargetMethod(),
-                "com.example.service.AbstractOrderService.fireEvent");
-
-        List<String> resolved = engine.resolveRunnableLambdaConstantsForPath(path, graph);
-        assertThat(resolved).containsExactly("OrderCommand.PAY");
-    }
 }