Extract call-site and link policy helpers to clarify analysis pipeline rules.

Move functional-interface checks, cross-frame Supplier resolution, and transition link resolution into dedicated types so call graph and linker share one DRY implementation.

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

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/enricher/CallChainLinkPolicy.java

@@ -0,0 +1,70 @@
+package click.kamil.springstatemachineexporter.analysis.enricher;
+
+import click.kamil.springstatemachineexporter.analysis.model.CallChain;
+import click.kamil.springstatemachineexporter.analysis.model.EntryPoint;
+import click.kamil.springstatemachineexporter.analysis.model.LinkResolution;
+import click.kamil.springstatemachineexporter.analysis.model.MatchedTransition;
+import click.kamil.springstatemachineexporter.analysis.model.TriggerPoint;
+import click.kamil.springstatemachineexporter.analysis.resolver.MachineEnumCanonicalizer;
+import click.kamil.springstatemachineexporter.analysis.service.ExternalTriggerPolicy;
+import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
+
+import java.util.List;
+
+/**
+ * Central rules for transition linking outcomes derived from existing trigger evidence only.
+ */
+public final class CallChainLinkPolicy {
+
+    private CallChainLinkPolicy() {
+    }
+
+    public static boolean shouldFailClosedOnAmbiguousCallGraphWiden(TriggerPoint trigger) {
+        if (trigger == null || trigger.isExternal() || !trigger.isAmbiguous()) {
+            return false;
+        }
+        List<String> poly = trigger.getPolymorphicEvents();
+        if (poly == null || poly.size() <= 1) {
+            return false;
+        }
+        String event = trigger.getEvent();
+        if (event != null && event.startsWith("ENUM_SET:")) {
+            return true;
+        }
+        return MachineEnumCanonicalizer.isDynamicTriggerExpression(event)
+                && event != null
+                && !event.contains("valueOf");
+    }
+
+    public static LinkResolution resolveLinkResolution(
+            TriggerPoint trigger,
+            List<MatchedTransition> matched,
+            boolean ambiguousSource) {
+        if (trigger != null && trigger.isExternal()) {
+            return LinkResolution.UNRESOLVED_EXTERNAL;
+        }
+        if (ambiguousSource
+                || (trigger != null && trigger.isAmbiguous())
+                || shouldFailClosedOnAmbiguousCallGraphWiden(trigger)) {
+            return LinkResolution.AMBIGUOUS_WIDEN;
+        }
+        if (matched != null && !matched.isEmpty()) {
+            return LinkResolution.RESOLVED;
+        }
+        return LinkResolution.NO_MATCH;
+    }
+
+    public static TriggerPoint applyRestExternalPolicy(TriggerPoint trigger, CallChain chain, CodebaseContext context) {
+        if (trigger == null || chain.getEntryPoint() == null) {
+            return trigger;
+        }
+        EntryPoint entryPoint = chain.getEntryPoint();
+        String entryMethod = entryPoint.getClassName() + "." + entryPoint.getMethodName();
+        boolean external = ExternalTriggerPolicy.isExternalFromSource(
+                entryPoint, trigger, entryMethod, trigger.getEvent(), context);
+        if (external == trigger.isExternal()) {
+            return trigger;
+        }
+        return trigger.toBuilder().external(external).build();
+    }
+}

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/enricher/TransitionLinkerEnricher.java

@@ -9,7 +9,6 @@ import click.kamil.springstatemachineexporter.analysis.model.LinkResolution;
 import click.kamil.springstatemachineexporter.analysis.model.MatchedTransition;
 import click.kamil.springstatemachineexporter.analysis.model.TriggerPoint;
 import click.kamil.springstatemachineexporter.analysis.service.CodebaseIntelligenceProvider;
-import click.kamil.springstatemachineexporter.analysis.service.ExternalTriggerPolicy;
 import click.kamil.springstatemachineexporter.analysis.service.JsonExportContextFactory;
 import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
 import click.kamil.springstatemachineexporter.model.State;
