C

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

@@ -347,6 +347,10 @@ public class CodebaseContext {
             }
         }
 
+        // Pass 2: register inheritance under canonical FQN keys once all types are indexed.
+        // Single-pass writes in indexType are unsafe when a subclass is scanned before its superclass.
+        indexInheritanceEdges();
+
         this.implementationsCache.clear();
         this.accessorIndex = inlineAccessors
                 ? new AccessorIndexBuilder().build(this)
@@ -388,21 +392,6 @@ public class CodebaseContext {
             simpleNameToFqn.put(simpleName, fqn);
         }
 
-        // Inheritance Mapping
-        // Track implementations (for both interfaces and superclasses)
-        if (td.superInterfaceTypes() != null) {
-            for (Object itfObj : td.superInterfaceTypes()) {
-                Type itf = (Type) itfObj;
-                String itfName = extractTypeName(itf);
-                interfaceToImpls.computeIfAbsent(itfName, k -> new ArrayList<>()).add(fqn);
-            }
-        }
-        Type superclass = td.getSuperclassType();
-        if (superclass != null) {
-            String superName = extractTypeName(superclass);
-            interfaceToImpls.computeIfAbsent(superName, k -> new ArrayList<>()).add(fqn);
-        }
-
         // Recursively index nested types
         for (Object decl : td.bodyDeclarations()) {
             if (decl instanceof TypeDeclaration nestedTd) {
@@ -415,6 +404,54 @@ public class CodebaseContext {
         }
     }
 
+    /**
+     * After all types are indexed, register extends/implements under a canonical key:
+     * resolved FQN when the supertype is in the codebase, otherwise the literal AST name.
+     */
+    private void indexInheritanceEdges() {
+        interfaceToImpls.clear();
+        // Snapshot keys: resolveSupertypeIndexKey -> getTypeDeclaration may cache into typeDeclarations.
+        for (String implFqn : new ArrayList<>(typeDeclarations.keySet())) {
+            TypeDeclaration td = typeDeclarations.get(implFqn);
+            if (td == null) {
+                continue;
+            }
+            CompilationUnit cu = td.getRoot() instanceof CompilationUnit root ? root : null;
+
+            if (td.superInterfaceTypes() != null) {
+                for (Object itfObj : td.superInterfaceTypes()) {
+                    if (itfObj instanceof Type itf) {
+                        String key = resolveSupertypeIndexKey(itf, cu);
+                        interfaceToImpls.computeIfAbsent(key, k -> new ArrayList<>()).add(implFqn);
+                    }
+                }
+            }
+            Type superclass = td.getSuperclassType();
+            if (superclass != null) {
+                String key = resolveSupertypeIndexKey(superclass, cu);
+                interfaceToImpls.computeIfAbsent(key, k -> new ArrayList<>()).add(implFqn);
+            }
+        }
+    }
+
+    /**
+     * Canonical index key for a supertype reference: FQN when resolvable in-context, else raw AST name.
+     */
+    private String resolveSupertypeIndexKey(Type type, CompilationUnit cu) {
+        String raw = extractTypeName(type);
+        if (raw == null || raw.isBlank()) {
+            return raw;
+        }
+        TypeDeclaration superTd = getTypeDeclaration(raw, cu);
+        if (superTd != null) {
+            String fqn = getFqn(superTd);
+            if (fqn != null && !fqn.isBlank()) {
+                return fqn;
+            }
+        }
+        return raw;
+    }
+
     private void indexRecord(org.eclipse.jdt.core.dom.RecordDeclaration rd, String parentFqn, Path javaFile) {
         String simpleName = rd.getName().getIdentifier();
         String fqn = parentFqn.isEmpty() ? simpleName : parentFqn + "." + simpleName;

state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/ast/common/CodebaseContextTest.java

@@ -301,6 +301,57 @@ class CodebaseContextTest {
                 .containsExactlyInAnyOrder("abcd.ccc.Av", "abcd.ccc.RealListener");
     }
 
+    @Test
+    void getImplementationsShouldResolveWhenSubclassScannedBeforeSuperclass() throws IOException {
+        Files.createDirectories(tempDir.resolve("abcd/ccc"));
+        // Names sort before AML so Files.walk typically indexes subclasses first.
+        Files.writeString(tempDir.resolve("abcd/ccc/AaaListener.java"), """
+                package abcd.ccc;
+                public class AaaListener extends AML {}
+                """);
+        Files.writeString(tempDir.resolve("abcd/ccc/Av.java"), """
+                package abcd.ccc;
+                public class Av extends abcd.ccc.AML {}
+                """);
+        Files.writeString(tempDir.resolve("abcd/ccc/AML.java"), """
+                package abcd.ccc;
+                public abstract class AML {}
+                """);
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        assertThat(context.getImplementations("abcd.ccc.AML"))
+                .containsExactlyInAnyOrder("abcd.ccc.AaaListener", "abcd.ccc.Av");
+        assertThat(context.getImplementations("AML"))
+                .containsExactlyInAnyOrder("abcd.ccc.AaaListener", "abcd.ccc.Av");
+    }
+
+    @Test
+    void getImplementationsShouldCanonicalizeImportAwareExtendsToSameFqnBucket() throws IOException {
+        Files.createDirectories(tempDir.resolve("abcd/ccc"));
+        Files.createDirectories(tempDir.resolve("other/pkg"));
+        Files.writeString(tempDir.resolve("abcd/ccc/AML.java"), """
+                package abcd.ccc;
+                public abstract class AML {}
+                """);
+        Files.writeString(tempDir.resolve("abcd/ccc/SamePackageListener.java"), """
+                package abcd.ccc;
+                public class SamePackageListener extends AML {}
+                """);
+        Files.writeString(tempDir.resolve("other/pkg/ImportedListener.java"), """
+                package other.pkg;
+                import abcd.ccc.AML;
+                public class ImportedListener extends AML {}
+                """);
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        assertThat(context.getImplementations("abcd.ccc.AML"))
+                .containsExactlyInAnyOrder("abcd.ccc.SamePackageListener", "other.pkg.ImportedListener");
+    }
+
     @Test
     void areSameTypeOrUnambiguousSimpleMatchShouldFailClosedForAmbiguousTypes() throws IOException {
         Files.createDirectories(tempDir.resolve("a"));