better finding state machines

settings.gradle

@@ -10,4 +10,5 @@ include ':state_machines:forkjoin_state_machine'
 include ':state_machines:inheritance_state_machine'
 include ':state_machines:inheritance_extra_functions_state_machine'
 include ':state_machines:inheritance_extra_functions2_state_machine'
+include ':state_machines:inheritance_extra_functions3_state_machine'
 include ':state_machines:simple_state_machine'

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/app/StateMachineAggregator.java

@@ -29,10 +29,12 @@ public class StateMachineAggregator {
             allTransitions.addAll(parseTransitionsWithMethodCalls(configureMethod, searchStartClass, visitedMethods, argsMap));
         }
 
+        CompilationUnit cu = (CompilationUnit) currentTd.getRoot();
+
         // Hierarchy - Superclass
         Type superclassType = currentTd.getSuperclassType();
         if (superclassType != null) {
-            TypeDeclaration superclassTd = context.getTypeDeclaration(getSimpleName(superclassType));
+            TypeDeclaration superclassTd = context.getTypeDeclaration(getSimpleName(superclassType), cu);
             if (superclassTd != null) {
                 aggregateTransitionsRecursive(superclassTd, searchStartClass, allTransitions, visitedClasses, visitedMethods, argsMap);
             }
@@ -41,7 +43,7 @@ public class StateMachineAggregator {
         // Hierarchy - Interfaces
         for (Object interfaceTypeObj : currentTd.superInterfaceTypes()) {
             if (interfaceTypeObj instanceof Type interfaceType) {
-                TypeDeclaration interfaceTd = context.getTypeDeclaration(getSimpleName(interfaceType));
+                TypeDeclaration interfaceTd = context.getTypeDeclaration(getSimpleName(interfaceType), cu);
                 if (interfaceTd != null) {
                     aggregateTransitionsRecursive(interfaceTd, searchStartClass, allTransitions, visitedClasses, visitedMethods, argsMap);
                 }
@@ -376,17 +378,19 @@ public class StateMachineAggregator {
         MethodDeclaration configureMethod = findConfigureStatesMethod(currentTd);
         if (configureMethod != null) parseStatesWithMethodCalls(configureMethod, searchStartClass, initialStates, endStates, visitedMethods, argsMap);
 
+        CompilationUnit cu = (CompilationUnit) currentTd.getRoot();
+
         // Hierarchy - Superclass
         Type superclassType = currentTd.getSuperclassType();
         if (superclassType != null) {
-            TypeDeclaration superclassTd = context.getTypeDeclaration(getSimpleName(superclassType));
+            TypeDeclaration superclassTd = context.getTypeDeclaration(getSimpleName(superclassType), cu);
             if (superclassTd != null) aggregateStatesRecursive(superclassTd, searchStartClass, initialStates, endStates, visitedClasses, visitedMethods, argsMap);
         }
 
         // Hierarchy - Interfaces
         for (Object interfaceTypeObj : currentTd.superInterfaceTypes()) {
             if (interfaceTypeObj instanceof Type interfaceType) {
-                TypeDeclaration interfaceTd = context.getTypeDeclaration(getSimpleName(interfaceType));
+                TypeDeclaration interfaceTd = context.getTypeDeclaration(getSimpleName(interfaceType), cu);
                 if (interfaceTd != null) aggregateStatesRecursive(interfaceTd, searchStartClass, initialStates, endStates, visitedClasses, visitedMethods, argsMap);
             }
         }

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

@@ -7,7 +7,11 @@ public final class AstUtils {
     }
     public static String extractSimpleTypeName(Type type) {
         if (type.isSimpleType()) {
-            return ((SimpleType) type).getName().getFullyQualifiedName();
+            Name name = ((SimpleType) type).getName();
+            if (name instanceof QualifiedName qn) {
+                return qn.getName().getIdentifier();
+            }
+            return name.toString();
         } else if (type.isParameterizedType()) {
             Type rawType = ((ParameterizedType) type).getType();
             return extractSimpleTypeName(rawType);

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

@@ -9,20 +9,22 @@ import java.util.*;
 
 public class CodebaseContext {
     private final Map<String, CompilationUnit> classes = new HashMap<>();
-    private final Map<String, String> classSources = new HashMap<>();
     private final Map<String, Path> classPaths = new HashMap<>();
+    private final Map<String, String> simpleNameToFqn = new HashMap<>();
 
     public void scan(Path rootDir) throws IOException {
         List<Path> javaFiles = FileUtils.findFilesWithExtension(Set.of(rootDir), ".java");
         for (Path javaFile : javaFiles) {
             String source = Files.readString(javaFile);
             CompilationUnit cu = parse(source);
+            String packageName = cu.getPackage() != null ? cu.getPackage().getName().getFullyQualifiedName() : "";
             for (Object type : cu.types()) {
                 if (type instanceof TypeDeclaration td) {
-                    String className = td.getName().getFullyQualifiedName();
+                    String simpleName = td.getName().getIdentifier();
-                    classes.put(className, cu);
+                    String fqn = packageName.isEmpty() ? simpleName : packageName + "." + simpleName;
-                    classSources.put(className, source);
+                    classes.put(fqn, cu);
-                    classPaths.put(className, javaFile);
+                    classPaths.put(fqn, javaFile);
+                    simpleNameToFqn.put(simpleName, fqn);
                 }
             }
         }
@@ -36,12 +38,67 @@ public class CodebaseContext {
         return (CompilationUnit) parser.createAST(null);
     }
 
-    public TypeDeclaration getTypeDeclaration(String className) {
+    public TypeDeclaration getTypeDeclaration(String name) {
-        CompilationUnit cu = classes.get(className);
+        // Try exact FQN match first
-        if (cu == null) return null;
+        CompilationUnit cu = classes.get(name);
+        if (cu != null) return findTypeInCu(cu, name);
+        
+        // If not found, it might be a simple name
+        String fqn = simpleNameToFqn.get(name);
+        if (fqn != null) {
+            cu = classes.get(fqn);
+            if (cu != null) return findTypeInCu(cu, fqn);
+        }
+        
+        return null;
+    }
+
+    public TypeDeclaration getTypeDeclaration(String name, CompilationUnit contextCu) {
+        if (name == null || name.isEmpty()) return null;
+        
+        // 1. Check if it's already an FQN
+        if (classes.containsKey(name)) return getTypeDeclaration(name);
+
+        // 2. Check imports in contextCu
+        if (contextCu != null) {
+            for (Object impObj : contextCu.imports()) {
+                ImportDeclaration imp = (ImportDeclaration) impObj;
+                String impName = imp.getName().getFullyQualifiedName();
+                if (!imp.isStatic() && !imp.isOnDemand()) {
+                    if (impName.endsWith("." + name)) return getTypeDeclaration(impName);
+                }
+            }
+            
+            // 3. Check same package
+            String packageName = contextCu.getPackage() != null ? contextCu.getPackage().getName().getFullyQualifiedName() : "";
+            String localFqn = packageName.isEmpty() ? name : packageName + "." + name;
+            if (classes.containsKey(localFqn)) return getTypeDeclaration(localFqn);
+            
+            // 4. Check on-demand imports (star imports)
+            for (Object impObj : contextCu.imports()) {
+                ImportDeclaration imp = (ImportDeclaration) impObj;
+                if (imp.isOnDemand()) {
+                    String starFqn = imp.getName().getFullyQualifiedName() + "." + name;
+                    if (classes.containsKey(starFqn)) return getTypeDeclaration(starFqn);
+                }
+            }
+            
+            // 5. Fallback to java.lang (common)
+            String langFqn = "java.lang." + name;
+            if (classes.containsKey(langFqn)) return getTypeDeclaration(langFqn);
+        }
+
+        // 6. Last resort: global simple name match (might be ambiguous but better than nothing)
+        return getTypeDeclaration(name);
+    }
+
+    private TypeDeclaration findTypeInCu(CompilationUnit cu, String fqn) {
+        String packageName = cu.getPackage() != null ? cu.getPackage().getName().getFullyQualifiedName() : "";
         for (Object type : cu.types()) {
-            if (type instanceof TypeDeclaration td && td.getName().getFullyQualifiedName().equals(className)) {
+            if (type instanceof TypeDeclaration td) {
-                return td;
+                String simpleName = td.getName().getIdentifier();
+                String typeFqn = packageName.isEmpty() ? simpleName : packageName + "." + simpleName;
+                if (typeFqn.equals(fqn)) return td;
             }
         }
         return null;
@@ -88,13 +145,14 @@ public class CodebaseContext {
 
     public boolean extendsStateMachineConfigurerAdapter(TypeDeclaration td) {
         Set<String> knownAdapters = Set.of("EnumStateMachineConfigurerAdapter", "StateMachineConfigurerAdapter", "StateMachineConfigurer");
-        
+        CompilationUnit cu = (CompilationUnit) td.getRoot();
+
         // Check superclass
         Type superclass = td.getSuperclassType();
         if (superclass != null) {
             String superclassName = getSimpleName(superclass);
             if (knownAdapters.contains(superclassName)) return true;
-            TypeDeclaration superTd = getTypeDeclaration(superclassName);
+            TypeDeclaration superTd = getTypeDeclaration(superclassName, cu);
             if (superTd != null && extendsStateMachineConfigurerAdapter(superTd)) return true;
         }
 
@@ -103,7 +161,7 @@ public class CodebaseContext {
             if (interfaceTypeObj instanceof Type interfaceType) {
                 String interfaceName = getSimpleName(interfaceType);
                 if (knownAdapters.contains(interfaceName)) return true;
-                TypeDeclaration interfaceTd = getTypeDeclaration(interfaceName);
+                TypeDeclaration interfaceTd = getTypeDeclaration(interfaceName, cu);
                 if (interfaceTd != null && extendsStateMachineConfigurerAdapter(interfaceTd)) return true;
             }
         }

state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/RegressionTest.java

@@ -61,6 +61,30 @@ public class RegressionTest {
                         Path.of("../state_machines/inheritance_extra_functions2_state_machine"),
                         Path.of("src/test/resources/golden/ThreeStateMachineConfiguration"),
                         "ThreeStateMachineConfiguration"
+                ),
+                new TestScenario(
+                        "F1 State Machine (Example 4)",
+                        Path.of("../state_machines/inheritance_extra_functions3_state_machine"),
+                        Path.of("src/test/resources/golden/F1StateMachineConfiguration"),
+                        "F1StateMachineConfiguration"
+                ),
+                new TestScenario(
+                        "F2 State Machine (Example 4)",
+                        Path.of("../state_machines/inheritance_extra_functions3_state_machine"),
+                        Path.of("src/test/resources/golden/F2StateMachineConfiguration"),
+                        "F2StateMachineConfiguration"
+                ),
+                new TestScenario(
+                        "G1 State Machine (Example 4)",
+                        Path.of("../state_machines/inheritance_extra_functions3_state_machine"),
+                        Path.of("src/test/resources/golden/G1StateMachineConfiguration"),
+                        "G1StateMachineConfiguration"
+                ),
+                new TestScenario(
+                        "G2 State Machine (Example 4)",
+                        Path.of("../state_machines/inheritance_extra_functions3_state_machine"),
+                        Path.of("src/test/resources/golden/G2StateMachineConfiguration"),
+                        "G2StateMachineConfiguration"
                 )
         );
     }

state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/ast/in/AstFileFinderTest.java

@@ -4,6 +4,7 @@ import org.junit.jupiter.api.Test;
 
 import java.util.Collections;
 import java.util.List;
+import java.util.Set;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy;
@@ -273,21 +274,38 @@ class AstFileFinderTest {
     }
 
     @Test
-    void shouldReturnTrue_whenClassHasAnnotationWithParameters() {
+    void shouldReturnTrue_whenFunctionReturnsTargetType() {
-        String classWithAnnotatedParams = """
+        String source = """
-                @Service("testService")
+                public class Test {
-                public class TestService {
+                    public StateMachine<S, E> build() { return null; }
-                    public void doSomething() {
-                        // implementation
-                    }
                 }
                 """;
         AstFileFinder finder = new AstFileFinder();
-        List<String> annotationNames = List.of("Service");
+        boolean result = finder.hasFunctionReturningType(source, Set.of("StateMachine"));
-
-        boolean result = finder.hasClassWithAnnotation(classWithAnnotatedParams, annotationNames);
-
         assertThat(result).isTrue();
     }
 
+    @Test
+    void shouldReturnFalse_whenNoFunctionReturnsTargetType() {
+        String source = """
+                public class Test {
+                    public void build() {}
+                }
+                """;
+        AstFileFinder finder = new AstFileFinder();
+        boolean result = finder.hasFunctionReturningType(source, Set.of("StateMachine"));
+        assertThat(result).isFalse();
+    }
+
+    @Test
+    void shouldReturnTrue_whenFunctionReturnsQualifiedTargetType() {
+        String source = """
+                public class Test {
+                    public org.springframework.statemachine.StateMachine build() { return null; }
+                }
+                """;
+        AstFileFinder finder = new AstFileFinder();
+        boolean result = finder.hasFunctionReturningType(source, Set.of("StateMachine"));
+        assertThat(result).isTrue();
+    }
 } 

state_machines/inheritance_extra_functions3_state_machine/src/main/java/click/kamil/examples/statemachine/inheritanceextrafunctionsstate3/AbstractGStateMachineConfiguration.java

@@ -0,0 +1,139 @@
+package click.kamil.examples.statemachine.inheritanceextrafunctionsstate3;
+
+import org.springframework.statemachine.action.Action;
+import org.springframework.statemachine.config.builders.StateMachineConfigurationConfigurer;
+import org.springframework.statemachine.config.builders.StateMachineStateConfigurer;
+import org.springframework.statemachine.config.builders.StateMachineTransitionConfigurer;
+import org.springframework.statemachine.config.configurers.ExternalTransitionConfigurer;
+import org.springframework.statemachine.config.configurers.StateConfigurer;
+import org.springframework.statemachine.guard.Guard;
+
+abstract class AbstractGStateMachineConfiguration extends AbstractStateMachineConfiguration<States, Events> {
+
+    protected abstract void addAdditionalStates(StateConfigurer<States, Events> states);
+
+    protected abstract void addAdditionalCancellationTransitions(StateMachineTransitionConfigurer<States, Events> transitions);
+
+    protected abstract void addAdditionalTransitions(ExternalTransitionConfigurer<States, Events> transitions) throws Exception;
+
+    protected void addAllCancellationTransitionsForSource(StateMachineTransitionConfigurer<States, Events> transitions, States source) {
+        addCancellationTransitionForCancelEvents(transitions, source, Events.EVENT_CANCEL);
+    }
+
+    private void addCancellationTransitionForCancelEvents(
+            StateMachineTransitionConfigurer<States, Events> transitions,
+            States source,
+            Events... cancelEvents) {
+        addTransitionForMultipleEvents(
+                transitions,
+                source,
+                States.CANCEL,
+                transitionConfigurer -> transitionConfigurer.action((context) -> System.out.println(context.getEvent().toString() + " " + context.getTransition().toString())),
+                cancelEvents
+        );
+    }
+
+    @Override
+    public void configure(StateMachineConfigurationConfigurer<States, Events> config) throws Exception {
+        config
+                .withConfiguration()
+                .autoStartup(true);
+    }
+
+    @Override
+    public void configure(StateMachineStateConfigurer<States, Events> states) throws Exception {
+        StateConfigurer<States, Events> end = states
+                .withStates()
+                .initial(States.STATE1)
+                .state(States.STATE2)
+                .state(States.STATE3)
+                .state(States.STATE4)
+                .state(States.STATE5)
+                .state(States.STATE6)
+                .state(States.STATE7)
+                .state(States.STATE8)
+                .state(States.STATE9)
+                .state(States.STATE10)
+                .state(States.STATEX)
+                .choice(States.STATEY)
+                .end(States.STATEZ);
+        addAdditionalStates(end);
+    }
+
+    @Override
+    public void configure(StateMachineTransitionConfigurer<States, Events> transitions) throws Exception {
+
+        // External transitions
+        ExternalTransitionConfigurer<States, Events> t1 = transitions
+                .withExternal().source(States.STATE1).target(States.STATE2).event(Events.EVENT1)
+                .and()
+                .withExternal().source(States.STATE2).target(States.STATE3).event(Events.EVENT2)
+                .and()
+                .withExternal().source(States.STATE3).target(States.STATE4).event(Events.EVENT3)
+                .and()
+                .withExternal().source(States.STATE4).target(States.STATE5).event(Events.EVENT4)
+                .and()
+                .withExternal().source(States.STATE5).target(States.STATE6).event(Events.EVENT5)
+                .and()
+                .withExternal().source(States.STATE6).target(States.STATE7).event(Events.EVENT6)
+                .and()
+                .withExternal().source(States.STATE7).target(States.STATE8).event(Events.EVENT7)
+                .and()
+                .withExternal().source(States.STATE8).target(States.STATE9).event(Events.EVENT8)
+                .and()
+                .withExternal().source(States.STATE9).target(States.STATE10).event(Events.EVENT9);
+
+        transitions.withExternal().source(States.STATE1).target(States.CANCEL).event(Events.EVENT_CANCEL)
+                .and()
+                .withExternal().source(States.STATE2).target(States.CANCEL).event(Events.EVENT_CANCEL)
+                .and()
+                .withExternal().source(States.STATE3).target(States.CANCEL).event(Events.EVENT_CANCEL)
+                .and()
+                .withExternal().source(States.STATE4).target(States.CANCEL).event(Events.EVENT_CANCEL)
+                .and()
+                .withExternal().source(States.STATE5).target(States.CANCEL).event(Events.EVENT_CANCEL);
+
+        // EXTRA CHOICE
+        transitions.withExternal().source(States.STATE4).target(States.STATEY).event(Events.EVENTY);
+        transitions.withChoice().source(States.STATEY).first(States.STATEX, guardVarEquals("value1")).last(States.STATEZ);
+
+        // Choice transitions
+        transitions
+                .withChoice()
+                .source(States.STATE9)
+                .first(States.STATE8, guardVarEquals("stepBack"))
+                .then(States.STATE7, guardVarEquals("stepBackMore"))
+                .last(States.STATE6)
+                .and()
+                .withChoice()
+                .source(States.STATE8)
+                .first(States.STATE9, guardVarEquals("forward9"))
+                .last(States.STATE10);
+
+        addAdditionalTransitions(t1);
+
+        addAdditionalCancellationTransitions(transitions);
+    }
+
+    protected Guard<States, Events> guardVarEquals(String expected) {
+        return context -> {
+            Object var = context.getExtendedState().getVariables().get("var");
+            return expected.equals(var);
+        };
+    }
+
+    protected Guard<States, Events> guardEventHeaderEquals(String headerName, String expected) {
+        return context -> {
+            Object header = context.getMessageHeader(headerName);
+            return expected.equals(header);
+        };
+    }
+
+    private Action<States, Events> logEntryAction(String msg) {
+        return context -> System.out.println(msg);
+    }
+
+    private Action<States, Events> logExitAction(String msg) {
+        return context -> System.out.println(msg);
+    }
+}

state_machines/inheritance_extra_functions3_state_machine/src/main/java/click/kamil/examples/statemachine/inheritanceextrafunctionsstate3/F2StateMachineConfiguration.java

@@ -0,0 +1,42 @@
+package click.kamil.examples.statemachine.inheritanceextrafunctionsstate3;
+
+import org.springframework.statemachine.config.builders.StateMachineTransitionConfigurer;
+import org.springframework.statemachine.config.configurers.ExternalTransitionConfigurer;
+import org.springframework.statemachine.config.configurers.StateConfigurer;
+
+class F2StateMachineConfiguration extends AbstractFStateMachineConfiguration {
+    @Override
+    protected void addAdditionalStates(StateConfigurer<States, Events> states) {
+        states.state(States.STATE_EXTRA_1_1);
+        states.state(States.STATE_EXTRA_1_2);
+        states.state(States.STATE_EXTRA_1);
+        states.state(States.STATE_EXTRA_1_3);
+    }
+
+    @Override
+    protected void addAdditionalCancellationTransitions(StateMachineTransitionConfigurer<States, Events> transitions) {
+        addAllCancellationTransitionsForSource(transitions, States.STATE_EXTRA_1_1);
+        addAllCancellationTransitionsForSource(transitions, States.STATE_EXTRA_1_2);
+        addAllCancellationTransitionsForSource(transitions, States.STATE_EXTRA_1);
+    }
+
+    @Override
+    protected void addAdditionalTransitions(ExternalTransitionConfigurer<States, Events> transitions) throws Exception {
+        transitions
+                .and()
+                .withExternal().event(Events.EVENT_1_1)
+                .source(States.STATE9)
+                .target(States.STATE_EXTRA_1_1)
+                .and()
+                .withChoice()
+                .source(States.STATE_EXTRA_1_2)
+                .first(States.STATE_EXTRA_1_1, guardEventHeaderEquals("header1", "foo"))
+                .last(States.STATE_EXTRA_1_3)
+                .and()
+                .withExternal()
+                .source(States.STATE_EXTRA_1_3)
+                .event(Events.EVENT_1_2)
+                .target(States.STATE_EXTRA_1);
+    }
+
+}

state_machines/inheritance_extra_functions3_state_machine/src/main/java/click/kamil/examples/statemachine/inheritanceextrafunctionsstate3/G1StateMachineConfiguration.java

@@ -0,0 +1,32 @@
+package click.kamil.examples.statemachine.inheritanceextrafunctionsstate3;
+
+import org.springframework.statemachine.config.builders.StateMachineTransitionConfigurer;
+import org.springframework.statemachine.config.configurers.ExternalTransitionConfigurer;
+import org.springframework.statemachine.config.configurers.StateConfigurer;
+
+class G1StateMachineConfiguration extends AbstractGStateMachineConfiguration {
+    @Override
+    protected void addAdditionalStates(StateConfigurer<States, Events> states) {
+        states.state(States.STATE_EXTRA_2);
+        states.state(States.STATE_EXTRA_3);
+    }
+
+    @Override
+    protected void addAdditionalCancellationTransitions(StateMachineTransitionConfigurer<States, Events> transitions) {
+        addAllCancellationTransitionsForSource(transitions, States.STATE_EXTRA_3);
+    }
+
+    @Override
+    protected void addAdditionalTransitions(ExternalTransitionConfigurer<States, Events> transitions) throws Exception {
+        transitions
+                .and()
+                .withExternal().event(Events.EVENT_1_2)
+                .source(States.STATE3)
+                .target(States.STATE_EXTRA_2)
+                .and()
+                .withExternal().event(Events.EVENT_1_2)
+                .source(States.STATE_EXTRA_2)
+                .target(States.STATE_EXTRA_3);
+    }
+
+}

state_machines/inheritance_extra_functions3_state_machine/src/main/java/click/kamil/examples/statemachine/inheritanceextrafunctionsstate3/G2StateMachineConfiguration.java

@@ -0,0 +1,32 @@
+package click.kamil.examples.statemachine.inheritanceextrafunctionsstate3;
+
+import org.springframework.statemachine.config.builders.StateMachineTransitionConfigurer;
+import org.springframework.statemachine.config.configurers.ExternalTransitionConfigurer;
+import org.springframework.statemachine.config.configurers.StateConfigurer;
+
+class G2StateMachineConfiguration extends AbstractGStateMachineConfiguration {
+    @Override
+    protected void addAdditionalStates(StateConfigurer<States, Events> states) {
+        states.state(States.STATE_EXTRA_1);
+        states.state(States.STATE_EXTRA_2);
+    }
+
+    @Override
+    protected void addAdditionalCancellationTransitions(StateMachineTransitionConfigurer<States, Events> transitions) {
+        addAllCancellationTransitionsForSource(transitions, States.STATE_EXTRA_1);
+        addAllCancellationTransitionsForSource(transitions, States.STATE_EXTRA_2);
+    }
+
+    @Override
+    protected void addAdditionalTransitions(ExternalTransitionConfigurer<States, Events> transitions) throws Exception {
+        transitions
+                .and()
+                .withExternal().event(Events.EVENT_1_2)
+                .source(States.STATE3)
+                .target(States.STATE_EXTRA_1)
+                .and()
+                .withExternal().event(Events.EVENT_1_3)
+                .source(States.STATE_EXTRA_1)
+                .target(States.STATE_EXTRA_2);
+    }
+}