diff --git a/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/resolver/ConstantResolver.java b/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/resolver/ConstantResolver.java index 57e8abb..29274ae 100644 --- a/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/resolver/ConstantResolver.java +++ b/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/resolver/ConstantResolver.java @@ -11,6 +11,7 @@ import java.util.Set; import java.util.LinkedHashSet; import java.util.Arrays; import java.util.List; +import java.util.Map; @Slf4j public class ConstantResolver { @@ -19,6 +20,36 @@ public class ConstantResolver { return resolveInternal(expr, context, new HashSet<>()); } + /** + * Resolves {@code map.get(key)} when the map receiver and key are statically known. + */ + public String resolveKeyedMapLookup(MethodInvocation getCall, CodebaseContext context) { + if (getCall == null || getCall.getExpression() == null || getCall.arguments().isEmpty()) { + return null; + } + String methodName = getCall.getName().getIdentifier(); + if (!"get".equals(methodName) && !"getOrDefault".equals(methodName)) { + return null; + } + return LiteralExpressionSupport.resolveMapGet( + getCall, context, new HashSet<>(), (e, v) -> resolveInternal(e, context, v)); + } + + /** + * Resolves a value from an encoded {@code MAP:...} literal or map-typed expression at {@code key}. + */ + public String resolveKeyedMapLookup(Expression mapExpr, String key, CodebaseContext context) { + if (mapExpr == null || key == null) { + return null; + } + String encoded = resolveInternal(mapExpr, context, new HashSet<>()); + if (encoded == null || !encoded.startsWith("MAP:")) { + return null; + } + Map entries = LiteralExpressionSupport.parseMapLiteral(encoded); + return entries != null ? entries.get(key) : null; + } + public static List decodeMapLiteralValues(String encoded) { return LiteralExpressionSupport.mapLiteralValues(encoded); } @@ -818,12 +849,13 @@ public class ConstantResolver { } } - td = typeFromCurrentCallPath(context); - if (td != null) { - String fqn = context.getFqn(td); - String result = resolveFieldInType(td, sn.getIdentifier(), fqn, context, visited); - if (result != null) { - return result; + for (String classFqn : classesFromCurrentCallPath()) { + TypeDeclaration pathTd = context.getTypeDeclaration(classFqn); + if (pathTd != null) { + String result = resolveFieldInType(pathTd, sn.getIdentifier(), classFqn, context, visited); + if (result != null) { + return result; + } } } @@ -901,21 +933,35 @@ public class ConstantResolver { } private TypeDeclaration typeFromCurrentCallPath(CodebaseContext context) { - String entryClass = getCurrentCallPathEntryClass(); - if (entryClass == null) { + List classes = classesFromCurrentCallPath(); + if (classes.isEmpty()) { return null; } - return context.getTypeDeclaration(entryClass); + return context.getTypeDeclaration(classes.get(0)); + } + + private List classesFromCurrentCallPath() { + List path = JdtDataFlowModel.getCurrentPath(); + if (path == null || path.isEmpty()) { + return List.of(); + } + List classes = new ArrayList<>(); + for (int i = path.size() - 1; i >= 0; i--) { + String methodFqn = path.get(i); + int dot = methodFqn.lastIndexOf('.'); + if (dot > 0) { + String className = methodFqn.substring(0, dot); + if (!classes.contains(className)) { + classes.add(className); + } + } + } + return classes; } private String getCurrentCallPathEntryClass() { - List path = JdtDataFlowModel.getCurrentPath(); - if (path == null || path.isEmpty()) { - return null; - } - String entryMethod = path.get(0); - int dot = entryMethod.lastIndexOf('.'); - return dot > 0 ? entryMethod.substring(0, dot) : null; + List classes = classesFromCurrentCallPath(); + return classes.isEmpty() ? null : classes.get(classes.size() - 1); } private String resolveFieldInType(TypeDeclaration td, String fieldName, String typeFqn, CodebaseContext context, Set visited) { diff --git a/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/AbstractCallGraphEngine.java b/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/AbstractCallGraphEngine.java index 79a8e78..b015a82 100644 --- a/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/AbstractCallGraphEngine.java +++ b/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/AbstractCallGraphEngine.java @@ -428,6 +428,8 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine { while (exprNode instanceof ParenthesizedExpression pe) { exprNode = pe.getExpression(); } + boolean[] keyedMapLookupOnInitializer = {false}; + boolean[] pathBoundMapKeyOnInitializer = {false}; if (exprNode instanceof Expression expressionNode) { String remappedGetter = ReactiveExpressionSupport.remapLambdaParameterGetter( @@ -461,8 +463,25 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine { if ("get".equals(methodName) || "getOrDefault".equals(methodName)) { boolean keyedLookup = expressionAccessClassifier.isKeyedLookup(mi, entryMethod); boolean functionalGet = "get".equals(methodName) && mi.arguments().isEmpty() && !keyedLookup; - if (keyedLookup || functionalGet) { - constantExtractor.extractConstantsFromExpression((Expression) exprNode, polymorphicEvents); + if (keyedLookup) { + List keyedResults = resolveKeyedMapLookup(mi, path, callGraph); + if (!keyedResults.isEmpty()) { + polymorphicEvents.addAll(keyedResults); + return buildTriggerPointWithPolymorphicEvents(tp, resolvedValue, polymorphicEvents); + } + boolean pathBoundKey = hasPathBindingForMapKey(mi, path, callGraph); + if (pathBoundKey) { + keyedMapLookupOnInitializer[0] = true; + pathBoundMapKeyOnInitializer[0] = true; + } + if (!pathBoundKey) { + addConstantsFromTracedExpression((Expression) exprNode, polymorphicEvents, path, callGraph, entryMethod); + if (!polymorphicEvents.isEmpty()) { + return buildTriggerPointWithPolymorphicEvents(tp, resolvedValue, polymorphicEvents); + } + } + } else if (functionalGet) { + addConstantsFromTracedExpression((Expression) exprNode, polymorphicEvents, path, callGraph, entryMethod); if (!polymorphicEvents.isEmpty()) { return TriggerPoint.builder() .event(resolvedValue) @@ -615,6 +634,12 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine { branches.add(cond.getThenExpression()); branches.add(cond.getElseExpression()); for (Expression branch : branches) { + if (branch instanceof ConditionalExpression nestedCond) { + // traceVariableAll collapses literal "true"/"false" conditions to a single branch; + // extract recursively so chained if-else assignments keep all enum variants. + constantExtractor.extractConstantsFromExpression(nestedCond, polymorphicEvents); + continue; + } List tracedBranch = variableTracer.traceVariableAll(branch); for (Expression tb : tracedBranch) { if (tb instanceof ClassInstanceCreation cic) { @@ -764,11 +789,9 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine { resolvedValue = extractedFinalTraced[0] + extractedFinalTraced[1]; if (extractedFinalTraced[0].contains("new ") && extractedFinalTraced[0].contains("{")) { ASTNode newNode = parseExpressionString(resolvedValue); - if (newNode instanceof Expression) { - List traced = variableTracer.traceVariableAll((Expression) newNode); - for (Expression ex : traced) { - constantExtractor.extractConstantsFromExpression(ex, polymorphicEvents); - } + if (newNode instanceof Expression newExpr) { + trackKeyedMapLookupFlags(newExpr, path, callGraph, entryMethod, keyedMapLookupOnInitializer, pathBoundMapKeyOnInitializer); + addConstantsFromTracedExpression(newExpr, polymorphicEvents, path, callGraph, entryMethod); } } } @@ -790,11 +813,9 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine { String initializer = variableTracer.traceLocalVariable(methodFqn, varName); if (initializer != null && !initializer.equals(varName)) { ASTNode mapNode = parseExpressionString(initializer); - if (mapNode instanceof Expression) { - List mapTraced = variableTracer.traceVariableAll((Expression) mapNode); - for (Expression t : mapTraced) { - constantExtractor.extractConstantsFromExpression(t, polymorphicEvents); - } + if (mapNode instanceof Expression initExpr) { + trackKeyedMapLookupFlags(initExpr, path, callGraph, methodFqn, keyedMapLookupOnInitializer, pathBoundMapKeyOnInitializer); + addConstantsFromTracedExpression(initExpr, polymorphicEvents, path, callGraph, methodFqn); } } } @@ -836,7 +857,8 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine { if (enumValues != null && !enumValues.isEmpty()) { if (!hasConcreteEnumConstants(polymorphicEvents, declaredType, context) && !tp.isExternal() - && !isRuntimeEnumParameter(exprNode instanceof Expression expression ? expression : null)) { + && !isRuntimeEnumParameter(exprNode instanceof Expression expression ? expression : null) + && !(keyedMapLookupOnInitializer[0] && pathBoundMapKeyOnInitializer[0])) { for (String ev : enumValues) { polymorphicEvents.add(constantExtractor.parseEnumSetElement(ev)); } @@ -885,8 +907,8 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine { } } - if (!hasValidConstant && exprNode instanceof Expression) { - constantExtractor.extractConstantsFromExpression((Expression) exprNode, polymorphicEvents); + if (!hasValidConstant && exprNode instanceof Expression expressionFallback) { + addConstantsFromTracedExpression(expressionFallback, polymorphicEvents, path, callGraph, entryMethod); } List newPolyEvents = new ArrayList<>(); @@ -964,6 +986,8 @@ public abstract class AbstractCallGraphEngine implements CallGraphEngine { polymorphicEvents = polymorphicEvents.stream().distinct().collect(java.util.stream.Collectors.toList()); + polymorphicEvents = narrowPolymorphicEventsWithPathBindings(polymorphicEvents, path, callGraph, tp); + Expression runtimeExpr = exprNode instanceof Expression expression ? expression : null; if (isRuntimeEnumParameter(runtimeExpr) || tp.isExternal()) { polymorphicEvents.removeIf(pe -> pe != null && pe.startsWith(" resolveKeyedMapLookup( + MethodInvocation getCall, List path, Map> callGraph) { + List results = new ArrayList<>(); + String literalResult = constantResolver.resolveKeyedMapLookup(getCall, context); + if (literalResult != null && !literalResult.isBlank()) { + addResolvedConstantToList(literalResult, results); + return results; + } + if (getCall.arguments().isEmpty()) { + return results; + } + Expression keyArg = (Expression) getCall.arguments().get(0); + String resolvedKey = resolveStringArgumentFromPath(keyArg, path, callGraph); + if (resolvedKey != null && getCall.getExpression() != null) { + String mapValue = resolveKeyedValueFromMapExpression(getCall.getExpression(), resolvedKey); + if (mapValue != null) { + addResolvedConstantToList(mapValue, results); + } + } + return results; + } + + private String resolveStringArgumentFromPath( + Expression keyArg, List path, Map> callGraph) { + if (keyArg instanceof StringLiteral sl) { + return sl.getLiteralValue(); + } + if (keyArg instanceof SimpleName sn) { + Map bindings = pathBindingEvaluator.traceBindings(path, callGraph, pathFinder); + String bound = bindings.get(sn.getIdentifier()); + if (bound != null) { + return unwrapStringLiteral(bound); + } + for (int i = path.size() - 1; i >= 0; i--) { + String traced = variableTracer.traceLocalVariable(path.get(i), sn.getIdentifier()); + if (traced != null && !traced.equals(sn.getIdentifier())) { + String unwrapped = unwrapStringLiteral(traced); + if (unwrapped != null && !unwrapped.isBlank()) { + return unwrapped; + } + } + } + } + String resolved = constantResolver.resolve(keyArg, context); + return unwrapStringLiteral(resolved); + } + + private boolean hasPathBindingForMapKey( + MethodInvocation getCall, List path, Map> callGraph) { + if (getCall.arguments().isEmpty()) { + return false; + } + Expression keyArg = (Expression) getCall.arguments().get(0); + if (!(keyArg instanceof SimpleName sn)) { + return false; + } + Map bindings = pathBindingEvaluator.traceBindings(path, callGraph, pathFinder); + if (bindings.containsKey(sn.getIdentifier())) { + return true; + } + for (String methodFqn : path) { + String traced = variableTracer.traceLocalVariable(methodFqn, sn.getIdentifier()); + if (traced != null && !traced.equals(sn.getIdentifier()) && unwrapStringLiteral(traced) != null) { + return true; + } + } + return false; + } + + private void addConstantsFromTracedExpression( + Expression expr, List constants, List path, Map> callGraph, String scopeMethod) { + List traced = variableTracer.traceVariableAll(expr); + for (Expression tracedExpr : traced) { + addConstantsFromSingleExpression(tracedExpr, constants, path, callGraph, scopeMethod); + } + if (constants.isEmpty()) { + addConstantsFromSingleExpression(expr, constants, path, callGraph, scopeMethod); + } + } + + private void addConstantsFromSingleExpression( + Expression expr, List constants, List path, Map> callGraph, String scopeMethod) { + if (expr instanceof MethodInvocation mi && expressionAccessClassifier.isKeyedLookup(mi, scopeMethod)) { + List keyed = resolveKeyedMapLookup(mi, path, callGraph); + if (!keyed.isEmpty()) { + for (String keyedValue : keyed) { + if (!constants.contains(keyedValue)) { + constants.add(keyedValue); + } + } + return; + } + if (!hasPathBindingForMapKey(mi, path, callGraph)) { + constantExtractor.extractConstantsFromExpression(expr, constants); + } + return; + } + constantExtractor.extractConstantsFromExpression(expr, constants); + } + + private List narrowPolymorphicEventsWithPathBindings( + List polymorphicEvents, + List path, + Map> callGraph, + TriggerPoint tp) { + if (polymorphicEvents == null || polymorphicEvents.isEmpty()) { + return polymorphicEvents; + } + Map bindings = pathBindingEvaluator.traceBindings(path, callGraph, pathFinder); + if (bindings.isEmpty()) { + return polymorphicEvents; + } + List narrowed = new ArrayList<>(polymorphicEvents); + boolean hasConcrete = narrowed.stream().anyMatch(this::looksLikeEnumConstant); + if (hasConcrete) { + narrowed.removeIf(pe -> pe != null && pe.startsWith(" pe != null && pe.startsWith("= 2) { + return trimmed.substring(1, trimmed.length() - 1); + } + return trimmed; + } + + private boolean looksLikeEnumConstant(String value) { + if (value == null || value.isBlank() || value.startsWith(" path, + Map> callGraph, + String scopeMethod, + boolean[] keyedMapLookupOnInitializer, + boolean[] pathBoundMapKeyOnInitializer) { + if (expr instanceof MethodInvocation mi && expressionAccessClassifier.isKeyedLookup(mi, scopeMethod)) { + keyedMapLookupOnInitializer[0] = true; + if (hasPathBindingForMapKey(mi, path, callGraph)) { + pathBoundMapKeyOnInitializer[0] = true; + } + } + } + + private void addResolvedConstantToList(String resolved, List results) { + if (resolved == null || resolved.isBlank()) { + return; + } + if (resolved.startsWith("ENUM_SET:")) { + for (String part : resolved.substring(9).split(",")) { + String parsed = constantExtractor.parseEnumSetElement(part); + if (!results.contains(parsed)) { + results.add(parsed); + } + } + return; + } + String parsed = constantExtractor.parseEnumSetElement(resolved); + if (!results.contains(parsed)) { + results.add(parsed); + } + } } \ No newline at end of file diff --git a/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/ExpressionAccessClassifier.java b/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/ExpressionAccessClassifier.java index 35e982b..e05eb92 100644 --- a/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/ExpressionAccessClassifier.java +++ b/state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/analysis/service/ExpressionAccessClassifier.java @@ -136,7 +136,11 @@ public final class ExpressionAccessClassifier { return "Map"; } if (receiver instanceof SimpleName sn) { - return lookupDeclaredType(scopeMethod, sn.getIdentifier()); + String declared = lookupDeclaredType(scopeMethod, sn.getIdentifier()); + if (declared != null) { + return declared; + } + return lookupFieldType(classNameFromScopeMethod(scopeMethod), sn.getIdentifier()); } if (receiver instanceof ThisExpression) { return classNameFromScopeMethod(scopeMethod); diff --git a/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/PerformanceCacheMultiMachineIsolationTest.java b/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/PerformanceCacheMultiMachineIsolationTest.java new file mode 100644 index 0000000..f548c58 --- /dev/null +++ b/state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/analysis/service/PerformanceCacheMultiMachineIsolationTest.java @@ -0,0 +1,128 @@ +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 org.junit.jupiter.api.io.TempDir; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Guards map-routed endpoint resolution and per-machine polymorphic event isolation when the + * call graph and accessor caches are reused across multiple {@link HeuristicCallGraphEngine#findChains} + * invocations on the same {@link CodebaseContext}. + */ +class PerformanceCacheMultiMachineIsolationTest { + + private static final String ORDER_CONTROLLER = "com.example.OrderController"; + private static final String SHIPMENT_CONTROLLER = "com.example.ShipmentController"; + private static final String ORDER_MACHINE = "com.example.OrderMachine"; + private static final String SHIPMENT_MACHINE = "com.example.ShipmentMachine"; + + @Test + void mapRoutedEndpointsShouldResolveSingleEventPerLiteralEndpoint(@TempDir Path tempDir) throws Exception { + String source = """ + package com.example; + import java.util.Map; + public class OrderController { + OrderMachine sm; + public void payOrder() { sm.sendEvent(Routes.ROUTES.get("order.pay")); } + public void shipOrder() { sm.sendEvent(Routes.ROUTES.get("order.ship")); } + } + class Routes { + static final Map ROUTES = Map.of( + "order.pay", OrderEvent.PAY, + "order.ship", OrderEvent.SHIP + ); + } + enum OrderEvent { PAY, SHIP, CANCEL } + class OrderMachine { public void sendEvent(OrderEvent e) {} } + """; + Files.writeString(tempDir.resolve("App.java"), source); + CodebaseContext context = new CodebaseContext(); + context.scan(tempDir); + + HeuristicCallGraphEngine engine = new HeuristicCallGraphEngine(context); + TriggerPoint trigger = TriggerPoint.builder() + .className(ORDER_MACHINE) + .methodName("sendEvent") + .event("e") + .build(); + + EntryPoint payEntry = EntryPoint.builder().className(ORDER_CONTROLLER).methodName("payOrder").build(); + EntryPoint shipEntry = EntryPoint.builder().className(ORDER_CONTROLLER).methodName("shipOrder").build(); + + CallChain payChain = engine.findChains(List.of(payEntry), List.of(trigger)).get(0); + CallChain shipChain = engine.findChains(List.of(shipEntry), List.of(trigger)).get(0); + + assertThat(payChain.getTriggerPoint().getPolymorphicEvents()).containsExactly("OrderEvent.PAY"); + assertThat(shipChain.getTriggerPoint().getPolymorphicEvents()).containsExactly("OrderEvent.SHIP"); + + assertThat(context.getCache()).containsKey("heuristicCallGraph"); + + CallChain payAgain = engine.findChains(List.of(payEntry), List.of(trigger)).get(0); + CallChain shipAgain = engine.findChains(List.of(shipEntry), List.of(trigger)).get(0); + + assertThat(payAgain.getTriggerPoint().getPolymorphicEvents()) + .isEqualTo(payChain.getTriggerPoint().getPolymorphicEvents()); + assertThat(shipAgain.getTriggerPoint().getPolymorphicEvents()) + .isEqualTo(shipChain.getTriggerPoint().getPolymorphicEvents()); + } + + @Test + void separateMachinesShouldNotCrossContaminateMapRoutedResolution(@TempDir Path tempDir) throws Exception { + String source = """ + package com.example; + import java.util.Map; + public class OrderController { + OrderMachine orderMachine; + public void pay() { orderMachine.sendEvent(OrderRoutes.ROUTES.get("order.pay")); } + } + public class ShipmentController { + ShipmentMachine shipmentMachine; + public void dispatch() { shipmentMachine.sendEvent(ShipmentRoutes.ROUTES.get("shipment.dispatch")); } + } + class OrderRoutes { + static final Map ROUTES = Map.of("order.pay", OrderEvent.PAY); + } + class ShipmentRoutes { + static final Map ROUTES = Map.of("shipment.dispatch", ShipmentEvent.DISPATCH); + } + enum OrderEvent { PAY } + enum ShipmentEvent { DISPATCH } + class OrderMachine { public void sendEvent(OrderEvent e) {} } + class ShipmentMachine { public void sendEvent(ShipmentEvent e) {} } + """; + Files.writeString(tempDir.resolve("App.java"), source); + CodebaseContext context = new CodebaseContext(); + context.scan(tempDir); + + HeuristicCallGraphEngine engine = new HeuristicCallGraphEngine(context); + + EntryPoint orderEntry = EntryPoint.builder().className(ORDER_CONTROLLER).methodName("pay").build(); + EntryPoint shipmentEntry = EntryPoint.builder().className(SHIPMENT_CONTROLLER).methodName("dispatch").build(); + TriggerPoint orderTrigger = TriggerPoint.builder() + .className(ORDER_MACHINE) + .methodName("sendEvent") + .event("e") + .build(); + TriggerPoint shipmentTrigger = TriggerPoint.builder() + .className(SHIPMENT_MACHINE) + .methodName("sendEvent") + .event("e") + .build(); + + engine.findChains(List.of(orderEntry), List.of(orderTrigger)); + CallChain shipmentChain = engine.findChains(List.of(shipmentEntry), List.of(shipmentTrigger)).get(0); + assertThat(shipmentChain.getTriggerPoint().getPolymorphicEvents()).containsExactly("ShipmentEvent.DISPATCH"); + + CallChain orderChain = engine.findChains(List.of(orderEntry), List.of(orderTrigger)).get(0); + assertThat(orderChain.getTriggerPoint().getPolymorphicEvents()).containsExactly("OrderEvent.PAY"); + } +}