Align enricher pipeline with context-aware enum canonicalization

Use label-only canonicalization for shared-infrastructure triggers, resolve valueOf enum types from JDT bindings when available, and allow trusted valueOf widens using the machine event type when the trigger omits type FQNs.

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

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

@@ -20,6 +20,12 @@ public final class CallChainLinkPolicy {
     }
 
     public static boolean shouldFailClosedOnAmbiguousCallGraphWiden(TriggerPoint trigger) {
+        return shouldFailClosedOnAmbiguousCallGraphWiden(trigger, null);
+    }
+
+    public static boolean shouldFailClosedOnAmbiguousCallGraphWiden(
+            TriggerPoint trigger,
+            String machineEventTypeFqn) {
         if (trigger == null || trigger.isExternal() || !trigger.isAmbiguous()) {
             return false;
         }
@@ -33,7 +39,7 @@ public final class CallChainLinkPolicy {
         }
         if (event != null && event.contains(".valueOf(")) {
             return MachineEnumCanonicalizer.isDynamicTriggerExpression(event)
-                    && !isTrustedEnumPolymorphicWiden(trigger);
+                    && !isTrustedEnumPolymorphicWiden(trigger, machineEventTypeFqn);
         }
         return MachineEnumCanonicalizer.isDynamicTriggerExpression(event)
                 && event != null;
@@ -45,10 +51,17 @@ public final class CallChainLinkPolicy {
      * cross-package or unqualified widens.
      */
     static boolean isTrustedEnumPolymorphicWiden(TriggerPoint trigger) {
+        return isTrustedEnumPolymorphicWiden(trigger, null);
+    }
+
+    static boolean isTrustedEnumPolymorphicWiden(TriggerPoint trigger, String machineEventTypeFqn) {
         if (trigger == null) {
             return false;
         }
         String eventTypeFqn = trigger.getEventTypeFqn();
+        if (eventTypeFqn == null || eventTypeFqn.isBlank()) {
+            eventTypeFqn = machineEventTypeFqn;
+        }
         if (eventTypeFqn == null || eventTypeFqn.isBlank()) {
             return false;
         }

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

@@ -68,7 +68,8 @@ public class TransitionLinkerEnricher implements AnalysisEnricher {
             }
 
             String triggerSource = tp.getSourceState() != null ? simplifySourceState(tp.getSourceState()) : null;
-            if (CallChainLinkPolicy.shouldFailClosedOnAmbiguousCallGraphWiden(tp)) {
+            if (CallChainLinkPolicy.shouldFailClosedOnAmbiguousCallGraphWiden(
+                    tp, machineTypes != null ? machineTypes.eventTypeFqn() : null)) {
                 updatedChains.add(chain.toBuilder()
                         .triggerPoint(tp)
                         .matchedTransitions(null)

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

@@ -45,11 +45,14 @@ public class TriggerCanonicalizationEnricher implements AnalysisEnricher {
                     if (chain.getTriggerPoint() == null) {
                         return chain;
                     }
-                    TriggerPoint canonical = MachineEnumCanonicalizer.canonicalizeTriggerPoint(
+                    TriggerPoint expanded = MachineEnumCanonicalizer.ensureCallChainPolymorphicEvents(
-                            chain.getTriggerPoint(), machineTypes);
+                            chain.getTriggerPoint(), machineTypes, context, result.getTransitions(), true);
-                    TriggerPoint enriched = MachineEnumCanonicalizer.ensureCallChainPolymorphicEvents(
+                    TriggerPoint canonical = expanded.getEventTypeFqn() == null && expanded.getStateTypeFqn() == null
-                            canonical, machineTypes, context, result.getTransitions(), true);
+                            ? MachineEnumCanonicalizer.canonicalizeTriggerLabelsForLinking(
-                    return chain.toBuilder().triggerPoint(enriched).build();
+                                    expanded, machineTypes, context)
+                            : MachineEnumCanonicalizer.canonicalizeTriggerPoint(
+                                    expanded, machineTypes, context);
+                    return chain.toBuilder().triggerPoint(canonical).build();
                 })
                 .collect(Collectors.toList());
 

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

@@ -2756,16 +2756,52 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
     }
 
     private String resolveValueOfEnumTypeName(MethodInvocation mi) {
-        if (mi.getExpression() instanceof SimpleName sn) {
+        if (mi == null || !"valueOf".equals(mi.getName().getIdentifier())) {
+            return null;
+        }
+        org.eclipse.jdt.core.dom.IMethodBinding binding = mi.resolveMethodBinding();
+        if (binding != null && binding.getDeclaringClass() != null) {
+            String fqn = binding.getDeclaringClass().getErasure().getQualifiedName();
+            if (isResolvedEnumTypeFqn(fqn)) {
+                return fqn;
+            }
+        }
+        Expression receiver = mi.getExpression();
+        if (receiver instanceof SimpleName sn) {
+            org.eclipse.jdt.core.dom.ITypeBinding typeBinding = sn.resolveTypeBinding();
+            if (typeBinding != null) {
+                String fqn = typeBinding.getErasure().getQualifiedName();
+                if (isResolvedEnumTypeFqn(fqn)) {
+                    return fqn;
+                }
+            }
             return sn.getIdentifier();
         }
-        if (mi.getExpression() instanceof QualifiedName qn) {
+        if (receiver instanceof QualifiedName qn) {
+            org.eclipse.jdt.core.dom.IBinding nameBinding = qn.resolveBinding();
+            if (nameBinding instanceof org.eclipse.jdt.core.dom.ITypeBinding typeBinding) {
+                String fqn = typeBinding.getErasure().getQualifiedName();
+                if (isResolvedEnumTypeFqn(fqn)) {
+                    return fqn;
+                }
+            }
             String full = qn.getFullyQualifiedName();
-            return full.contains(".") ? full.substring(full.lastIndexOf('.') + 1) : full;
+            if (full.contains(".")) {
+                String typePart = full.substring(0, full.lastIndexOf('.'));
+                if (isResolvedEnumTypeFqn(typePart)) {
+                    return typePart;
+                }
+                return full.substring(full.lastIndexOf('.') + 1);
+            }
+            return full;
         }
         return null;
     }
 
+    private static boolean isResolvedEnumTypeFqn(String fqn) {
+        return fqn != null && !fqn.isBlank() && fqn.contains(".") && !fqn.startsWith("<");
+    }
+
     private String resolveEnumConstantFromArgument(String argName, List<String> path, Map<String, List<CallEdge>> callGraph) {
         if (argName == null || path == null || path.size() < 2) {
             return null;

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

@@ -60,6 +60,20 @@ class CallChainLinkPolicyTest {
         assertThat(CallChainLinkPolicy.shouldFailClosedOnAmbiguousCallGraphWiden(trigger)).isFalse();
     }
 
+    @Test
+    void shouldAllowTrustedWidenUsingMachineEventTypeWhenTriggerTypeIsNull() {
+        TriggerPoint trigger = TriggerPoint.builder()
+                .ambiguous(true)
+                .event("OrderEvent.valueOf(eventStr)")
+                .polymorphicEvents(List.of("com.example.OrderEvent.PAY", "com.example.OrderEvent.SHIP"))
+                .build();
+
+        assertThat(CallChainLinkPolicy.isTrustedEnumPolymorphicWiden(
+                trigger, "com.example.OrderEvent")).isTrue();
+        assertThat(CallChainLinkPolicy.shouldFailClosedOnAmbiguousCallGraphWiden(
+                trigger, "com.example.OrderEvent")).isFalse();
+    }
+
     @Test
     void shouldNotTrustImportStylePolymorphicWidenWithoutPackage() {
         TriggerPoint trigger = TriggerPoint.builder()

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

@@ -1,6 +1,7 @@
 package click.kamil.springstatemachineexporter.analysis.service;
 
 import click.kamil.springstatemachineexporter.analysis.enricher.TransitionLinkerEnricher;
+import click.kamil.springstatemachineexporter.analysis.enricher.TriggerCanonicalizationEnricher;
 import click.kamil.springstatemachineexporter.analysis.model.AnalysisResult;
 import click.kamil.springstatemachineexporter.analysis.model.CallChain;
 import click.kamil.springstatemachineexporter.analysis.model.CodebaseMetadata;
@@ -224,5 +225,39 @@ class AmbiguousSimpleEnumLinkingTest {
         assertThat(linked.getMatchedTransitions()).isNullOrEmpty();
         assertThat(linked.getLinkResolution()).isIn(LinkResolution.NO_MATCH, LinkResolution.AMBIGUOUS_WIDEN);
     }
+
+    @Test
+    void shouldNotRewriteAmbiguousImportStylePolyDuringTriggerCanonicalization(@TempDir Path tempDir) throws Exception {
+        Files.createDirectories(tempDir.resolve("a"));
+        Files.createDirectories(tempDir.resolve("b"));
+        Files.writeString(tempDir.resolve("a/OrderEvent.java"),
+                "package a; public enum OrderEvent { PAY }");
+        Files.writeString(tempDir.resolve("b/OrderEvent.java"),
+                "package b; public enum OrderEvent { CANCEL }");
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        TriggerPoint trigger = TriggerPoint.builder()
+                .event("OrderEvent.valueOf(eventStr)")
+                .polymorphicEvents(List.of("OrderEvent.PAY", "OrderEvent.SHIP"))
+                .ambiguous(true)
+                .build();
+
+        AnalysisResult result = AnalysisResult.builder()
+                .name("com.example.OrderStateMachineConfig")
+                .eventTypeFqn("a.OrderEvent")
+                .transitions(List.of())
+                .metadata(CodebaseMetadata.builder()
+                        .callChains(List.of(CallChain.builder().triggerPoint(trigger).build()))
+                        .build())
+                .build();
+
+        new TriggerCanonicalizationEnricher().enrich(result, context, null);
+
+        TriggerPoint canonical = result.getMetadata().getCallChains().get(0).getTriggerPoint();
+        assertThat(canonical.getPolymorphicEvents())
+                .containsExactly("OrderEvent.PAY", "OrderEvent.SHIP");
+    }
 }