Thread context through transition and state canonicalization

Machine-scoped transition/state labels and polymorphic event harvesting now use context-aware canonicalizeLabel so ambiguous import-style enum names are not rewritten during export or validation.

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

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/resolver/MachineEnumCanonicalizer.java

@@ -29,38 +29,53 @@ public final class MachineEnumCanonicalizer {
     public static void canonicalizeTransitions(
             List<Transition> transitions,
             StateMachineTypeResolver.MachineTypes machineTypes) {
+        canonicalizeTransitions(transitions, machineTypes, null);
+    }
+
+    public static void canonicalizeTransitions(
+            List<Transition> transitions,
+            StateMachineTypeResolver.MachineTypes machineTypes,
+            CodebaseContext context) {
         if (transitions == null || machineTypes == null) {
             return;
         }
         for (Transition transition : transitions) {
             if (transition.getEvent() != null) {
-                transition.setEvent(canonicalizeEvent(transition.getEvent(), machineTypes.eventTypeFqn()));
+                transition.setEvent(canonicalizeEvent(transition.getEvent(), machineTypes.eventTypeFqn(), context));
             }
             if (transition.getSourceStates() != null) {
-                transition.setSourceStates(canonicalizeStates(transition.getSourceStates(), machineTypes.stateTypeFqn()));
+                transition.setSourceStates(canonicalizeStates(transition.getSourceStates(), machineTypes.stateTypeFqn(), context));
             }
             if (transition.getTargetStates() != null) {
-                transition.setTargetStates(canonicalizeStates(transition.getTargetStates(), machineTypes.stateTypeFqn()));
+                transition.setTargetStates(canonicalizeStates(transition.getTargetStates(), machineTypes.stateTypeFqn(), context));
             }
         }
     }
 
     public static Set<String> canonicalizeStateLabels(Set<String> labels, String stateTypeFqn) {
+        return canonicalizeStateLabels(labels, stateTypeFqn, null);
+    }
+
+    public static Set<String> canonicalizeStateLabels(Set<String> labels, String stateTypeFqn, CodebaseContext context) {
         if (labels == null || labels.isEmpty()) {
             return labels;
         }
         return labels.stream()
-                .map(label -> canonicalizeLabel(label, stateTypeFqn))
+                .map(label -> canonicalizeLabel(label, stateTypeFqn, context))
                 .collect(Collectors.toCollection(LinkedHashSet::new));
     }
 
     public static Set<State> canonicalizeStates(Set<State> states, String stateTypeFqn) {
+        return canonicalizeStates(states, stateTypeFqn, null);
+    }
+
+    public static Set<State> canonicalizeStates(Set<State> states, String stateTypeFqn, CodebaseContext context) {
         if (states == null || states.isEmpty()) {
             return states;
         }
         LinkedHashMap<String, State> byFullIdentifier = new LinkedHashMap<>();
         for (State state : states) {
-            State canonical = canonicalizeState(state, stateTypeFqn);
+            State canonical = canonicalizeState(state, stateTypeFqn, context);
             byFullIdentifier.putIfAbsent(canonical.fullIdentifier(), canonical);
         }
         return new LinkedHashSet<>(byFullIdentifier.values());
@@ -154,7 +169,7 @@ public final class MachineEnumCanonicalizer {
             return polymorphicEvents;
         }
 
-        List<String> transitionEvents = polymorphicEventsFromTransitions(machineTransitions, machineEventTypeFqn);
+        List<String> transitionEvents = polymorphicEventsFromTransitions(machineTransitions, machineEventTypeFqn, context);
 
         List<String> expanded = new ArrayList<>();
         for (String pe : polymorphicEvents) {
@@ -399,7 +414,7 @@ public final class MachineEnumCanonicalizer {
             String eventTypeFqn,
             List<Transition> machineTransitions,
             CodebaseContext context) {
-        List<String> transitionEvents = polymorphicEventsFromTransitions(machineTransitions, eventTypeFqn);
+        List<String> transitionEvents = polymorphicEventsFromTransitions(machineTransitions, eventTypeFqn, context);
         if (transitionEvents.isEmpty()) {
             return List.of();
         }
@@ -420,7 +435,7 @@ public final class MachineEnumCanonicalizer {
             String eventTypeFqn,
             List<Transition> machineTransitions,
             CodebaseContext context) {
-        List<String> transitionEvents = polymorphicEventsFromTransitions(machineTransitions, eventTypeFqn);
+        List<String> transitionEvents = polymorphicEventsFromTransitions(machineTransitions, eventTypeFqn, context);
         List<String> result = current == null ? List.of() : new ArrayList<>(current);
 
         if (EnumMemberPredicateEvaluator.hasEnumMemberPredicates(constraint)) {
@@ -503,6 +518,13 @@ public final class MachineEnumCanonicalizer {
     public static List<String> polymorphicEventsFromTransitions(
             List<Transition> machineTransitions,
             String eventTypeFqn) {
+        return polymorphicEventsFromTransitions(machineTransitions, eventTypeFqn, null);
+    }
+
+    public static List<String> polymorphicEventsFromTransitions(
+            List<Transition> machineTransitions,
+            String eventTypeFqn,
+            CodebaseContext context) {
         if (machineTransitions == null || machineTransitions.isEmpty() || eventTypeFqn == null) {
             return List.of();
         }
@@ -517,7 +539,7 @@ public final class MachineEnumCanonicalizer {
             if (identifier == null || identifier.isBlank()) {
                 continue;
             }
-            String canonical = canonicalizeLabel(identifier, eventTypeFqn);
+            String canonical = canonicalizeLabel(identifier, eventTypeFqn, context);
             if (classifyTriggerEvent(canonical) == TriggerEventKind.CANONICAL_ENUM
                     && !events.contains(canonical)) {
                 events.add(canonical);
@@ -685,25 +707,30 @@ public final class MachineEnumCanonicalizer {
         return !event.equals(denormalized);
     }
 
-    private static List<State> canonicalizeStates(List<State> states, String stateTypeFqn) {
+    private static List<State> canonicalizeStates(List<State> states, String stateTypeFqn, CodebaseContext context) {
         if (states == null) {
             return null;
         }
         LinkedHashMap<String, State> byFullIdentifier = new LinkedHashMap<>();
         for (State state : states) {
-            State canonical = canonicalizeState(state, stateTypeFqn);
+            State canonical = canonicalizeState(state, stateTypeFqn, context);
             byFullIdentifier.putIfAbsent(canonical.fullIdentifier(), canonical);
         }
         return new ArrayList<>(byFullIdentifier.values());
     }
 
     static Event canonicalizeEvent(Event event, String enumTypeFqn) {
+        return canonicalizeEvent(event, enumTypeFqn, null);
+    }
+
+    static Event canonicalizeEvent(Event event, String enumTypeFqn, CodebaseContext context) {
         if (event == null || enumTypeFqn == null || enumTypeFqn.isBlank()) {
             return event;
         }
         String canonical = canonicalizeLabel(
                 event.fullIdentifier() != null ? event.fullIdentifier() : event.rawName(),
-                enumTypeFqn);
+                enumTypeFqn,
+                context);
         String fnRaw = toFnForm(canonical, enumTypeFqn);
         if (fnRaw == null) {
             fnRaw = event.rawName();
@@ -715,12 +742,16 @@ public final class MachineEnumCanonicalizer {
     }
 
     static State canonicalizeState(State state, String enumTypeFqn) {
+        return canonicalizeState(state, enumTypeFqn, null);
+    }
+
+    static State canonicalizeState(State state, String enumTypeFqn, CodebaseContext context) {
         if (state == null || enumTypeFqn == null || enumTypeFqn.isBlank()) {
             return state;
         }
         String full = resolvePlaceholder(
                 state.fullIdentifier() != null ? state.fullIdentifier() : state.rawName());
-        String canonical = canonicalizeLabel(full, enumTypeFqn);
+        String canonical = canonicalizeLabel(full, enumTypeFqn, context);
 
         String raw;
         if (isStringOrPrimitiveType(enumTypeFqn) && canonical.startsWith(enumTypeFqn + ".")) {

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

@@ -86,16 +86,16 @@ public final class AnalysisResultFinalizer {
             AnalysisResult result,
             CodebaseContext context,
             StateMachineTypeResolver.MachineTypes machineTypes) {
-        MachineEnumCanonicalizer.canonicalizeTransitions(result.getTransitions(), machineTypes);
+        MachineEnumCanonicalizer.canonicalizeTransitions(result.getTransitions(), machineTypes, context);
 
         if (result.getStates() != null) {
             result.setStates(MachineEnumCanonicalizer.canonicalizeStates(
-                    result.getStates(), machineTypes.stateTypeFqn()));
+                    result.getStates(), machineTypes.stateTypeFqn(), context));
         }
         result.setStartStates(MachineEnumCanonicalizer.canonicalizeStateLabels(
-                result.getStartStates(), machineTypes.stateTypeFqn()));
+                result.getStartStates(), machineTypes.stateTypeFqn(), context));
         result.setEndStates(MachineEnumCanonicalizer.canonicalizeStateLabels(
-                result.getEndStates(), machineTypes.stateTypeFqn()));
+                result.getEndStates(), machineTypes.stateTypeFqn(), context));
 
         if (result.getMetadata() != null) {
             CodebaseMetadata metadata = result.getMetadata();

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/validation/AnalysisCanonicalFormValidator.java

@@ -365,7 +365,7 @@ public final class AnalysisCanonicalFormValidator {
             return;
         }
         List<String> transitionEvents =
-                MachineEnumCanonicalizer.polymorphicEventsFromTransitions(transitions, eventTypeFqn);
+                MachineEnumCanonicalizer.polymorphicEventsFromTransitions(transitions, eventTypeFqn, context);
         if (transitionEvents.isEmpty()) {
             return;
         }

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/service/ExportService.java

@@ -272,13 +272,13 @@ public class ExportService {
         List<Transition> transitions = aggregator.aggregateTransitions(td);
 
         StateMachineTypeResolver.MachineTypes machineTypes = StateMachineTypeResolver.resolveTypes(className, context);
-        MachineEnumCanonicalizer.canonicalizeTransitions(transitions, machineTypes);
+        MachineEnumCanonicalizer.canonicalizeTransitions(transitions, machineTypes, context);
 
         aggregator.aggregateStates(td);
         Set<String> initialStatesAst = MachineEnumCanonicalizer.canonicalizeStateLabels(
-                aggregator.getInitialStates(), machineTypes.stateTypeFqn());
+                aggregator.getInitialStates(), machineTypes.stateTypeFqn(), context);
         Set<String> endStatesAst = MachineEnumCanonicalizer.canonicalizeStateLabels(
-                aggregator.getEndStates(), machineTypes.stateTypeFqn());
+                aggregator.getEndStates(), machineTypes.stateTypeFqn(), context);
 
         log.debug("Start States Ast: {}", initialStatesAst);
         log.debug("End States Ast: {}", endStatesAst);
@@ -287,7 +287,7 @@ public class ExportService {
         Set<String> endStates = TransitionStateUtils.findEndStates(transitions, endStatesAst);
         Set<click.kamil.springstatemachineexporter.model.State> allStates = MachineEnumCanonicalizer.canonicalizeStates(
                 TransitionStateUtils.findAllStates(transitions, initialStatesAst, endStatesAst),
-                machineTypes.stateTypeFqn());
+                machineTypes.stateTypeFqn(), context);
 
         if (allStates.isEmpty() && transitions.isEmpty()) {
             log.info("Skipping empty state machine config: {}", className);
@@ -323,13 +323,13 @@ public class ExportService {
         List<Transition> transitions = AstTransitionParser.parseTransitions(m, context);
 
         StateMachineTypeResolver.MachineTypes machineTypes = StateMachineTypeResolver.resolveTypes(parentFqn, context);
-        MachineEnumCanonicalizer.canonicalizeTransitions(transitions, machineTypes);
+        MachineEnumCanonicalizer.canonicalizeTransitions(transitions, machineTypes, context);
 
         Set<String> startStates = TransitionStateUtils.findStartStates(transitions, null);
         Set<String> endStates = TransitionStateUtils.findEndStates(transitions, null);
         Set<click.kamil.springstatemachineexporter.model.State> allStates = MachineEnumCanonicalizer.canonicalizeStates(
                 TransitionStateUtils.findAllStates(transitions, null, null),
-                machineTypes.stateTypeFqn());
+                machineTypes.stateTypeFqn(), context);
 
         if (allStates.isEmpty() && transitions.isEmpty()) {
             log.info("Skipping empty state machine bean: {}", uniqueName);

state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/resolver/MachineEnumCanonicalizerCrossPackageTest.java

@@ -110,4 +110,35 @@ class MachineEnumCanonicalizerCrossPackageTest {
         assertThat(MachineEnumCanonicalizer.isMachineEnumReference(
                 "OrderEvent.PAY", "a.OrderEvent", context)).isFalse();
     }
+
+    @Test
+    void shouldNotCanonicalizeAmbiguousImportStyleTransitionEvents(@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 }");
+        Files.writeString(tempDir.resolve("a/OrderState.java"),
+                "package a; public enum OrderState { NEW, PAID }");
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        click.kamil.springstatemachineexporter.model.Transition transition =
+                new click.kamil.springstatemachineexporter.model.Transition();
+        transition.setEvent(click.kamil.springstatemachineexporter.model.Event.of(
+                "OrderEvent.PAY", "OrderEvent.PAY"));
+        transition.setSourceStates(List.of(click.kamil.springstatemachineexporter.model.State.of(
+                "OrderState.NEW", "OrderState.NEW")));
+        transition.setTargetStates(List.of(click.kamil.springstatemachineexporter.model.State.of(
+                "OrderState.PAID", "OrderState.PAID")));
+
+        StateMachineTypeResolver.MachineTypes types = new StateMachineTypeResolver.MachineTypes(
+                "a.OrderState", "a.OrderEvent");
+        MachineEnumCanonicalizer.canonicalizeTransitions(List.of(transition), types, context);
+
+        assertThat(transition.getEvent().fullIdentifier()).isEqualTo("OrderEvent.PAY");
+        assertThat(transition.getSourceStates().get(0).fullIdentifier()).isEqualTo("a.OrderState.NEW");
+    }
 }