new format, code not pretty

TODO.md

@@ -0,0 +1,2 @@
+1.extract from all possible formats (ask chatgpt)
+

src/main/java/click/kamil/springstatemachineexporter/ast/Main.java

@@ -14,6 +14,7 @@ import click.kamil.springstatemachineexporter.ast.out.StateMachineExporter;
 import org.eclipse.jdt.core.dom.CompilationUnit;
 import org.eclipse.jdt.core.dom.MethodDeclaration;
 
+import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.io.PrintWriter;
 import java.nio.file.Files;
@@ -27,25 +28,9 @@ public class Main {
     private static final String START_DIR = ".";
     private static final String EXTENSION = ".java";
     private static final List<String> TARGET_ANNOTATIONS = List.of("EnableStateMachineFactory", "EnableStateMachine");
+    public static final Set<String> STATE_MACHINE_RETURN_TYPES = Set.of("StateMachine");
 
-    public static void main(String[] args) throws IOException {
+    static void processAnnotationFoundSource(String source, Path javaFile, List<StateMachineExporter> outputs) throws IOException {
-        AstFileFinder finder = new AstFileFinder();
-        List<Path> javaFiles = FileUtils.findFilesWithExtension(Set.of(Paths.get(START_DIR)), EXTENSION);
-
-        // List of output handlers
-        List<StateMachineExporter> outputs = List.of(
-                new PlantUml(),
-                new Dot(),
-                new Scxml()
-        );
-
-        for (Path javaFile : javaFiles) {
-            String source = Files.readString(javaFile);
-
-            if (finder.hasClassWithAnnotation(source, TARGET_ANNOTATIONS)) {
-                System.out.println("Found annotation in: " + javaFile.toAbsolutePath());
-            }
-
         MethodDeclaration configureMethod = new StateMachineTransitionConfigurationMethodFinder(source).findConfigureMethod();
         if (configureMethod != null) {
             System.out.println("Found configure method in: " + javaFile.toAbsolutePath());
@@ -86,5 +71,74 @@ public class Main {
             }
         }
     }
+
+    public static void main(String[] args) throws IOException {
+        AstFileFinder finder = new AstFileFinder();
+        List<Path> javaFiles = FileUtils.findFilesWithExtension(Set.of(Paths.get(START_DIR)), EXTENSION);
+
+        // List of output handlers
+        List<StateMachineExporter> outputs = List.of(
+                new PlantUml(),
+                new Dot(),
+                new Scxml()
+        );
+
+        for (Path javaFile : javaFiles) {
+            String source = Files.readString(javaFile);
+            if (finder.hasClassWithAnnotation(source, TARGET_ANNOTATIONS)) {
+                System.out.println("Found annotation in: " + javaFile.toAbsolutePath());
+                processAnnotationFoundSource(source, javaFile, outputs);
+            } else if (finder.hasFunctionReturningType(source, STATE_MACHINE_RETURN_TYPES)) {
+                System.out.println("Found StateMachine return type in: " + javaFile.toAbsolutePath());
+                processReturnTypeFoundSource(source, javaFile, outputs);
+            }
+        }
+    }
+
+    static void processReturnTypeFoundSource(String source, Path javaFile, List<StateMachineExporter> outputs) {
+        Set<MethodDeclaration> methodsWithReturnType = new StateMachineTransitionConfigurationMethodFinder(source).findMethodsWithReturnType(STATE_MACHINE_RETURN_TYPES);
+        if (!methodsWithReturnType.isEmpty()) {
+            methodsWithReturnType.forEach(m -> {
+                System.out.println("Found method returning StateMachine in: " + javaFile.toAbsolutePath());
+                System.out.println("Method body:\n" + m.getBody());
+                List<Transition> transitions = AstTransitionParser.parseTransitions2(m);
+                for (Transition t : transitions) {
+                    System.out.println(t);
+                }
+
+                Set<String> startStatesTransitions = TransitionStateUtils.findStartStates(transitions);
+                Set<String> endStatesTransitions = TransitionStateUtils.findEndStates(transitions);
+
+                CompilationUnit cu = StateMachineStateConfigurationMethodFinder.parse(source);
+                StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor visitor = new StateMachineStateConfigurationMethodFinder.StateConfigurerVisitor();
+                cu.accept(visitor);
+
+                System.out.println("Start States Ast: " + visitor.getInitialStates());
+                System.out.println("End States Ast: " + visitor.getEndStates());
+                System.out.println("Start States: " + startStatesTransitions);
+                System.out.println("End States: " + endStatesTransitions);
+
+                Set<String> startStates = new HashSet<>(startStatesTransitions);
+                startStates.addAll(visitor.getInitialStates());
+
+                Set<String> endStates = new HashSet<>(endStatesTransitions);
+                endStates.addAll(visitor.getEndStates());
+
+                String baseName = javaFile.getFileName().toString().replaceFirst("[.][^.]+$", "");
+
+                // Iterate over outputs and generate files dynamically
+                for (StateMachineExporter output : outputs) {
+                    String content = output.export(transitions, startStates, endStates, true);
+                    String fileName = baseName + output.getFileExtension();
+                    try (PrintWriter out = new PrintWriter(fileName)) {
+                        out.println(content);
+                    } catch (FileNotFoundException e) {
+                        throw new RuntimeException(e);
+                    }
+                }
+
+            });
+
+        }
     }
 }

src/main/java/click/kamil/springstatemachineexporter/ast/app/AstTransitionParser.java

@@ -27,6 +27,52 @@ public class AstTransitionParser {
                 });
         return transitions;
     }
