Fail closed on ambiguous import-style enum canonicalization

Stop rewriting OrderEvent.PAY to the machine package when the simple name is ambiguous, and apply the same guard when capping polymorphic events to configured transitions.

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

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

@@ -69,6 +69,13 @@ public final class MachineEnumCanonicalizer {
     public static TriggerPoint canonicalizeTriggerPoint(
             TriggerPoint trigger,
             StateMachineTypeResolver.MachineTypes machineTypes) {
+        return canonicalizeTriggerPoint(trigger, machineTypes, null);
+    }
+
+    public static TriggerPoint canonicalizeTriggerPoint(
+            TriggerPoint trigger,
+            StateMachineTypeResolver.MachineTypes machineTypes,
+            CodebaseContext context) {
         if (trigger == null) {
             return null;
         }
@@ -81,14 +88,14 @@ public final class MachineEnumCanonicalizer {
 
         List<String> polymorphicEvents = trigger.getPolymorphicEvents() == null ? null
                 : trigger.getPolymorphicEvents().stream()
-                        .map(event -> canonicalizeLabel(event, eventTypeFqn))
+                        .map(event -> canonicalizeLabel(event, eventTypeFqn, context))
                         .collect(Collectors.toList());
 
         return trigger.toBuilder()
                 .eventTypeFqn(eventTypeFqn)
                 .stateTypeFqn(stateTypeFqn)
-                .event(canonicalizeLabel(trigger.getEvent(), eventTypeFqn))
+                .event(canonicalizeLabel(trigger.getEvent(), eventTypeFqn, context))
-                .sourceState(canonicalizeLabel(trigger.getSourceState(), stateTypeFqn))
+                .sourceState(canonicalizeLabel(trigger.getSourceState(), stateTypeFqn, context))
                 .polymorphicEvents(polymorphicEvents)
                 .build();
     }
@@ -162,7 +169,7 @@ public final class MachineEnumCanonicalizer {
             TriggerPoint trigger,
             StateMachineTypeResolver.MachineTypes machineTypes,
             CodebaseContext context) {
-        TriggerPoint canonical = canonicalizeTriggerPoint(trigger, machineTypes);
+        TriggerPoint canonical = canonicalizeTriggerPoint(trigger, machineTypes, context);
         if (canonical == null || machineTypes == null || machineTypes.eventTypeFqn() == null) {
             return canonical;
         }
@@ -404,7 +411,7 @@ public final class MachineEnumCanonicalizer {
         }
 
         if (!transitionEvents.isEmpty()) {
-            result = capToConfiguredTransitionEvents(result, transitionEvents, eventTypeFqn);
+            result = capToConfiguredTransitionEvents(result, transitionEvents, eventTypeFqn, context);
         }
         return result;
     }
@@ -416,6 +423,14 @@ public final class MachineEnumCanonicalizer {
             List<String> candidates,
             List<String> transitionEvents,
             String eventTypeFqn) {
+        return capToConfiguredTransitionEvents(candidates, transitionEvents, eventTypeFqn, null);
+    }
+
+    static List<String> capToConfiguredTransitionEvents(
+            List<String> candidates,
+            List<String> transitionEvents,
+            String eventTypeFqn,
+            CodebaseContext context) {
         if (transitionEvents == null || transitionEvents.isEmpty()) {
             return candidates == null ? List.of() : candidates;
         }
@@ -440,12 +455,12 @@ public final class MachineEnumCanonicalizer {
                 String transitionType = enumTypeFromRef(transitionEvent);
                 if (candidate.contains(".") && transitionEvent.contains(".")) {
                     if (candidateType != null && transitionType != null
-                            && enumTypesMatch(candidateType, transitionType)) {
+                            && enumTypesMatch(candidateType, transitionType, context)) {
                         capped.add(transitionEvent);
                         break;
                     }
                 } else if (eventTypeFqn != null && transitionType != null
-                        && enumTypesMatch(eventTypeFqn, transitionType)) {
+                        && enumTypesMatch(eventTypeFqn, transitionType, context)) {
                     capped.add(transitionEvent);
                     break;
                 }
@@ -723,6 +738,10 @@ public final class MachineEnumCanonicalizer {
     }
 
     public static String canonicalizeLabel(String value, String enumTypeFqn) {
+        return canonicalizeLabel(value, enumTypeFqn, null);
+    }
+
+    public static String canonicalizeLabel(String value, String enumTypeFqn, CodebaseContext context) {
         if (value == null || value.isBlank() || enumTypeFqn == null || enumTypeFqn.isBlank()) {
             return value;
         }
@@ -749,11 +768,15 @@ public final class MachineEnumCanonicalizer {
         String constant = constantName(stripped);
         String typePart = enumTypeFromRef(stripped);
 
-        if (typePart != null && enumTypesMatch(enumTypeFqn, typePart)) {
+        if (typePart != null && context != null && context.isAmbiguousSimpleName(typePart)) {
+            return stripped;
+        }
+
+        if (typePart != null && enumTypesMatch(enumTypeFqn, typePart, context)) {
             return enumTypeFqn + "." + constant;
         }
 
-        if (typePart != null && !typePart.contains(".") && !enumTypesMatch(enumTypeFqn, typePart)
+        if (typePart != null && !typePart.contains(".") && !enumTypesMatch(enumTypeFqn, typePart, context)
                 && importStyleEnumTypeMatches(typePart, enumTypeFqn)
                 && constant.matches("[A-Z_][A-Z0-9_]*")) {
             return enumTypeFqn + "." + constant;

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

@@ -1,6 +1,12 @@
 package click.kamil.springstatemachineexporter.analysis.resolver;
 
+import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -23,4 +29,39 @@ class MachineEnumCanonicalizerCrossPackageTest {
                 "OrderEvents.PAY", "com.bar.OrderEvents"))
                 .isEqualTo("com.bar.OrderEvents.PAY");
     }
+
+    @Test
+    void shouldNotRewriteImportStyleWhenSimpleNameIsAmbiguous(@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);
+
+        assertThat(MachineEnumCanonicalizer.canonicalizeLabel("OrderEvent.PAY", "a.OrderEvent", context))
+                .isEqualTo("OrderEvent.PAY");
+    }
+
+    @Test
+    void shouldNotCapImportStyleWhenSimpleNameIsAmbiguous(@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);
+
+        assertThat(MachineEnumCanonicalizer.capToConfiguredTransitionEvents(
+                List.of("OrderEvent.PAY"),
+                List.of("a.OrderEvent.PAY"),
+                "a.OrderEvent",
+                context)).isEmpty();
+    }
 }