better tests

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

@@ -352,7 +352,7 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
                             resolvedValue = cic.toString() + "." + methodName + "()";
                         }
                     }
-                    if (!polymorphicEvents.isEmpty()) {
+                    System.out.println("Returning early with events: " + polymorphicEvents + " for " + resolvedValue); if (!polymorphicEvents.isEmpty()) {
                         return TriggerPoint.builder()
                             .event(resolvedValue)
                             .className(tp.getClassName())
@@ -438,6 +438,21 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine {
                             if (values != null && !values.isEmpty()) {
                                 polymorphicEvents.addAll(values);
                             }
+                            if (polymorphicEvents.isEmpty() && variableTracer != null) {
+                                String initializer = variableTracer.traceLocalVariable(methodFqn, varName);
+                                if (initializer != null && !initializer.equals(varName)) {
+                                    ASTParser mapParser = ASTParser.newParser(AST.JLS17);
+                                    mapParser.setKind(ASTParser.K_EXPRESSION);
+                                    mapParser.setSource(initializer.toCharArray());
+                                    ASTNode mapNode = mapParser.createAST(null);
+                                    if (mapNode instanceof Expression) {
+                                        List<Expression> mapTraced = variableTracer.traceVariableAll((Expression) mapNode);
+                                        for (Expression t : mapTraced) {
+                                            constantExtractor.extractConstantsFromExpression(t, polymorphicEvents);
+                                        }
+                                    }
+                                }
+                            }
                             sourceMethod = methodFqn;
                             break;
                         }

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

@@ -72,11 +72,20 @@ public class ConstantExtractor {
         } else if (expr instanceof MethodInvocation mi) {
             String methodName = mi.getName().getIdentifier();
             
-            if (methodName.equals("get") && mi.arguments().size() == 1 && mi.getExpression() != null) {
+            System.out.println("Processing mi: " + mi); if (methodName.equals("get") || methodName.equals("getOrDefault")) {
+                if (mi.getExpression() != null) {
+                    if (variableTracer != null) {
+                        List<Expression> traced = variableTracer.traceVariableAll(mi.getExpression());
+                        for (Expression t : traced) {
+                            extractConstantsFromExpression(t, constants);
+                        }
+                    } else {
                         extractConstantsFromExpression(mi.getExpression(), constants);
+                    }
+                }
                 return;
             }
-            if ((methodName.equals("of") || methodName.equals("asList")) && mi.getExpression() instanceof SimpleName) {
+            if (methodName.equals("of") || methodName.equals("ofEntries") || methodName.equals("asList") || methodName.equals("entry")) {
                 for (Object argObj : mi.arguments()) {
                     extractConstantsFromExpression((Expression) argObj, constants);
                 }

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

@@ -341,19 +341,7 @@ public class VariableTracer {
                         Expression traced = traceVariable(expr);
                         if (traced instanceof MethodInvocation mi) {
                             Expression innerMost = unwrapMethodInvocation(mi, 0);
-                            if (innerMost instanceof MethodInvocation innerMi) {
-                                if (innerMi.getExpression() instanceof SimpleName sn) {
-                                    stringified.add(sn.getIdentifier() + "." + innerMi.getName().getIdentifier() + "()");
-                                } else {
-                                    stringified.add(innerMi.getName().getIdentifier() + "()");
-                                }
-                            } else if (innerMost instanceof SimpleName sn) {
-                                stringified.add(sn.getIdentifier());
-                            } else {
                             stringified.add(innerMost.toString());
-                            }
-                        } else if (traced instanceof SimpleName sn) {
-                            stringified.add(sn.getIdentifier());
                         } else if (traced instanceof ArrayInitializer) {
                             stringified.add("new Object[]" + traced.toString());
                         } else {

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

@@ -0,0 +1,69 @@
+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 NewBugTest {
+
+    @Test
+    void shouldResolveMapLookup() throws IOException {
+        String source = """
+                package com.example;
+                import java.util.Map;
+                
+                public class Controller {
+                    private StateMachine machine;
+                    private static final Map<InputEnum, TransitionEnum> MAP = Map.of(
+                        InputEnum.IN_A, TransitionEnum.STATE_X,
+                        InputEnum.IN_B, TransitionEnum.STATE_Y
+                    );
+                
+                    public void process(InputEnum input) {
+                        TransitionEnum trans = MAP.get(input);
+                        machine.fire(trans);
+                    }
+                }
+                
+                class StateMachine {
+                    public void fire(TransitionEnum event) {}
+                }
+                
+                enum InputEnum { IN_A, IN_B }
+                enum TransitionEnum { STATE_X, STATE_Y }
+                """;
+
+        Path tempDir = Files.createTempDirectory("callgraph_test_map");
+        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", "TransitionEnum.STATE_Y");
+    }
+}