@@ -59,7 +58,7 @@ public class TransitionLinkerEnricher implements AnalysisEnricher {
 
             tp = MachineEnumCanonicalizer.ensureCallChainPolymorphicEvents(
                     tp, machineTypes, context, stateMachineTransitions, true);
-            tp = markRestEndpointExternal(tp, chain, context);
+            tp = CallChainLinkPolicy.applyRestExternalPolicy(tp, chain, context);
 
             if (LifecycleTriggerMarkers.isLifecycle(tp.getEvent())) {
                 updatedChains.add(chain);
@@ -67,7 +66,7 @@ public class TransitionLinkerEnricher implements AnalysisEnricher {
             }
 
             String triggerSource = tp.getSourceState() != null ? simplify(tp.getSourceState()) : null;
-            if (shouldFailClosedOnAmbiguousCallGraphWiden(tp)) {
+            if (CallChainLinkPolicy.shouldFailClosedOnAmbiguousCallGraphWiden(tp)) {
                 updatedChains.add(chain.toBuilder()
                         .triggerPoint(tp)
                         .matchedTransitions(null)
@@ -153,7 +152,7 @@ public class TransitionLinkerEnricher implements AnalysisEnricher {
                 tp = tp.toBuilder().ambiguous(true).build();
             }
 
-            LinkResolution linkResolution = resolveLinkResolution(tp, matched, ambiguousSource);
+            LinkResolution linkResolution = CallChainLinkPolicy.resolveLinkResolution(tp, matched, ambiguousSource);
 
             if (!matched.isEmpty()) {
                 CallChain newChain = chain.toBuilder()
@@ -236,51 +235,4 @@ public class TransitionLinkerEnricher implements AnalysisEnricher {
         return event.fullIdentifier() != null ? event.fullIdentifier() : event.rawName();
     }
 
-    private static LinkResolution resolveLinkResolution(
-            TriggerPoint trigger, List<MatchedTransition> matched, boolean ambiguousSource) {
-        if (trigger != null && trigger.isExternal()) {
-            return LinkResolution.UNRESOLVED_EXTERNAL;
-        }
-        if (ambiguousSource
-                || (trigger != null && trigger.isAmbiguous())
-                || shouldFailClosedOnAmbiguousCallGraphWiden(trigger)) {
-            return LinkResolution.AMBIGUOUS_WIDEN;
-        }
-        if (matched != null && !matched.isEmpty()) {
-            return LinkResolution.RESOLVED;
-        }
-        return LinkResolution.NO_MATCH;
-    }
-
-    private static boolean shouldFailClosedOnAmbiguousCallGraphWiden(TriggerPoint tp) {
-        if (tp == null || tp.isExternal() || !tp.isAmbiguous()) {
-            return false;
-        }
-        List<String> poly = tp.getPolymorphicEvents();
-        if (poly == null || poly.size() <= 1) {
-            return false;
-        }
-        String event = tp.getEvent();
-        if (event != null && event.startsWith("ENUM_SET:")) {
-            return true;
-        }
-        return MachineEnumCanonicalizer.isDynamicTriggerExpression(event)
-                && event != null
-                && !event.contains("valueOf");
-    }
-
-    private TriggerPoint markRestEndpointExternal(TriggerPoint trigger, CallChain chain, CodebaseContext context) {
-        if (trigger == null || chain.getEntryPoint() == null) {
-            return trigger;
-        }
-        EntryPoint entryPoint = chain.getEntryPoint();
-        String entryMethod = entryPoint.getClassName() + "." + entryPoint.getMethodName();
-        boolean external = ExternalTriggerPolicy.isExternalFromSource(
-                entryPoint, trigger, entryMethod, trigger.getEvent(), context);
-        if (external == trigger.isExternal()) {
-            return trigger;
-        }
-        return trigger.toBuilder().external(external).build();
-    }
-
 }

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

@@ -334,7 +334,7 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
                         }
                         if (expectedType != null && paramIndex < edge.getArguments().size()) {
                             String argValue = edge.getArguments().get(paramIndex);
-                            if (!isProvablyResolvedCallSiteArgument(argValue, expectedType)) {
+                            if (!FunctionalInterfaceTypes.isProvablyResolvedCallSiteArgument(argValue, expectedType)) {
                                 String actualType = null;
                                 if (argValue.contains(".") && !argValue.contains("(")) {
                                     String prefix = argValue.substring(0, argValue.lastIndexOf('.'));
@@ -3053,7 +3053,8 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
                 && "get".equals(mi.getName().getIdentifier())
                 && mi.arguments().isEmpty()
                 && !expressionAccessClassifier.isKeyedLookup(mi, scopeMethod)) {
-            List<String> callerFrameConstants = resolveSupplierConstantsFromCallerFrame(mi, path, callGraph, 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)) {
@@ -3486,100 +3487,19 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
         return name.matches("[A-Z_][A-Z0-9_]*");
     }
 
-    private boolean isProvablyResolvedCallSiteArgument(String argValue, String expectedType) {
+    private CallSiteArgumentResolver.ExpressionHooks callSiteHooks() {
-        if (argValue == null || argValue.isBlank() || expectedType == null) {
+        return new CallSiteArgumentResolver.ExpressionHooks() {
-            return false;
+            @Override
-        }
+            public String resolveArgument(Expression expr) {
-        if (!isFunctionalInterfaceType(expectedType)) {
+                return AbstractCallGraphEngine.this.resolveArgument(expr);
-            return false;
-        }
-        if (looksLikeEnumConstant(argValue)) {
-            return true;
-        }
-        return isLambdaArgument(argValue);
             }
 
-    private boolean isLambdaArgument(String argValue) {
+            @Override
-        String trimmed = argValue.trim();
+            public Expression parseExpression(String expr) {
-        return trimmed.contains("->")
+                ASTNode node = parseExpressionString(expr);
-                && (trimmed.startsWith("(") || trimmed.matches("^\\w+\\s*->.*"));
+                return node instanceof Expression e ? e : null;
             }
-
+        };
-    private boolean isFunctionalInterfaceType(String typeName) {
-        if (typeName == null || typeName.isBlank()) {
-            return false;
-        }
-        String simple = typeName.contains(".") ? typeName.substring(typeName.lastIndexOf('.') + 1) : typeName;
-        if (simple.contains("<")) {
-            simple = simple.substring(0, simple.indexOf('<'));
-        }
-        return "Supplier".equals(simple) || "Function".equals(simple) || "Callable".equals(simple);
-    }
-
-    private List<String> resolveSupplierConstantsFromCallerFrame(
-            MethodInvocation getCall,
-            List<String> path,
-            Map<String, List<CallEdge>> callGraph,
-            String scopeMethod) {
-        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 = -1;
-        for (int i = path.size() - 1; i >= 0; i--) {
-            if (scopeMethod.equals(path.get(i))) {
-                scopeIndex = i;
-                break;
-            }
-        }
-        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();
-        }
-        List<CallEdge> edges = callGraph.get(caller);
-        if (edges == null) {
-            return List.of();
-        }
-        for (CallEdge edge : edges) {
-            if (!edge.getTargetMethod().equals(callee) && !pathFinder.isHeuristicMatch(edge.getTargetMethod(), callee)) {
-                continue;
-            }
-            if (paramIndex >= edge.getArguments().size()) {
-                continue;
-            }
-            return extractConstantsFromResolvedCallSiteArgument(edge.getArguments().get(paramIndex));
-        }
-        return List.of();
-    }
-
-    private List<String> extractConstantsFromResolvedCallSiteArgument(String argValue) {
-        if (argValue == null || argValue.isBlank()) {
-            return List.of();
-        }
-        ASTNode parsed = parseExpressionString(argValue);
-        if (!(parsed instanceof Expression expr)) {
-            if (looksLikeEnumConstant(argValue)) {
-                return List.of(argValue);
-            }
-            return List.of();
-        }
-        String resolved = resolveArgument(expr);
-        ASTNode resolvedNode = parseExpressionString(resolved);
-        List<String> constants = new ArrayList<>();
-        if (resolvedNode instanceof Expression resolvedExpr) {
-            constantExtractor.extractConstantsFromExpression(resolvedExpr, constants);
-        }
-        if (constants.isEmpty() && looksLikeEnumConstant(resolved)) {
-            constants.add(resolved);
-        }
-        return constants;
     }
 
     private void trackKeyedMapLookupFlags(

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

@@ -0,0 +1,120 @@
+package click.kamil.springstatemachineexporter.analysis.service;
+
+import click.kamil.springstatemachineexporter.analysis.model.CallEdge;
+import org.eclipse.jdt.core.dom.ASTNode;
+import org.eclipse.jdt.core.dom.Expression;
+import org.eclipse.jdt.core.dom.MethodInvocation;
+import org.eclipse.jdt.core.dom.SimpleName;
+
+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);
+    }
+
+    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) {
+            constantExtractor.extractConstantsFromExpression(resolvedExpr, constants);
+        }
+        if (constants.isEmpty() && FunctionalInterfaceTypes.looksLikeEnumConstant(resolved)) {
+            constants.add(resolved);
+        }
+        return constants;
+    }
+}

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

@@ -0,0 +1,56 @@
+package click.kamil.springstatemachineexporter.analysis.service;
+
+/**
+ * Recognizes JDK functional interface parameters used as delayed event providers
+ * ({@code Supplier.get()}, etc.) in call-graph argument tracing.
+ */
+public final class FunctionalInterfaceTypes {
+
+    private FunctionalInterfaceTypes() {
+    }
+
+    public static boolean isFunctionalInterface(String typeName) {
+        if (typeName == null || typeName.isBlank()) {
+            return false;
+        }
+        String simple = typeName.contains(".") ? typeName.substring(typeName.lastIndexOf('.') + 1) : typeName;
+        if (simple.contains("<")) {
+            simple = simple.substring(0, simple.indexOf('<'));
+        }
+        return "Supplier".equals(simple) || "Function".equals(simple) || "Callable".equals(simple);
+    }
+
+    public static boolean isLambdaArgument(String argValue) {
+        if (argValue == null) {
+            return false;
+        }
+        String trimmed = argValue.trim();
+        return trimmed.contains("->")
+                && (trimmed.startsWith("(") || trimmed.matches("^\\w+\\s*->.*"));
+    }
+
+    /**
+     * When the caller passes a lambda or enum literal into a functional parameter, skip
+     * incompatible-type rejection between {@code Supplier<T>} and resolved {@code T}.
+     */
+    public static boolean isProvablyResolvedCallSiteArgument(String argValue, String expectedType) {
+        if (argValue == null || argValue.isBlank() || expectedType == null) {
+            return false;
+        }
+        if (!isFunctionalInterface(expectedType)) {
+            return false;
+        }
+        if (looksLikeEnumConstant(argValue)) {
+            return true;
+        }
+        return isLambdaArgument(argValue);
+    }
+
+    static boolean looksLikeEnumConstant(String value) {
+        if (value == null || value.isBlank()) {
+            return false;
+        }
+        String name = value.contains(".") ? value.substring(value.lastIndexOf('.') + 1) : value;
+        return name.matches("[A-Z_][A-Z0-9_]*");
+    }
+}

state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/enricher/CallChainLinkPolicyTest.java

@@ -0,0 +1,36 @@
+package click.kamil.springstatemachineexporter.analysis.enricher;
+
+import click.kamil.springstatemachineexporter.analysis.model.LinkResolution;
+import click.kamil.springstatemachineexporter.analysis.model.TriggerPoint;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+class CallChainLinkPolicyTest {
+
+    @Test
+    void shouldFailClosedOnEnumSetAmbiguousWiden() {
+        TriggerPoint trigger = TriggerPoint.builder()
+                .ambiguous(true)
+                .event("ENUM_SET:com.example.OrderEvent.PAY,com.example.OrderEvent.SHIP")
+                .polymorphicEvents(List.of("com.example.OrderEvent.PAY", "com.example.OrderEvent.SHIP"))
+                .build();
+
+        assertThat(CallChainLinkPolicy.shouldFailClosedOnAmbiguousCallGraphWiden(trigger)).isTrue();
+        assertThat(CallChainLinkPolicy.resolveLinkResolution(trigger, List.of(), false))
+                .isEqualTo(LinkResolution.AMBIGUOUS_WIDEN);
+    }
+
+    @Test
+    void shouldMarkExternalTriggersUnresolved() {
+        TriggerPoint trigger = TriggerPoint.builder()
+                .external(true)
+                .event("eventString")
+                .build();
+
+        assertThat(CallChainLinkPolicy.resolveLinkResolution(trigger, List.of(), false))
+                .isEqualTo(LinkResolution.UNRESOLVED_EXTERNAL);
+    }
+}

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

@@ -0,0 +1,28 @@
+package click.kamil.springstatemachineexporter.analysis.service;
+
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+class FunctionalInterfaceTypesTest {
+
+    @Test
+    void shouldRecognizeSupplierParameter() {
+        assertThat(FunctionalInterfaceTypes.isFunctionalInterface("java.util.function.Supplier"))
+                .isTrue();
+    }
+
+    @Test
+    void shouldAcceptLambdaIntoSupplierParameter() {
+        assertThat(FunctionalInterfaceTypes.isProvablyResolvedCallSiteArgument(
+                "() -> OrderEvent.CANCEL", "java.util.function.Supplier"))
+                .isTrue();
+    }
+
+    @Test
+    void shouldRejectNonFunctionalParameter() {
+        assertThat(FunctionalInterfaceTypes.isProvablyResolvedCallSiteArgument(
+                "OrderEvent.CANCEL", "com.example.OrderEvent"))
+                .isFalse();
+    }
+}