more bugs

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

@@ -737,6 +737,23 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
             if (receiver instanceof SimpleName sn) {
                 receiverName = sn.getIdentifier();
                 receiverType = variableTracer.getVariableDeclaredType(entryMethod, receiverName);
+            } else if (receiver instanceof MethodInvocation receiverMi) {
+                String propName = getterName.startsWith("get") ? getterName.substring(3) : getterName;
+                Expression currentMi = receiverMi;
+                while (currentMi instanceof MethodInvocation chainMi) {
+                    String mName = chainMi.getName().getIdentifier();
+                    if (mName.equalsIgnoreCase("set" + propName) || mName.equalsIgnoreCase(propName)) {
+                        if (!chainMi.arguments().isEmpty()) {
+                            Expression arg = (Expression) chainMi.arguments().get(0);
+                            List<String> builderConstants = new ArrayList<>();
+                            constantExtractor.extractConstantsFromExpression(arg, builderConstants);
+                            if (!builderConstants.isEmpty()) {
+                                return builderConstants;
+                            }
+                        }
+                    }
+                    currentMi = chainMi.getExpression();
+                }
             }
 
             if (receiverType == null && path.size() > 1) {

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

@@ -367,7 +367,23 @@ public class HeuristicCallGraphEngine extends AbstractCallGraphEngine {
                 cic = findReturnedCIC(receiverType, factoryMi.getName().getIdentifier(), context);
             }
         }
-        if (cic == null) return null;
+        if (cic == null) {
+            // Support builder pattern: e.g. EventWrapper.builder().event(TransitionEnum.STATE_X).build()
+            if (initExpr[0] instanceof MethodInvocation initMi) {
+                String propName = getterName.startsWith("get") ? getterName.substring(3) : getterName;
+                Expression currentMi = initMi;
+                while (currentMi instanceof MethodInvocation chainMi) {
+                    String mName = chainMi.getName().getIdentifier();
+                    if (mName.equalsIgnoreCase("set" + propName) || mName.equalsIgnoreCase(propName)) {
+                        if (!chainMi.arguments().isEmpty()) {
+                            return chainMi.arguments().get(0).toString();
+                        }
+                    }
+                    currentMi = chainMi.getExpression();
+                }
+            }
+            return null;
+        }
 
         // Resolve the variable's declared type to a TypeDeclaration
         CompilationUnit cu = (CompilationUnit) getterCall.getRoot();

state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/BuilderBugTest.java

@@ -0,0 +1,79 @@
+package click.kamil.springstatemachineexporter.analysis.service;
+
+import click.kamil.springstatemachineexporter.analysis.model.CallChain;
+import click.kamil.springstatemachineexporter.analysis.model.EntryPoint;
+import click.kamil.springstatemachineexporter.analysis.model.TriggerPoint;
+import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
+import org.junit.jupiter.api.Test;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class BuilderBugTest {
+
+    @Test
+    void shouldResolveEventFromBuilder() throws IOException {
+        String source = """
+                package com.example;
+                
+                public class Controller {
+                    private StateMachine machine;
+                    
+                    public void process() {
+                        EventWrapper wrapper = EventWrapper.builder()
+                                .event(TransitionEnum.STATE_X)
+                                .build();
+                        machine.fire(wrapper.getEvent());
+                    }
+                }
+                
+                class StateMachine {
+                    public void fire(TransitionEnum event) {}
+                }
+                
+                class EventWrapper {
+                    private TransitionEnum event;
+                    public TransitionEnum getEvent() { return event; }
+                    public static Builder builder() { return new Builder(); }
+                    
+                    public static class Builder {
+                        private TransitionEnum event;
+                        public Builder event(TransitionEnum e) { this.event = e; return this; }
+                        public EventWrapper build() { 
+                            EventWrapper w = new EventWrapper();
+                            w.event = this.event;
+                            return w;
+                        }
+                    }
+                }
+                
+                enum TransitionEnum { STATE_X, STATE_Y }
+                """;
+
+        Path tempDir = Files.createTempDirectory("callgraph_test_builder");
+        Files.writeString(tempDir.resolve("Config.java"), source);
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        HeuristicCallGraphEngine engine = new HeuristicCallGraphEngine(context);
+
+        EntryPoint entryPoint = EntryPoint.builder()
+                .className("com.example.Controller")
+                .methodName("process")
+                .build();
+
+        TriggerPoint trigger = TriggerPoint.builder()
+                .className("com.example.StateMachine")
+                .methodName("fire")
+                .event("event")
+                .build();
+
+        List<CallChain> chains = engine.findChains(List.of(entryPoint), List.of(trigger));
+        assertThat(chains).hasSize(1);
+        System.out.println("Poly events: " + chains.get(0).getTriggerPoint().getPolymorphicEvents());
+        assertThat(chains.get(0).getTriggerPoint().getPolymorphicEvents()).contains("TransitionEnum.STATE_X");
+    }
+}