+    public static List<Transition> parseTransitions2(MethodDeclaration method) {
+        List<Transition> transitions = new ArrayList<>();
+
+        Optional.ofNullable(method)
+                .map(MethodDeclaration::getBody)
+                .ifPresent(body -> {
+                    for (Object stmtObj : body.statements()) {
+                        if (stmtObj instanceof ExpressionStatement es
+                                && es.getExpression() instanceof MethodInvocation mi) {
+
+                            // recursively collect all method invocations in this chain
+                            List<MethodInvocation> chain = collectMethodInvocationChain(mi);
+
+                            // Find the index of `configureTransitions` in the chain
+                            for (int i = 0; i < chain.size(); i++) {
+                                MethodInvocation current = chain.get(i);
+                                if (current.getName().getIdentifier().equals("configureTransitions")) {
+                                    List<MethodInvocation> afterConfigure = chain.subList(i + 1, chain.size());
+                                    transitions.addAll(parseTransitionsFromExpression(afterConfigure.getLast()));
+                                    break;
+                                }
+                            }
+                        }
+                    }
+                });
+
+        return transitions;
+    }
+
+    private static List<MethodInvocation> collectMethodInvocationChain(MethodInvocation root) {
+        List<MethodInvocation> chain = new ArrayList<>();
+
+        MethodInvocation current = root;
+        while (current != null) {
+            chain.add(0, current); // add to the front to preserve order
+
+            Expression expr = current.getExpression();
+            if (expr instanceof MethodInvocation parent) {
+                current = parent;
+            } else {
+                break;
+            }
+        }
+
+        return chain;
+    }
 
     private static List<Transition> parseTransitionsFromExpression(MethodInvocation rootCall) {
         List<Transition> transitions = new ArrayList<>();

src/main/java/click/kamil/springstatemachineexporter/ast/app/StateMachineStateConfigurationMethodFinder.java

@@ -24,9 +24,7 @@ public class StateMachineStateConfigurationMethodFinder {
         private final Set<String> initialStates = new HashSet<>();
         private final Set<String> endStates = new HashSet<>();
 
-        @Override
+        private boolean visitConfigureMethod(MethodDeclaration node) {
-        public boolean visit(MethodDeclaration node) {
-            // Visit method named "configure" only
             if (!"configure".equals(node.getName().getIdentifier())) {
                 return super.visit(node);
             }
@@ -48,6 +46,94 @@ public class StateMachineStateConfigurationMethodFinder {
             }
             return false;  // no need to visit further inside configure
         }
+        private boolean visitBuilderMethod(MethodDeclaration node) {
+            Block body = node.getBody();
+            if (body == null) return false;
+
+            // Traverse statements inside buildStateMachine method
+            for (Object stmtObj : body.statements()) {
+                if (!(stmtObj instanceof ExpressionStatement exprStmt)) continue;
+                Expression expr = exprStmt.getExpression();
+                if (expr instanceof MethodInvocation rootCall) {
+                    // Detect when the chain starts with "configureStates"
+                    // Start looping through the method chain
+                    findConfigureStates(rootCall);
+                }
+            }
+            return false;
+        }
+
+
+        private void findConfigureStates(MethodInvocation methodInvocation) {
+            // Start at the outermost method invocation
+            MethodInvocation currentInvocation = methodInvocation;
+
+            while (currentInvocation != null) {
+                // Check if the current method in the chain is 'configureStates'
+                if ("configureStates".equals(currentInvocation.getName().getIdentifier())) {
+                    // Once found, call the parsing method
+                    parseStateMachineMethodChain(methodInvocation);
+                    return;  // Exit once we've handled the method
+                }
+
+                // Move to the next method invocation in the chain (moving forward)
+                Expression nextExpression = currentInvocation.getExpression();
+
+                // Check if the next expression is another method invocation
+                if (nextExpression instanceof MethodInvocation) {
+                    // Continue to the next method invocation
+                    currentInvocation = (MethodInvocation) nextExpression;
+                } else {
+                    // No more method calls in the chain, break out of the loop
+                    currentInvocation = null;
+                }
+            }
+        }
+
+        // Check for method chain starting with "configureStates"
+        private void parseStateMachineMethodChain(MethodInvocation mi) {
+            List<MethodInvocation> chain = new ArrayList<>();
+            Expression current = mi;
+            while (current instanceof MethodInvocation m) {
+                chain.addFirst(m);
+                current = m.getExpression();
+            }
+
+            for (MethodInvocation call : chain) {
+                String methodName = call.getName().getIdentifier();
+                List<?> args = call.arguments();
+                if (args.isEmpty()) {
+                    continue;
+                }
+                Expression arg = (Expression) args.get(0);
+                switch (methodName) {
+                    case "initial" -> {
+                        if (!isInsideRegionOrFork(call)) {
+                            initialStates.add(QuotedExpression.of(arg).toStringWithoutQuotes());
+                        }
+                    }
+                    case "end" -> {
+                        endStates.add(QuotedExpression.of(arg).toStringWithoutQuotes());
+                    }
+                    // You can also add cases for other methods if needed (e.g., stateEntry, stateExit)
+                    default -> {}
+                }
+            }
+        }
+
+        @Override
+        public boolean visit(MethodDeclaration node) {
+            // Visit method named "configure" only
+            if ("configure".equals(node.getName().getIdentifier())) {
+                return visitConfigureMethod(node);
+            }
+            // Handle builder method configuration (e.g., for buildStateMachine)
+            if ("buildStateMachine".equals(node.getName().getIdentifier())) {
+                return visitBuilderMethod(node);
+            }
+            return false;
+        }
+
 
         private boolean isWithStatesChain(MethodInvocation mi) {
             MethodInvocation current = mi;

src/main/java/click/kamil/springstatemachineexporter/ast/app/StateMachineTransitionConfigurationMethodFinder.java

@@ -1,9 +1,12 @@
 package click.kamil.springstatemachineexporter.ast.app;
 
+import click.kamil.springstatemachineexporter.ast.common.AstUtils;
 import lombok.Getter;
 import org.eclipse.jdt.core.dom.*;
 
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 
 public class StateMachineTransitionConfigurationMethodFinder {
     private final CompilationUnit cu;
@@ -24,6 +27,23 @@ public class StateMachineTransitionConfigurationMethodFinder {
         return visitor.getConfigureMethod();
     }
 
+    public Set<MethodDeclaration> findMethodsWithReturnType(Set<String> returnTypes) {
+        Set<MethodDeclaration> matchedMethods = new HashSet<>();
+        cu.accept(new ASTVisitor() {
+            @Override
+            public boolean visit(MethodDeclaration node) {
+                Type returnType = node.getReturnType2();
+                String typeName = (returnType == null) ? "void" : AstUtils.extractSimpleTypeName(returnType);
+                if (returnTypes.contains(typeName)) {
+                    matchedMethods.add(node);
+                }
+                return true;
+            }
+        });
+        return matchedMethods;
+    }
+
+
     @Getter
     private static class ConfigureMethodVisitor extends ASTVisitor {
         private MethodDeclaration configureMethod = null;

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

@@ -0,0 +1,26 @@
+package click.kamil.springstatemachineexporter.ast.common;
+
+import org.eclipse.jdt.core.dom.*;
+
+public final class AstUtils {
+    private AstUtils() {
+    }
+    public static String extractSimpleTypeName(Type type) {
+        if (type.isSimpleType()) {
+            return ((SimpleType) type).getName().getFullyQualifiedName();
+        } else if (type.isParameterizedType()) {
+            Type rawType = ((ParameterizedType) type).getType();
+            return extractSimpleTypeName(rawType);
+        } else if (type.isQualifiedType()) {
+            return ((QualifiedType) type).getName().getIdentifier();
+        } else if (type.isNameQualifiedType()) {
+            return ((NameQualifiedType) type).getName().getIdentifier();
+        } else if (type.isArrayType()) {
+            return extractSimpleTypeName(((ArrayType) type).getElementType());
+        } else if (type.isPrimitiveType()) {
+            return ((PrimitiveType) type).getPrimitiveTypeCode().toString();
+        } else {
+            return type.toString(); // fallback
+        }
+    }
+}

src/main/java/click/kamil/springstatemachineexporter/ast/in/AstFileFinder.java

@@ -1,8 +1,10 @@
 package click.kamil.springstatemachineexporter.ast.in;
 
+import click.kamil.springstatemachineexporter.ast.common.AstUtils;
 import org.eclipse.jdt.core.dom.*;
 
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 public class AstFileFinder {
@@ -35,4 +37,26 @@ public class AstFileFinder {
         });
         return found.get();
     }
+    public boolean hasFunctionReturningType(String source, Set<String> targetReturnTypes) {
+        parser.setSource(source.toCharArray());
+        CompilationUnit cu = (CompilationUnit) parser.createAST(null);
+        AtomicBoolean found = new AtomicBoolean(false);
+
+        cu.accept(new ASTVisitor() {
+            @Override
+            public boolean visit(MethodDeclaration node) {
+                Type returnType = node.getReturnType2();
+                String typeName = (returnType == null) ? "void" : AstUtils.extractSimpleTypeName(returnType);
+                
+                if (targetReturnTypes.contains(typeName)) {
+                    found.set(true);
+                    return false; // stop visiting
+                }
+
+                return true;
+            }
+        });
+
+        return found.get();
+    }
 }

src/main/java/click/kamil/springstatemachineexporter/statemachine/DynamicStateMachineExample.java

@@ -1,10 +1,6 @@
-1. extract state machines from this as well
+package click.kamil.springstatemachineexporter.statemachine;
-2. extract from all possible formats (ask chatgpt)
-
-```java
-
 import org.springframework.statemachine.StateMachine;
-import org.springframework.statemachine.config.builders.StateMachineBuilder;
+import org.springframework.statemachine.config.StateMachineBuilder;
 
 import java.util.EnumSet;
 
@@ -24,6 +20,7 @@ public class DynamicStateMachineExample {
         builder.configureStates()
                 .withStates()
                 .initial(States.DRAFT)
+                .end(States.REVIEW)
                 .states(EnumSet.allOf(States.class));
 
         builder.configureTransitions()
@@ -49,5 +46,3 @@ public class DynamicStateMachineExample {
         System.out.println("After APPROVE: " + sm.getState().getId());
     }
 }
-
-```