From b5f6e8bccc4e8f383cfcd4dc365d5da3239ec8db Mon Sep 17 00:00:00 2001 From: Daniel1464 Date: Sat, 30 May 2026 14:51:32 -0400 Subject: [PATCH 01/14] add files --- .../diagrams/MakeStateMachineGraphsTask.java | 245 +++++++++++++++++ .../diagrams/ReturnConditionAnalyzer.java | 256 ++++++++++++++++++ 2 files changed, 501 insertions(+) create mode 100644 src/main/java/org/wpilib/gradlerio/diagrams/MakeStateMachineGraphsTask.java create mode 100644 src/main/java/org/wpilib/gradlerio/diagrams/ReturnConditionAnalyzer.java diff --git a/src/main/java/org/wpilib/gradlerio/diagrams/MakeStateMachineGraphsTask.java b/src/main/java/org/wpilib/gradlerio/diagrams/MakeStateMachineGraphsTask.java new file mode 100644 index 00000000..4a093c4d --- /dev/null +++ b/src/main/java/org/wpilib/gradlerio/diagrams/MakeStateMachineGraphsTask.java @@ -0,0 +1,245 @@ +package org.wpilib.gradlerio.diagrams; + +import com.github.javaparser.JavaParser; +import com.github.javaparser.ParserConfiguration; +import com.github.javaparser.ast.body.MethodDeclaration; +import com.github.javaparser.ast.expr.Expression; +import com.github.javaparser.ast.expr.MethodCallExpr; +import com.github.javaparser.ast.expr.VariableDeclarationExpr; +import org.gradle.api.DefaultTask; +import org.gradle.api.provider.Property; +import org.gradle.api.tasks.Input; +import org.gradle.api.tasks.Optional; +import org.gradle.api.tasks.TaskAction; +import org.gradle.work.DisableCachingByDefault; + +import java.io.File; +import java.io.FileWriter; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Paths; +import java.util.*; +import java.util.regex.Pattern; + +@DisableCachingByDefault +public abstract class MakeStateMachineGraphsTask extends DefaultTask { + private record Transition(String fromState, String toState, String transitionCond) {} + private static class Diagram { + List transitions = new ArrayList<>(); + String initialState; + List stateDefinitionOrder = new ArrayList<>(); + } + + private final Pattern unusedVarPattern = Pattern.compile("(?<=[(,\\s])_(?=[\\s,)->])"); + + @Input + @Optional + public abstract Property getJavaRoot(); + + @TaskAction + public void run() throws IOException { + extractFromDirectory(getJavaRoot().getOrElse("src/main/java")); + } + + private void extractFromFile(File sourceFile, Map diagrams) throws IOException { + // 1. Initialize configuration + var config = new ParserConfiguration(); + config.setLanguageLevel(ParserConfiguration.LanguageLevel.JAVA_21); + config.setPreprocessUnicodeEscapes(true); + var parser = new JavaParser(config); + + var content = Files.readString(sourceFile.toPath()); + // JavaParser 3.28.1 (and earlier) has issues with unnamed variables (JEP 456). + // We replace them with a valid identifier before parsing. + content = unusedVarPattern + .matcher(content) + .replaceAll(res -> "unused_var_" + UUID.randomUUID().toString().substring(0, 8)); + + var cu = parser.parse(content).getResult().orElseThrow(); + cu.findAll(MethodDeclaration.class).forEach(method -> { + var annotationOpt = method.getAnnotationByName("MakeStateMachineGraph"); + if (!method.getTypeAsString().equals("StateMachine") || annotationOpt.isEmpty()) return; + + var variableDefs = method.findAll(VariableDeclarationExpr.class); + var stateMachineDefs = variableDefs + .stream() + .filter(v -> { + var def = v.getVariable(0).getInitializer(); + return def.isPresent() && + def.get().isObjectCreationExpr() && + def.get().asObjectCreationExpr().getType().asString().equals("StateMachine"); + }) + .toList(); + if (stateMachineDefs.isEmpty()) { + throw new RuntimeException("No StateMachine declaration found in " + method.getNameAsString()); + } else if (stateMachineDefs.size() > 1) { + throw new RuntimeException( + "Multiple StateMachine declarations found in " + method.getNameAsString() + + ". Currently, the @MakeStateMachineGraph annotation doesn't support multiple StateMachine declarations in the same method." + ); + } + + var stateMachineDef = stateMachineDefs.getFirst().getVariable(0).getInitializer().orElseThrow(); + var rawDiagramName = stateMachineDef.asObjectCreationExpr().getArgument(0); + if (!rawDiagramName.isStringLiteralExpr()) { + throw new RuntimeException( + "The diagram generator requires the name " + + "of the state machine in '" + method.getNameAsString() + + "' to be a single string literal (e.g. new StateMachine(\"My Diagram\"))." + ); + } + var diagramName = rawDiagramName.toString().substring(1, rawDiagramName.toString().length() - 1); + if (diagrams.containsKey(diagramName)) { + throw new RuntimeException("You already have a State Machine named '" + diagramName + "'"); + } + var diagram = new Diagram(); + diagrams.put(diagramName, diagram); + + variableDefs.stream() + .filter(v -> { + var initializer = v.getVariable(0).getInitializer(); + return initializer.isPresent() && + initializer.get().isMethodCallExpr() && + initializer.get().asMethodCallExpr().getNameAsString().equals("addState"); + }) + .forEachOrdered(v -> diagram.stateDefinitionOrder.add(v.getVariable(0).getNameAsString())); + + method.findAll(MethodCallExpr.class).forEach(call -> { + if (call.getNameAsString().equals("setInitialState")) { + diagram.initialState = call.getArguments().getFirst().orElseThrow().toString(); + return; + } + + boolean isWhen = call.getNameAsString().equals("when"); + boolean isWhenComplete = call.getNameAsString().equals("whenComplete"); + boolean isWhenCompleteAnd = call.getNameAsString().equals("whenCompleteAnd"); + if (!isWhen && !isWhenComplete && !isWhenCompleteAnd) return; + + var toStateCallOpt = call.getScope() + .filter(s -> s instanceof MethodCallExpr) + .map(s -> ((MethodCallExpr) s)); + if (toStateCallOpt.isEmpty()) return; + + var transitionCondExpr = call.getArguments().getFirst(); + var transitionCond = ""; + if (isWhenComplete) { + transitionCond = "when complete"; + } else { + var prefix = isWhen ? "" : "when complete and "; + var expr = transitionCondExpr.orElseThrow().toString(); + expr = expr.replace("() -> ", ""); + expr = expr.replace(".getAsBoolean()", ""); + if (expr.contains("||") || expr.contains(".or(")) { + expr = "(" + expr + ")"; + } + transitionCond = prefix + expr; + } + + var toStateCall = toStateCallOpt.get(); + var toState = toStateCall.getArguments().getFirst().orElseThrow(); + if (toStateCall.getNameAsString().equals("switchTo")) { + var fromState = toStateCall.getScope().orElseThrow().toString(); + if (toState.isLambdaExpr()) { + diagram.transitions.addAll( + transitionsFromLambdaExpr(toState, fromState, transitionCond) + ); + } else { + diagram.transitions.add(new Transition(fromState, toState.toString(), transitionCond)); + } + } else if (toStateCall.getNameAsString().equals("to")) { + // If it's just a regular "to", there must be a switchFromAny before it. + var argList = toStateCall.getScope() + .filter(s -> s instanceof MethodCallExpr) + .map(s -> ((MethodCallExpr) s)) + .map(MethodCallExpr::getArguments); + if (argList.isEmpty()) return; + for (var fromState: argList.get()) { + if (toState.isLambdaExpr()) { + diagram.transitions.addAll( + transitionsFromLambdaExpr(toState, fromState.toString(), transitionCond) + ); + } else { + var t = new Transition(fromState.toString(), toState.toString(), transitionCond); + diagram.transitions.add(t); + } + } + } + }); + }); + } + + private List transitionsFromLambdaExpr( + Expression toState, + String fromState, + String transitionCond + ) { + var transitions = new ArrayList(); + for (var entry: ReturnConditionAnalyzer.analyze(toState.asLambdaExpr()).entrySet()) { + var innerToState = entry.getKey(); + var additionalCond = entry.getValue(); + if (additionalCond.contains("||") || additionalCond.contains(".or(")) { + additionalCond = "(" + additionalCond + ")"; + } + var fullCond = transitionCond + " and " + additionalCond; + transitions.add(new Transition(fromState, innerToState, fullCond)); + } + return transitions; + } + + private void extractFromDirectory(String rootDir) throws IOException { + Map diagrams = new LinkedHashMap<>(); + + Files.walk(Paths.get(rootDir)) + .filter(path -> path.toString().endsWith(".java")) + .forEach(path -> { + try { + extractFromFile(path.toFile(), diagrams); + } catch (IOException e) { + System.err.println("Failed to parse: " + path + " — " + e.getMessage()); + } + }); + + for (var entry : diagrams.entrySet()) { + var methodName = entry.getKey(); + var diagram = entry.getValue(); + var deployDir = getJavaRoot().getOrElse("src/main/java/") + "../deploy/stateMachineGraphs/"; + var file = new File(deployDir + methodName + ".mermaid"); + if (file.getParentFile() != null) { + file.getParentFile().mkdirs(); + } + try (var writer = new FileWriter(file)) { + writer.write(generateMermaid(diagram)); + } + } + } + + private String generateMermaid(Diagram diagram) { + StringBuilder sb = new StringBuilder(); + sb.append("---\n"); + sb.append("state_definition_order: "); + sb.append(diagram.stateDefinitionOrder); + sb.append("\n---\n"); + sb.append("stateDiagram-v2\n"); + sb.append(" direction LR\n\n"); + for (var t : diagram.transitions) { + sb.append(" ") + .append(sanitize(t.fromState())) + .append(" --> ") + .append(sanitize(t.toState())) + .append(" : ") + .append(sanitize(t.transitionCond())) + .append("\n"); + } + if (diagram.initialState != null) { + sb.append("\n classDef initialState color: #00FF00"); + sb.append("\n class "); + sb.append(diagram.initialState); + sb.append(" initialState\n"); + } + return sb.toString(); + } + + private String sanitize(String s) { + return s.replaceAll("[\"'\\n\\r]", "").trim(); + } +} \ No newline at end of file diff --git a/src/main/java/org/wpilib/gradlerio/diagrams/ReturnConditionAnalyzer.java b/src/main/java/org/wpilib/gradlerio/diagrams/ReturnConditionAnalyzer.java new file mode 100644 index 00000000..82180299 --- /dev/null +++ b/src/main/java/org/wpilib/gradlerio/diagrams/ReturnConditionAnalyzer.java @@ -0,0 +1,256 @@ +package org.wpilib.gradlerio.diagrams; + +import com.github.javaparser.ast.expr.Expression; +import com.github.javaparser.ast.expr.LambdaExpr; +import com.github.javaparser.ast.expr.SwitchExpr; +import com.github.javaparser.ast.stmt.*; + +import java.util.*; +import java.util.stream.Collectors; + +/** + * Statically analyzes a lambda body and builds a map of: returnedVariableName -> condition string + * that must be true for that return to execute. + */ +public class ReturnConditionAnalyzer { + /** + * Analyzes a lambda expression to determine + * + * @param expr full source of a single lambda + * @return LinkedHashMap preserving insertion order, keyed by returned variable name + */ + public static Map analyze(LambdaExpr expr) { + Map result = new LinkedHashMap<>(); + var body = expr.getBody(); + if (body.isBlockStmt()) { + walkBlock(body.asBlockStmt().getStatements(), Collections.emptyList(), result); + } else { + walkStatement(body, Collections.emptyList(), result); + } + return result; + } + + // ----------------------------------------------------------------------- + // Recursive statement walker + // ----------------------------------------------------------------------- + + private static void walkBlock(List stmts, + List pathConds, + Map result) { + for (Statement stmt : stmts) { + walkStatement(stmt, pathConds, result); + } + } + + private static void walkStatement(Statement stmt, + List pathConds, + Map result) { + if (stmt.isReturnStmt()) { + handleReturn(stmt.asReturnStmt(), pathConds, result); + } else if (stmt.isIfStmt()) { + handleIf(stmt.asIfStmt(), pathConds, result); + } else if (stmt.isSwitchStmt()) { + handleSwitchStmt(stmt.asSwitchStmt(), pathConds, result); + } else if (stmt.isYieldStmt()) { + handleYield(stmt.asYieldStmt(), pathConds, result); + } else if (stmt.isBlockStmt()) { + walkBlock(stmt.asBlockStmt().getStatements(), pathConds, result); + } else if (stmt.isExpressionStmt()) { + var expr = stmt.asExpressionStmt().getExpression(); + if (!expr.isConditionalExpr()) return; + handleInlineConditional(expr, pathConds, result); + } else if (!stmt.isEmptyStmt()) { + // For any other statement type, descend into direct child statements + stmt.findAll(Statement.class, s -> s != stmt && s.getParentNode().map(p -> p == stmt).orElse(false)) + .forEach(child -> walkStatement(child, pathConds, result)); + } + } + + // ----------------------------------------------------------------------- + // Return — handles both plain returns and "return " + // ----------------------------------------------------------------------- + + private static void handleReturn(ReturnStmt ret, + List pathConds, + Map result) { + ret.getExpression().ifPresent(expr -> { + if (expr.isSwitchExpr()) { + // return switch (x) { case 1 -> a; ... } + handleSwitchExpr(expr.asSwitchExpr(), pathConds, result); + } else { + String varName = expr.toString().trim(); + String condition = joinConditions(pathConds); + mergeCondition(result, varName, condition.isEmpty() ? "true" : condition); + } + }); + } + + private static void handleInlineConditional(Expression expr, + List pathConds, + Map result) { + if (!expr.isConditionalExpr()) { + mergeCondition(result, expr.toString().trim(), joinConditions(pathConds)); + return; + } + var conditional = expr.asConditionalExpr(); + String cond = conditional.getCondition().toString(); + + handleInlineConditional(conditional.getThenExpr(), append(pathConds, cond), result); + handleInlineConditional(conditional.getElseExpr(), append(pathConds, negate(cond)), result); + } + + // ----------------------------------------------------------------------- + // Yield — used inside arrow-switch arms that contain a block + // e.g. case 1 -> { yield x; } + // ----------------------------------------------------------------------- + + private static void handleYield(YieldStmt yield, + List pathConds, + Map result) { + String varName = yield.getExpression().toString().trim(); + String condition = joinConditions(pathConds); + mergeCondition(result, varName, condition.isEmpty() ? "true" : condition); + } + + // ----------------------------------------------------------------------- + // Classic switch statement (case N: ... return x;) + // ----------------------------------------------------------------------- + + private static void handleSwitchStmt(SwitchStmt sw, + List pathConds, + Map result) { + walkSwitchEntries(sw.getSelector().toString(), sw.getEntries(), pathConds, result); + } + + // ----------------------------------------------------------------------- + // Modern switch expression (case N -> x / case N -> { yield x; }) + // ----------------------------------------------------------------------- + + private static void handleSwitchExpr(SwitchExpr sw, + List pathConds, + Map result) { + walkSwitchEntries(sw.getSelector().toString(), sw.getEntries(), pathConds, result); + } + + // ----------------------------------------------------------------------- + // Shared switch-entry walker (handles both SwitchStmt and SwitchExpr) + // ----------------------------------------------------------------------- + + private static void walkSwitchEntries(String selector, + List entries, + List pathConds, + Map result) { + // Collect every non-default label up front for building default negations + List nonDefaultLabels = entries.stream() + .flatMap(e -> e.getLabels().stream()) + .map(Expression::toString) + .toList(); + + for (SwitchEntry entry : entries) { + List entryConds; + + if (entry.getLabels().isEmpty()) { + // default / default -> + List negs = nonDefaultLabels.stream() + .map(l -> selector + " != " + l) + .collect(Collectors.toList()); + entryConds = append(pathConds, joinConditions(negs)); + } else { + // One or more comma-separated labels (arrow style) or stacked case: lines + List labelConds = entry.getLabels().stream() + .map(l -> selector + " == " + l) + .collect(Collectors.toList()); + String labelCond = labelConds.size() == 1 + ? labelConds.get(0) + : "(" + String.join(" || ", labelConds) + ")"; + entryConds = append(pathConds, labelCond); + } + + // Arrow entries with a direct expression (no braces) appear as an + // ExpressionStmt wrapping the value — treat that as an implicit yield. + if (isArrowExpressionEntry(entry)) { + String varName = entry.getStatements().get(0) + .asExpressionStmt().getExpression().toString().trim(); + String condition = joinConditions(entryConds); + mergeCondition(result, varName, condition.isEmpty() ? "true" : condition); + } else { + walkBlock(entry.getStatements(), entryConds, result); + } + } + } + + /** + * Returns true when this is an arrow-case entry whose body is a single + * bare expression (not a block and not a throw/yield statement). + * e.g. case 1 -> x; + */ + private static boolean isArrowExpressionEntry(SwitchEntry entry) { + if (entry.getType() != SwitchEntry.Type.EXPRESSION) return false; + return entry.getStatements().size() == 1 + && entry.getStatements().get(0).isExpressionStmt(); + } + + // ----------------------------------------------------------------------- + // If / else-if / else + // ----------------------------------------------------------------------- + + private static void handleIf(IfStmt ifStmt, + List pathConds, + Map result) { + String cond = ifStmt.getCondition().toString(); + + List thenConds = append(pathConds, cond); + walkStatement(ifStmt.getThenStmt(), thenConds, result); + + ifStmt.getElseStmt().ifPresent(elseStmt -> { + List elseConds = append(pathConds, negate(cond)); + walkStatement(elseStmt, elseConds, result); + }); + } + + // ----------------------------------------------------------------------- + // Helpers + // ----------------------------------------------------------------------- + + private static String negate(String cond) { + cond = cond.trim(); + if (cond.startsWith("!(") && cond.endsWith(")")) return cond.substring(2, cond.length() - 1); + if (cond.endsWith(".negate()")) return cond.substring(0, cond.length() - 9); + if (cond.contains(" == ")) return cond.replace(" == ", " != "); + if (cond.contains(" != ")) return cond.replace(" != ", " == "); + if (cond.contains(" >= ")) return cond.replace(" >= ", " < "); + if (cond.contains(" <= ")) return cond.replace(" <= ", " > "); + if (cond.contains(" > ")) return cond.replace(" > ", " <= "); + if (cond.contains(" < ")) return cond.replace(" < ", " >= "); + + boolean hasLogicalOps = cond.contains("&&") || cond.contains("||") || cond.contains(".and(") || cond.contains(".or("); + if (hasLogicalOps) return "!(" + cond + ")"; + if (cond.startsWith("!") && !cond.startsWith("!(")) { + return cond.substring(1); + } + return "!" + cond; + } + + private static List append(List conds, String newCond) { + if (newCond == null || newCond.isEmpty()) return conds; + List copy = new ArrayList<>(conds); + copy.add(newCond); + return copy; + } + + private static String joinConditions(List conds) { + if (conds.isEmpty()) return ""; + return String.join(" && ", conds); + } + + /** + * Merges a new condition for a variable into the result map. + * If the variable already has a condition, the two are OR-ed together: + * (existing) || (incoming) + * This handles the case where the same variable is returned from multiple branches. + */ + private static void mergeCondition(Map result, String varName, String newCond) { + result.merge(varName, newCond, (existing, incoming) -> + "(" + existing + ") || (" + incoming + ")"); + } +} From b0169987850fe1c44885c7c7a1ee8199e14e5f75 Mon Sep 17 00:00:00 2001 From: Daniel1464 Date: Sat, 30 May 2026 15:01:14 -0400 Subject: [PATCH 02/14] Cleanup human & AI-generated output --- .../diagrams/MakeStateMachineGraphsTask.java | 3 ++ .../diagrams/ReturnConditionAnalyzer.java | 50 ++----------------- 2 files changed, 8 insertions(+), 45 deletions(-) diff --git a/src/main/java/org/wpilib/gradlerio/diagrams/MakeStateMachineGraphsTask.java b/src/main/java/org/wpilib/gradlerio/diagrams/MakeStateMachineGraphsTask.java index 4a093c4d..ba4a8b1b 100644 --- a/src/main/java/org/wpilib/gradlerio/diagrams/MakeStateMachineGraphsTask.java +++ b/src/main/java/org/wpilib/gradlerio/diagrams/MakeStateMachineGraphsTask.java @@ -21,6 +21,9 @@ import java.util.*; import java.util.regex.Pattern; +/** + * A gradle task that generates state machine names + */ @DisableCachingByDefault public abstract class MakeStateMachineGraphsTask extends DefaultTask { private record Transition(String fromState, String toState, String transitionCond) {} diff --git a/src/main/java/org/wpilib/gradlerio/diagrams/ReturnConditionAnalyzer.java b/src/main/java/org/wpilib/gradlerio/diagrams/ReturnConditionAnalyzer.java index 82180299..1c8acbff 100644 --- a/src/main/java/org/wpilib/gradlerio/diagrams/ReturnConditionAnalyzer.java +++ b/src/main/java/org/wpilib/gradlerio/diagrams/ReturnConditionAnalyzer.java @@ -30,10 +30,6 @@ public static Map analyze(LambdaExpr expr) { return result; } - // ----------------------------------------------------------------------- - // Recursive statement walker - // ----------------------------------------------------------------------- - private static void walkBlock(List stmts, List pathConds, Map result) { @@ -66,10 +62,6 @@ private static void walkStatement(Statement stmt, } } - // ----------------------------------------------------------------------- - // Return — handles both plain returns and "return " - // ----------------------------------------------------------------------- - private static void handleReturn(ReturnStmt ret, List pathConds, Map result) { @@ -99,11 +91,6 @@ private static void handleInlineConditional(Expression expr, handleInlineConditional(conditional.getElseExpr(), append(pathConds, negate(cond)), result); } - // ----------------------------------------------------------------------- - // Yield — used inside arrow-switch arms that contain a block - // e.g. case 1 -> { yield x; } - // ----------------------------------------------------------------------- - private static void handleYield(YieldStmt yield, List pathConds, Map result) { @@ -112,30 +99,18 @@ private static void handleYield(YieldStmt yield, mergeCondition(result, varName, condition.isEmpty() ? "true" : condition); } - // ----------------------------------------------------------------------- - // Classic switch statement (case N: ... return x;) - // ----------------------------------------------------------------------- - private static void handleSwitchStmt(SwitchStmt sw, List pathConds, Map result) { walkSwitchEntries(sw.getSelector().toString(), sw.getEntries(), pathConds, result); } - // ----------------------------------------------------------------------- - // Modern switch expression (case N -> x / case N -> { yield x; }) - // ----------------------------------------------------------------------- - private static void handleSwitchExpr(SwitchExpr sw, List pathConds, Map result) { walkSwitchEntries(sw.getSelector().toString(), sw.getEntries(), pathConds, result); } - // ----------------------------------------------------------------------- - // Shared switch-entry walker (handles both SwitchStmt and SwitchExpr) - // ----------------------------------------------------------------------- - private static void walkSwitchEntries(String selector, List entries, List pathConds, @@ -150,7 +125,7 @@ private static void walkSwitchEntries(String selector, List entryConds; if (entry.getLabels().isEmpty()) { - // default / default -> + // "default" / "default ->" clauses List negs = nonDefaultLabels.stream() .map(l -> selector + " != " + l) .collect(Collectors.toList()); @@ -179,21 +154,15 @@ private static void walkSwitchEntries(String selector, } } - /** - * Returns true when this is an arrow-case entry whose body is a single - * bare expression (not a block and not a throw/yield statement). - * e.g. case 1 -> x; - */ + // Returns true when this is an arrow-case entry whose body is a single + // bare expression (not a block and not a throw/yield statement). + // e.g. case 1 -> x; private static boolean isArrowExpressionEntry(SwitchEntry entry) { if (entry.getType() != SwitchEntry.Type.EXPRESSION) return false; return entry.getStatements().size() == 1 && entry.getStatements().get(0).isExpressionStmt(); } - // ----------------------------------------------------------------------- - // If / else-if / else - // ----------------------------------------------------------------------- - private static void handleIf(IfStmt ifStmt, List pathConds, Map result) { @@ -208,10 +177,7 @@ private static void handleIf(IfStmt ifStmt, }); } - // ----------------------------------------------------------------------- - // Helpers - // ----------------------------------------------------------------------- - + // parses a boolean, Trigger, or BooleanSupplier, and returns the negation of that value. private static String negate(String cond) { cond = cond.trim(); if (cond.startsWith("!(") && cond.endsWith(")")) return cond.substring(2, cond.length() - 1); @@ -243,12 +209,6 @@ private static String joinConditions(List conds) { return String.join(" && ", conds); } - /** - * Merges a new condition for a variable into the result map. - * If the variable already has a condition, the two are OR-ed together: - * (existing) || (incoming) - * This handles the case where the same variable is returned from multiple branches. - */ private static void mergeCondition(Map result, String varName, String newCond) { result.merge(varName, newCond, (existing, incoming) -> "(" + existing + ") || (" + incoming + ")"); From 2b9e99041fec0a816415a8c5d02a24ec98acf504 Mon Sep 17 00:00:00 2001 From: Daniel1464 Date: Sat, 30 May 2026 16:09:37 -0400 Subject: [PATCH 03/14] More cleanup + format --- .../MakeStateMachineGraphsTask.java | 108 +++++++++--------- .../ReturnConditionAnalyzer.java | 77 ++++++++----- .../gradlerio/wpi/java/WPIJavaExtension.java | 8 +- 3 files changed, 111 insertions(+), 82 deletions(-) rename src/main/java/org/wpilib/gradlerio/{diagrams => graphgen}/MakeStateMachineGraphsTask.java (84%) rename src/main/java/org/wpilib/gradlerio/{diagrams => graphgen}/ReturnConditionAnalyzer.java (82%) diff --git a/src/main/java/org/wpilib/gradlerio/diagrams/MakeStateMachineGraphsTask.java b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java similarity index 84% rename from src/main/java/org/wpilib/gradlerio/diagrams/MakeStateMachineGraphsTask.java rename to src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java index ba4a8b1b..039f0e1a 100644 --- a/src/main/java/org/wpilib/gradlerio/diagrams/MakeStateMachineGraphsTask.java +++ b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java @@ -1,4 +1,5 @@ -package org.wpilib.gradlerio.diagrams; + +package org.wpilib.gradlerio.graphgen; import com.github.javaparser.JavaParser; import com.github.javaparser.ParserConfiguration; @@ -27,7 +28,7 @@ @DisableCachingByDefault public abstract class MakeStateMachineGraphsTask extends DefaultTask { private record Transition(String fromState, String toState, String transitionCond) {} - private static class Diagram { + private static class StateMachineGraph { List transitions = new ArrayList<>(); String initialState; List stateDefinitionOrder = new ArrayList<>(); @@ -44,13 +45,43 @@ public void run() throws IOException { extractFromDirectory(getJavaRoot().getOrElse("src/main/java")); } - private void extractFromFile(File sourceFile, Map diagrams) throws IOException { - // 1. Initialize configuration + private void extractFromDirectory(String rootDir) throws IOException { + var graphs = new LinkedHashMap(); + // Init configuration for javaparser var config = new ParserConfiguration(); config.setLanguageLevel(ParserConfiguration.LanguageLevel.JAVA_21); config.setPreprocessUnicodeEscapes(true); var parser = new JavaParser(config); + Files.walk(Paths.get(rootDir)) + .filter(path -> path.toString().endsWith(".java")) + .forEach(path -> { + try { + extractFromFile(parser, path.toFile(), graphs); + } catch (IOException e) { + System.err.println("Failed to parse: " + path + " — " + e.getMessage()); + } + }); + + for (var entry: graphs.entrySet()) { + var methodName = entry.getKey(); + var graph = entry.getValue(); + var deployDir = getJavaRoot().getOrElse("src/main/java/") + "../deploy/stateMachineGraphs/"; + var file = new File(deployDir + methodName + ".mermaid"); + if (file.getParentFile() != null) { + file.getParentFile().mkdirs(); + } + try (var writer = new FileWriter(file)) { + writer.write(generateMermaid(graph)); + } + } + } + + private void extractFromFile( + JavaParser parser, + File sourceFile, + Map graphs + ) throws IOException { var content = Files.readString(sourceFile.toPath()); // JavaParser 3.28.1 (and earlier) has issues with unnamed variables (JEP 456). // We replace them with a valid identifier before parsing. @@ -83,20 +114,20 @@ private void extractFromFile(File sourceFile, Map diagrams) thr } var stateMachineDef = stateMachineDefs.getFirst().getVariable(0).getInitializer().orElseThrow(); - var rawDiagramName = stateMachineDef.asObjectCreationExpr().getArgument(0); - if (!rawDiagramName.isStringLiteralExpr()) { + var rawGraphName = stateMachineDef.asObjectCreationExpr().getArgument(0); + if (!rawGraphName.isStringLiteralExpr()) { throw new RuntimeException( - "The diagram generator requires the name " + + "The graph generator requires the name " + "of the state machine in '" + method.getNameAsString() + - "' to be a single string literal (e.g. new StateMachine(\"My Diagram\"))." + "' to be a single string literal (e.g. new StateMachine(\"My Auto\"))." ); } - var diagramName = rawDiagramName.toString().substring(1, rawDiagramName.toString().length() - 1); - if (diagrams.containsKey(diagramName)) { - throw new RuntimeException("You already have a State Machine named '" + diagramName + "'"); + var graphName = rawGraphName.toString().substring(1, rawGraphName.toString().length() - 1); + if (graphs.containsKey(graphName)) { + throw new RuntimeException("You already have a State Machine named '" + graphName + "'"); } - var diagram = new Diagram(); - diagrams.put(diagramName, diagram); + var graph = new StateMachineGraph(); + graphs.put(graphName, graph); variableDefs.stream() .filter(v -> { @@ -105,11 +136,11 @@ private void extractFromFile(File sourceFile, Map diagrams) thr initializer.get().isMethodCallExpr() && initializer.get().asMethodCallExpr().getNameAsString().equals("addState"); }) - .forEachOrdered(v -> diagram.stateDefinitionOrder.add(v.getVariable(0).getNameAsString())); + .forEachOrdered(v -> graph.stateDefinitionOrder.add(v.getVariable(0).getNameAsString())); method.findAll(MethodCallExpr.class).forEach(call -> { if (call.getNameAsString().equals("setInitialState")) { - diagram.initialState = call.getArguments().getFirst().orElseThrow().toString(); + graph.initialState = call.getArguments().getFirst().orElseThrow().toString(); return; } @@ -143,11 +174,11 @@ private void extractFromFile(File sourceFile, Map diagrams) thr if (toStateCall.getNameAsString().equals("switchTo")) { var fromState = toStateCall.getScope().orElseThrow().toString(); if (toState.isLambdaExpr()) { - diagram.transitions.addAll( + graph.transitions.addAll( transitionsFromLambdaExpr(toState, fromState, transitionCond) ); } else { - diagram.transitions.add(new Transition(fromState, toState.toString(), transitionCond)); + graph.transitions.add(new Transition(fromState, toState.toString(), transitionCond)); } } else if (toStateCall.getNameAsString().equals("to")) { // If it's just a regular "to", there must be a switchFromAny before it. @@ -158,12 +189,12 @@ private void extractFromFile(File sourceFile, Map diagrams) thr if (argList.isEmpty()) return; for (var fromState: argList.get()) { if (toState.isLambdaExpr()) { - diagram.transitions.addAll( + graph.transitions.addAll( transitionsFromLambdaExpr(toState, fromState.toString(), transitionCond) ); } else { var t = new Transition(fromState.toString(), toState.toString(), transitionCond); - diagram.transitions.add(t); + graph.transitions.add(t); } } } @@ -189,42 +220,15 @@ private List transitionsFromLambdaExpr( return transitions; } - private void extractFromDirectory(String rootDir) throws IOException { - Map diagrams = new LinkedHashMap<>(); - - Files.walk(Paths.get(rootDir)) - .filter(path -> path.toString().endsWith(".java")) - .forEach(path -> { - try { - extractFromFile(path.toFile(), diagrams); - } catch (IOException e) { - System.err.println("Failed to parse: " + path + " — " + e.getMessage()); - } - }); - - for (var entry : diagrams.entrySet()) { - var methodName = entry.getKey(); - var diagram = entry.getValue(); - var deployDir = getJavaRoot().getOrElse("src/main/java/") + "../deploy/stateMachineGraphs/"; - var file = new File(deployDir + methodName + ".mermaid"); - if (file.getParentFile() != null) { - file.getParentFile().mkdirs(); - } - try (var writer = new FileWriter(file)) { - writer.write(generateMermaid(diagram)); - } - } - } - - private String generateMermaid(Diagram diagram) { + private String generateMermaid(StateMachineGraph graph) { StringBuilder sb = new StringBuilder(); sb.append("---\n"); sb.append("state_definition_order: "); - sb.append(diagram.stateDefinitionOrder); + sb.append(graph.stateDefinitionOrder); sb.append("\n---\n"); sb.append("stateDiagram-v2\n"); sb.append(" direction LR\n\n"); - for (var t : diagram.transitions) { + for (var t : graph.transitions) { sb.append(" ") .append(sanitize(t.fromState())) .append(" --> ") @@ -233,10 +237,10 @@ private String generateMermaid(Diagram diagram) { .append(sanitize(t.transitionCond())) .append("\n"); } - if (diagram.initialState != null) { + if (graph.initialState != null) { sb.append("\n classDef initialState color: #00FF00"); sb.append("\n class "); - sb.append(diagram.initialState); + sb.append(graph.initialState); sb.append(" initialState\n"); } return sb.toString(); @@ -245,4 +249,4 @@ private String generateMermaid(Diagram diagram) { private String sanitize(String s) { return s.replaceAll("[\"'\\n\\r]", "").trim(); } -} \ No newline at end of file +} diff --git a/src/main/java/org/wpilib/gradlerio/diagrams/ReturnConditionAnalyzer.java b/src/main/java/org/wpilib/gradlerio/graphgen/ReturnConditionAnalyzer.java similarity index 82% rename from src/main/java/org/wpilib/gradlerio/diagrams/ReturnConditionAnalyzer.java rename to src/main/java/org/wpilib/gradlerio/graphgen/ReturnConditionAnalyzer.java index 1c8acbff..33cfe191 100644 --- a/src/main/java/org/wpilib/gradlerio/diagrams/ReturnConditionAnalyzer.java +++ b/src/main/java/org/wpilib/gradlerio/graphgen/ReturnConditionAnalyzer.java @@ -1,4 +1,5 @@ -package org.wpilib.gradlerio.diagrams; + +package org.wpilib.gradlerio.graphgen; import com.github.javaparser.ast.expr.Expression; import com.github.javaparser.ast.expr.LambdaExpr; @@ -30,17 +31,21 @@ public static Map analyze(LambdaExpr expr) { return result; } - private static void walkBlock(List stmts, - List pathConds, - Map result) { + private static void walkBlock( + List stmts, + List pathConds, + Map result + ) { for (Statement stmt : stmts) { walkStatement(stmt, pathConds, result); } } - private static void walkStatement(Statement stmt, - List pathConds, - Map result) { + private static void walkStatement( + Statement stmt, + List pathConds, + Map result + ) { if (stmt.isReturnStmt()) { handleReturn(stmt.asReturnStmt(), pathConds, result); } else if (stmt.isIfStmt()) { @@ -62,9 +67,11 @@ private static void walkStatement(Statement stmt, } } - private static void handleReturn(ReturnStmt ret, - List pathConds, - Map result) { + private static void handleReturn( + ReturnStmt ret, + List pathConds, + Map result + ) { ret.getExpression().ifPresent(expr -> { if (expr.isSwitchExpr()) { // return switch (x) { case 1 -> a; ... } @@ -77,9 +84,11 @@ private static void handleReturn(ReturnStmt ret, }); } - private static void handleInlineConditional(Expression expr, - List pathConds, - Map result) { + private static void handleInlineConditional( + Expression expr, + List pathConds, + Map result + ) { if (!expr.isConditionalExpr()) { mergeCondition(result, expr.toString().trim(), joinConditions(pathConds)); return; @@ -91,30 +100,38 @@ private static void handleInlineConditional(Expression expr, handleInlineConditional(conditional.getElseExpr(), append(pathConds, negate(cond)), result); } - private static void handleYield(YieldStmt yield, - List pathConds, - Map result) { + private static void handleYield( + YieldStmt yield, + List pathConds, + Map result + ) { String varName = yield.getExpression().toString().trim(); String condition = joinConditions(pathConds); mergeCondition(result, varName, condition.isEmpty() ? "true" : condition); } - private static void handleSwitchStmt(SwitchStmt sw, - List pathConds, - Map result) { + private static void handleSwitchStmt( + SwitchStmt sw, + List pathConds, + Map result + ) { walkSwitchEntries(sw.getSelector().toString(), sw.getEntries(), pathConds, result); } - private static void handleSwitchExpr(SwitchExpr sw, - List pathConds, - Map result) { + private static void handleSwitchExpr( + SwitchExpr sw, + List pathConds, + Map result + ) { walkSwitchEntries(sw.getSelector().toString(), sw.getEntries(), pathConds, result); } - private static void walkSwitchEntries(String selector, - List entries, - List pathConds, - Map result) { + private static void walkSwitchEntries( + String selector, + List entries, + List pathConds, + Map result + ) { // Collect every non-default label up front for building default negations List nonDefaultLabels = entries.stream() .flatMap(e -> e.getLabels().stream()) @@ -163,9 +180,11 @@ private static boolean isArrowExpressionEntry(SwitchEntry entry) { && entry.getStatements().get(0).isExpressionStmt(); } - private static void handleIf(IfStmt ifStmt, - List pathConds, - Map result) { + private static void handleIf( + IfStmt ifStmt, + List pathConds, + Map result + ) { String cond = ifStmt.getCondition().toString(); List thenConds = append(pathConds, cond); diff --git a/src/main/java/org/wpilib/gradlerio/wpi/java/WPIJavaExtension.java b/src/main/java/org/wpilib/gradlerio/wpi/java/WPIJavaExtension.java index bd44d91d..b3d56a07 100644 --- a/src/main/java/org/wpilib/gradlerio/wpi/java/WPIJavaExtension.java +++ b/src/main/java/org/wpilib/gradlerio/wpi/java/WPIJavaExtension.java @@ -1,3 +1,4 @@ + package org.wpilib.gradlerio.wpi.java; import java.io.File; @@ -28,6 +29,7 @@ import org.gradle.jvm.tasks.Jar; import org.gradle.process.JavaForkOptions; import org.gradle.internal.os.OperatingSystem; +import org.wpilib.gradlerio.graphgen.MakeStateMachineGraphsTask; import org.wpilib.gradlerio.simulation.HalSimPair; import org.wpilib.gradlerio.simulation.JavaExternalSimulationTask; import org.wpilib.gradlerio.simulation.JavaSimulationTask; @@ -129,7 +131,6 @@ private void configureSimulationTask(JavaSimulationTask t, boolean debug, t.jvmArgs(jvmArgs); t.doFirst(new Action() { - @Override public void execute(Task task) { File ldpath = extract.get().getDestinationDirectory().get().getAsFile(); @@ -282,5 +283,10 @@ public WPIJavaExtension(Project project, SimulationExtension sim, WPIVersionsExt }); t.dependsOn(simTask); }); + + project.getTasks().register("generateGraphs", MakeStateMachineGraphsTask.class, task -> { + task.setGroup("GradleRIO"); + task.setDescription("Generates Graphs for state machine-returning methods annotated with @MakeStateMachineGraph."); + }); } } From f155cdeb0acd9faafef93236c96892c3d0184194 Mon Sep 17 00:00:00 2001 From: Daniel1464 Date: Sat, 30 May 2026 23:34:46 -0400 Subject: [PATCH 04/14] add exit state machine handling --- .../graphgen/MakeStateMachineGraphsTask.java | 21 ++++++++++--------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java index 039f0e1a..43a33c5e 100644 --- a/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java +++ b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java @@ -170,17 +170,19 @@ private void extractFromFile( } var toStateCall = toStateCallOpt.get(); - var toState = toStateCall.getArguments().getFirst().orElseThrow(); - if (toStateCall.getNameAsString().equals("switchTo")) { + var toStateExpr = toStateCall.getArguments().getFirst(); + var toState = toStateExpr.map(Expression::toString).orElse("Exit_State_Machine"); + boolean toStateIsLambda = toStateExpr.isPresent() && toStateExpr.get().isLambdaExpr(); + if (List.of("switchTo", "exitStateMachine").contains(toStateCall.getNameAsString())) { var fromState = toStateCall.getScope().orElseThrow().toString(); - if (toState.isLambdaExpr()) { + if (toStateIsLambda) { graph.transitions.addAll( - transitionsFromLambdaExpr(toState, fromState, transitionCond) + transitionsFromLambdaExpr(toStateExpr.get(), fromState, transitionCond) ); } else { - graph.transitions.add(new Transition(fromState, toState.toString(), transitionCond)); + graph.transitions.add(new Transition(fromState, toState, transitionCond)); } - } else if (toStateCall.getNameAsString().equals("to")) { + } else if (List.of("to", "toExitStateMachine").contains(toStateCall.getNameAsString())) { // If it's just a regular "to", there must be a switchFromAny before it. var argList = toStateCall.getScope() .filter(s -> s instanceof MethodCallExpr) @@ -188,13 +190,12 @@ private void extractFromFile( .map(MethodCallExpr::getArguments); if (argList.isEmpty()) return; for (var fromState: argList.get()) { - if (toState.isLambdaExpr()) { + if (toStateIsLambda) { graph.transitions.addAll( - transitionsFromLambdaExpr(toState, fromState.toString(), transitionCond) + transitionsFromLambdaExpr(toStateExpr.get(), fromState.toString(), transitionCond) ); } else { - var t = new Transition(fromState.toString(), toState.toString(), transitionCond); - graph.transitions.add(t); + graph.transitions.add(new Transition(fromState.toString(), toState, transitionCond)); } } } From 74ab8e42326431a86110375b7807d5c0aa503051 Mon Sep 17 00:00:00 2001 From: Daniel1464 Date: Sun, 31 May 2026 08:59:56 -0400 Subject: [PATCH 05/14] Add tests --- .../graphgen/MakeStateMachineGraphsTask.java | 16 +- .../MakeStateMachineGraphsTaskTest.java | 218 ++++++++++++++++++ 2 files changed, 229 insertions(+), 5 deletions(-) create mode 100644 src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java index 43a33c5e..4b584ed1 100644 --- a/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java +++ b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java @@ -40,12 +40,19 @@ private static class StateMachineGraph { @Optional public abstract Property getJavaRoot(); + @Input + @Optional + public abstract Property getDeployDirectory(); + @TaskAction public void run() throws IOException { - extractFromDirectory(getJavaRoot().getOrElse("src/main/java")); + extractFromDirectory( + getJavaRoot().getOrElse("src/main/java/"), + getDeployDirectory().getOrElse("src/main/deploy/") + ); } - private void extractFromDirectory(String rootDir) throws IOException { + private void extractFromDirectory(String javaRoot, String deployDir) throws IOException { var graphs = new LinkedHashMap(); // Init configuration for javaparser var config = new ParserConfiguration(); @@ -53,7 +60,7 @@ private void extractFromDirectory(String rootDir) throws IOException { config.setPreprocessUnicodeEscapes(true); var parser = new JavaParser(config); - Files.walk(Paths.get(rootDir)) + Files.walk(Paths.get(javaRoot)) .filter(path -> path.toString().endsWith(".java")) .forEach(path -> { try { @@ -66,8 +73,7 @@ private void extractFromDirectory(String rootDir) throws IOException { for (var entry: graphs.entrySet()) { var methodName = entry.getKey(); var graph = entry.getValue(); - var deployDir = getJavaRoot().getOrElse("src/main/java/") + "../deploy/stateMachineGraphs/"; - var file = new File(deployDir + methodName + ".mermaid"); + var file = new File(deployDir, "/stateMachineGraphs/" + methodName + ".mermaid"); if (file.getParentFile() != null) { file.getParentFile().mkdirs(); } diff --git a/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java b/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java new file mode 100644 index 00000000..ab59aaa2 --- /dev/null +++ b/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java @@ -0,0 +1,218 @@ +package org.wpilib.gradlerio.graphgen; + +import org.gradle.api.tasks.TaskProvider; +import org.gradle.testfixtures.ProjectBuilder; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; + +import static org.junit.jupiter.api.Assertions.*; + +class MakeStateMachineGraphsTaskTest { + + @TempDir + Path tempDir; + + private TaskProvider taskProvider; + private Path javaRoot; + private Path outputDir; + + @BeforeEach + void setup() throws IOException { + var project = ProjectBuilder.builder().withProjectDir(tempDir.toFile()).build(); + taskProvider = project.getTasks().register("makeGraphs", MakeStateMachineGraphsTask.class); + + javaRoot = tempDir.resolve("src/main/java").toAbsolutePath(); + Files.createDirectories(javaRoot); + + var deployDir = tempDir.resolve("src/main/deploy").toAbsolutePath(); + outputDir = deployDir.resolve("stateMachineGraphs"); + + // Use the same path format as default to avoid any normalization issues in the task + taskProvider.configure(task -> { + task.getJavaRoot().set(javaRoot.toString()); + task.getDeployDirectory().set(deployDir.toAbsolutePath().toString()); + }); + } + + @Test + void testBasicStateMachine() throws IOException { + String content = """ + package frc.robot; + + import org.wpilib.command3.Command; + import org.wpilib.gradlerio.graphgen.MakeStateMachineGraph; + + public class Robot { + @MakeStateMachineGraph + public StateMachine simpleAuto() { + System.out.println("Processing simpleAuto"); + StateMachine sm = new StateMachine("Simple Auto"); + var state1 = sm.addState(cmd("State1")); + var state2 = sm.addState(cmd("State2")); + + sm.setInitialState(state1); + + state1.switchTo(state2).when(() -> true); + state2.exitStateMachine().whenComplete(); + + return sm; + } + + private Command cmd(String name) { + return Command.noRequirements(coro -> while(true) coro.yield()).named(name); + } + } + """; + Files.writeString(javaRoot.resolve("Robot.java"), content); + + taskProvider.get().run(); + + Path mermaidFile = outputDir.resolve("Simple Auto.mermaid"); + assertTrue(Files.exists(mermaidFile), "Mermaid file should be generated"); + + String mermaidContent = Files.readString(mermaidFile); + assertTrue(mermaidContent.contains("state_definition_order: [state1, state2]")); + assertTrue(mermaidContent.contains("state1 --> state2 : true")); + assertTrue(mermaidContent.contains("state2 --> Exit_State_Machine : when complete")); + assertTrue(mermaidContent.contains("class state1 initialState")); + } + + @Test + void testComplexTransitions() throws IOException { + String content = """ + package frc.robot; + + import org.wpilib.command3.Command; + import org.wpilib.gradlerio.graphgen.MakeStateMachineGraph; + + public class Robot { + @MakeStateMachineGraph + public StateMachine complexAuto() { + StateMachine sm = new StateMachine("Complex Auto"); + var s1 = sm.addState(cmd("S1")); + var s2 = sm.addState(cmd("S2")); + var s3 = sm.addState(cmd("S3")); + + s1.switchTo(s2).whenCompleteAnd(() -> someCondition()); + sm.switchFromAny(s1, s2).to(s3).when(() -> emergency()); + sm.switchFromAny(s3).toExitStateMachine().when(() -> done()); + + return sm; + } + + private Command cmd(String name) { + return Command.noRequirements(coro -> while(true) coro.yield()).named(name); + } + } + """; + Files.writeString(javaRoot.resolve("Robot.java"), content); + + taskProvider.get().run(); + + String mermaidContent = Files.readString(outputDir.resolve("Complex Auto.mermaid")); + System.out.println(mermaidContent); + assertTrue(mermaidContent.contains("s1 --> s2 : when complete and someCondition()")); + assertTrue(mermaidContent.contains("s1 --> s3 : emergency()")); + assertTrue(mermaidContent.contains("s2 --> s3 : emergency()")); + assertTrue(mermaidContent.contains("s3 --> Exit_State_Machine : done()")); + } + + @Test + void testLambdaTransitions() throws IOException { + // This tests the interaction with ReturnConditionAnalyzer + String content = """ + package frc.robot; + + import org.wpilib.command3.Command; + import org.wpilib.gradlerio.graphgen.MakeStateMachineGraph; + + public class Robot { + @MakeStateMachineGraph + public StateMachine lambdaAuto() { + StateMachine sm = new StateMachine("Lambda Auto"); + var start = sm.addState(cmd("Start")); + var left = sm.addState(cmd("Left")); + var right = sm.addState(cmd("Right")); + + start.switchTo(cond -> { + if (cond) return left; + else return right; + }).when(() -> check()); + + return sm; + } + + private Command cmd(String name) { + return Command.noRequirements(coro -> while(true) coro.yield()).named(name); + } + } + """; + Files.writeString(javaRoot.resolve("Robot.java"), content); + + taskProvider.get().run(); + + String mermaidContent = Files.readString(outputDir.resolve("Lambda Auto.mermaid")); + assertTrue(mermaidContent.contains("start --> left : check() and cond")); + assertTrue(mermaidContent.contains("start --> right : check() and !cond")); + } + + @Test + void testUnnamedVariables() throws IOException { + String content = """ + package frc.robot; + + import org.wpilib.command3.Command; + import org.wpilib.gradlerio.graphgen.MakeStateMachineGraph; + + public class Robot { + @MakeStateMachineGraph + public StateMachine unnamedAuto() { + StateMachine sm = new StateMachine("Unnamed Auto"); + var s1 = sm.addState(cmd("S1")); + var s2 = sm.addState(cmd("S2")); + + s1.switchTo(_ -> s2).when(() -> true); + + return sm; + } + + private Command cmd(String name) { + return Command.noRequirements(coro -> while(true) coro.yield()).named(name); + } + } + """; + Files.writeString(javaRoot.resolve("Robot.java"), content); + + // Should not throw exception despite '_' being invalid in some Java versions for identifiers + assertDoesNotThrow(() -> taskProvider.get().run()); + + String mermaidContent = Files.readString(outputDir.resolve("Unnamed Auto.mermaid")); + System.out.println(mermaidContent); + assertTrue(mermaidContent.contains("s1 --> s2")); + } + + @Test + void testMultipleStateMachinesInMethodThrows() throws IOException { + String content = """ + package frc.robot; + import org.wpilib.gradlerio.graphgen.MakeStateMachineGraph; + + public class Robot { + @MakeStateMachineGraph + public StateMachine doubleSM() { + StateMachine sm1 = new StateMachine("SM1"); + StateMachine sm2 = new StateMachine("SM2"); + return sm1; + } + } + """; + Files.writeString(javaRoot.resolve("Robot.java"), content); + + assertThrows(RuntimeException.class, () -> taskProvider.get().run()); + } +} \ No newline at end of file From daa5db6737fe0acbbb1717727ffe4212891ea07d Mon Sep 17 00:00:00 2001 From: Daniel1464 Date: Sun, 31 May 2026 15:42:55 -0400 Subject: [PATCH 06/14] Fix tests + other small changes --- .../graphgen/MakeStateMachineGraphsTask.java | 20 ++++++++++++--- .../graphgen/ReturnConditionAnalyzer.java | 18 ++++++++++--- .../MakeStateMachineGraphsTaskTest.java | 25 +++++++++++++++++++ 3 files changed, 56 insertions(+), 7 deletions(-) diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java index 4b584ed1..7ee93606 100644 --- a/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java +++ b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java @@ -101,13 +101,24 @@ private void extractFromFile( if (!method.getTypeAsString().equals("StateMachine") || annotationOpt.isEmpty()) return; var variableDefs = method.findAll(VariableDeclarationExpr.class); + var stateSupplierTypes = List.of("Supplier", "Supplier"); + boolean hasStateSupplierVar = variableDefs.stream() + .filter(dec -> dec.getVariables().size() == 1) + .anyMatch(dec -> stateSupplierTypes.contains(dec.getVariable(0).getTypeAsString())); + if (hasStateSupplierVar) { + throw new RuntimeException( + "Methods annotated with @MakeStateMachineGraph cannot contain variables of type Supplier." + ); + } var stateMachineDefs = variableDefs .stream() .filter(v -> { + // we cannot check the type itself, because of var declarations - + // so, we check for the instantiation of the StateMachine class var def = v.getVariable(0).getInitializer(); return def.isPresent() && - def.get().isObjectCreationExpr() && - def.get().asObjectCreationExpr().getType().asString().equals("StateMachine"); + def.get().isObjectCreationExpr() && + def.get().asObjectCreationExpr().getType().asString().endsWith("StateMachine"); }) .toList(); if (stateMachineDefs.isEmpty()) { @@ -125,7 +136,7 @@ private void extractFromFile( throw new RuntimeException( "The graph generator requires the name " + "of the state machine in '" + method.getNameAsString() + - "' to be a single string literal (e.g. new StateMachine(\"My Auto\"))." + "' to be a simple string literal (e.g. new StateMachine(\"My Auto\"))." ); } var graphName = rawGraphName.toString().substring(1, rawGraphName.toString().length() - 1); @@ -193,6 +204,7 @@ private void extractFromFile( var argList = toStateCall.getScope() .filter(s -> s instanceof MethodCallExpr) .map(s -> ((MethodCallExpr) s)) + .filter(s -> s.getNameAsString().equals("switchFromAny")) .map(MethodCallExpr::getArguments); if (argList.isEmpty()) return; for (var fromState: argList.get()) { @@ -221,7 +233,7 @@ private List transitionsFromLambdaExpr( if (additionalCond.contains("||") || additionalCond.contains(".or(")) { additionalCond = "(" + additionalCond + ")"; } - var fullCond = transitionCond + " and " + additionalCond; + var fullCond = transitionCond + (additionalCond.isEmpty() ? "" : (" and " + additionalCond)); transitions.add(new Transition(fromState, innerToState, fullCond)); } return transitions; diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/ReturnConditionAnalyzer.java b/src/main/java/org/wpilib/gradlerio/graphgen/ReturnConditionAnalyzer.java index 33cfe191..dc757284 100644 --- a/src/main/java/org/wpilib/gradlerio/graphgen/ReturnConditionAnalyzer.java +++ b/src/main/java/org/wpilib/gradlerio/graphgen/ReturnConditionAnalyzer.java @@ -26,7 +26,7 @@ public static Map analyze(LambdaExpr expr) { if (body.isBlockStmt()) { walkBlock(body.asBlockStmt().getStatements(), Collections.emptyList(), result); } else { - walkStatement(body, Collections.emptyList(), result); + walkStatement(body, Collections.emptyList(), result, true); } return result; } @@ -45,6 +45,15 @@ private static void walkStatement( Statement stmt, List pathConds, Map result + ) { + walkStatement(stmt, pathConds, result, false); + } + + private static void walkStatement( + Statement stmt, + List pathConds, + Map result, + boolean isLambdaInline ) { if (stmt.isReturnStmt()) { handleReturn(stmt.asReturnStmt(), pathConds, result); @@ -58,8 +67,11 @@ private static void walkStatement( walkBlock(stmt.asBlockStmt().getStatements(), pathConds, result); } else if (stmt.isExpressionStmt()) { var expr = stmt.asExpressionStmt().getExpression(); - if (!expr.isConditionalExpr()) return; - handleInlineConditional(expr, pathConds, result); + if (expr.isConditionalExpr()) { + handleInlineConditional(expr, pathConds, result); + } else if (isLambdaInline) { + result.put(expr.toString(), ""); + } } else if (!stmt.isEmptyStmt()) { // For any other statement type, descend into direct child statements stmt.findAll(Statement.class, s -> s != stmt && s.getParentNode().map(p -> p == stmt).orElse(false)) diff --git a/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java b/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java index ab59aaa2..1ec3f79a 100644 --- a/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java +++ b/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java @@ -215,4 +215,29 @@ public StateMachine doubleSM() { assertThrows(RuntimeException.class, () -> taskProvider.get().run()); } + + @Test + void stateSupplierInStateMachineThrows() throws IOException { + String content = """ + package frc.robot; + import org.wpilib.gradlerio.graphgen.MakeStateMachineGraph; + + public class Robot { + @MakeStateMachineGraph + public StateMachine doubleSM() { + StateMachine sm1 = new StateMachine("SM1"); + State state1 = sm1.addState(cmd("state1")); + Supplier derived = () -> state1; + return sm1; + } + + private Command cmd(String name) { + return Command.noRequirements(coro -> while(true) coro.yield()).named(name); + } + } + """; + Files.writeString(javaRoot.resolve("Robot.java"), content); + + assertThrows(RuntimeException.class, () -> taskProvider.get().run()); + } } \ No newline at end of file From d916a94f4a6a0a568b2ecc7a7bcf852fb1742091 Mon Sep 17 00:00:00 2001 From: Daniel1464 Date: Mon, 1 Jun 2026 23:18:04 -0400 Subject: [PATCH 07/14] Fix early return statements and reduced parentheses stacking --- .../graphgen/MakeStateMachineGraphsTask.java | 4 +- .../graphgen/ReturnConditionAnalyzer.java | 61 ++++++++-- .../org/wpilib/gradlerio/graphgen/Utils.java | 10 ++ .../MakeStateMachineGraphsTaskTest.java | 111 +++++++++++++++--- 4 files changed, 161 insertions(+), 25 deletions(-) create mode 100644 src/main/java/org/wpilib/gradlerio/graphgen/Utils.java diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java index 7ee93606..04d3f36e 100644 --- a/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java +++ b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java @@ -181,7 +181,7 @@ private void extractFromFile( expr = expr.replace("() -> ", ""); expr = expr.replace(".getAsBoolean()", ""); if (expr.contains("||") || expr.contains(".or(")) { - expr = "(" + expr + ")"; + expr = Utils.addParentheses(expr); } transitionCond = prefix + expr; } @@ -231,7 +231,7 @@ private List transitionsFromLambdaExpr( var innerToState = entry.getKey(); var additionalCond = entry.getValue(); if (additionalCond.contains("||") || additionalCond.contains(".or(")) { - additionalCond = "(" + additionalCond + ")"; + additionalCond = Utils.addParentheses(additionalCond); } var fullCond = transitionCond + (additionalCond.isEmpty() ? "" : (" and " + additionalCond)); transitions.add(new Transition(fromState, innerToState, fullCond)); diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/ReturnConditionAnalyzer.java b/src/main/java/org/wpilib/gradlerio/graphgen/ReturnConditionAnalyzer.java index dc757284..e1396a0b 100644 --- a/src/main/java/org/wpilib/gradlerio/graphgen/ReturnConditionAnalyzer.java +++ b/src/main/java/org/wpilib/gradlerio/graphgen/ReturnConditionAnalyzer.java @@ -1,4 +1,3 @@ - package org.wpilib.gradlerio.graphgen; import com.github.javaparser.ast.expr.Expression; @@ -13,7 +12,7 @@ * Statically analyzes a lambda body and builds a map of: returnedVariableName -> condition string * that must be true for that return to execute. */ -public class ReturnConditionAnalyzer { +class ReturnConditionAnalyzer { /** * Analyzes a lambda expression to determine * @@ -36,11 +35,59 @@ private static void walkBlock( List pathConds, Map result ) { + List currentConds = pathConds; for (Statement stmt : stmts) { - walkStatement(stmt, pathConds, result); + walkStatement(stmt, currentConds, result); + // If this statement only covers some paths (e.g. if-without-else, switch-without-default), + // subsequent statements are only reached when none of those paths fired. + String fallThrough = fallThroughCondition(stmt); + if (fallThrough != null) { + currentConds = append(currentConds, fallThrough); + } } } + private static String fallThroughCondition(Statement stmt) { + if (stmt.isIfStmt()) { + // Walk the if/else-if chain collecting every condition; if there is a + // final else the chain is exhaustive and nothing falls through. + List branchConds = new ArrayList<>(); + Statement cur = stmt; + while (cur.isIfStmt()) { + IfStmt ifStmt = cur.asIfStmt(); + branchConds.add(ifStmt.getCondition().toString()); + if (ifStmt.getElseStmt().isEmpty()) { + // No final else — fall-through is possible when all conditions are false. + String combined = branchConds.size() == 1 + ? branchConds.get(0) + : Utils.addParentheses(String.join(" || ", branchConds)); + return negate(combined); + } + cur = ifStmt.getElseStmt().get(); + } + // Ended with a concrete else block — all paths covered, no fall-through. + return null; + } + + if (stmt.isSwitchStmt()) { + SwitchStmt sw = stmt.asSwitchStmt(); + boolean hasDefault = sw.getEntries().stream().anyMatch(e -> e.getLabels().isEmpty()); + if (hasDefault) return null; + // No default: falls through when no case label matches. + String selector = sw.getSelector().toString(); + List nonDefaultLabels = sw.getEntries().stream() + .flatMap(e -> e.getLabels().stream()) + .map(Expression::toString) + .toList(); + List negs = nonDefaultLabels.stream() + .map(l -> selector + " != " + l) + .collect(Collectors.toList()); + return joinConditions(negs); + } + + return null; + } + private static void walkStatement( Statement stmt, List pathConds, @@ -166,7 +213,7 @@ private static void walkSwitchEntries( .collect(Collectors.toList()); String labelCond = labelConds.size() == 1 ? labelConds.get(0) - : "(" + String.join(" || ", labelConds) + ")"; + : Utils.addParentheses(String.join(" || ", labelConds)); entryConds = append(pathConds, labelCond); } @@ -221,7 +268,7 @@ private static String negate(String cond) { if (cond.contains(" < ")) return cond.replace(" < ", " >= "); boolean hasLogicalOps = cond.contains("&&") || cond.contains("||") || cond.contains(".and(") || cond.contains(".or("); - if (hasLogicalOps) return "!(" + cond + ")"; + if (hasLogicalOps) return "!" + Utils.addParentheses(cond); if (cond.startsWith("!") && !cond.startsWith("!(")) { return cond.substring(1); } @@ -242,6 +289,6 @@ private static String joinConditions(List conds) { private static void mergeCondition(Map result, String varName, String newCond) { result.merge(varName, newCond, (existing, incoming) -> - "(" + existing + ") || (" + incoming + ")"); + Utils.addParentheses(existing) + " || " + Utils.addParentheses(incoming)); } -} +} \ No newline at end of file diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/Utils.java b/src/main/java/org/wpilib/gradlerio/graphgen/Utils.java new file mode 100644 index 00000000..18594e12 --- /dev/null +++ b/src/main/java/org/wpilib/gradlerio/graphgen/Utils.java @@ -0,0 +1,10 @@ +package org.wpilib.gradlerio.graphgen; + +class Utils { + public static String addParentheses(String condition) { + if ((condition.startsWith("!(") || condition.startsWith("(")) && condition.endsWith(")")) { + return condition; + } + return "(" + condition + ")"; + } +} diff --git a/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java b/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java index 1ec3f79a..7e8b3520 100644 --- a/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java +++ b/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java @@ -76,10 +76,20 @@ private Command cmd(String name) { assertTrue(Files.exists(mermaidFile), "Mermaid file should be generated"); String mermaidContent = Files.readString(mermaidFile); - assertTrue(mermaidContent.contains("state_definition_order: [state1, state2]")); - assertTrue(mermaidContent.contains("state1 --> state2 : true")); - assertTrue(mermaidContent.contains("state2 --> Exit_State_Machine : when complete")); - assertTrue(mermaidContent.contains("class state1 initialState")); + String expected = """ + --- + state_definition_order: [state1, state2] + --- + stateDiagram-v2 + direction LR + + state1 --> state2 : true + state2 --> Exit_State_Machine : when complete + + classDef initialState color: #00FF00 + class state1 initialState + """; + assertEquals(mermaidContent.strip(), expected.strip()); } @Test @@ -115,16 +125,23 @@ private Command cmd(String name) { taskProvider.get().run(); String mermaidContent = Files.readString(outputDir.resolve("Complex Auto.mermaid")); - System.out.println(mermaidContent); - assertTrue(mermaidContent.contains("s1 --> s2 : when complete and someCondition()")); - assertTrue(mermaidContent.contains("s1 --> s3 : emergency()")); - assertTrue(mermaidContent.contains("s2 --> s3 : emergency()")); - assertTrue(mermaidContent.contains("s3 --> Exit_State_Machine : done()")); + String expected = """ + --- + state_definition_order: [s1, s2, s3] + --- + stateDiagram-v2 + direction LR + + s1 --> s2 : when complete and someCondition() + s1 --> s3 : emergency() + s2 --> s3 : emergency() + s3 --> Exit_State_Machine : done() + """; + assertEquals(mermaidContent.strip(), expected.strip()); } @Test - void testLambdaTransitions() throws IOException { - // This tests the interaction with ReturnConditionAnalyzer + void testIfElseInLambdaTransition() throws IOException { String content = """ package frc.robot; @@ -157,8 +174,62 @@ private Command cmd(String name) { taskProvider.get().run(); String mermaidContent = Files.readString(outputDir.resolve("Lambda Auto.mermaid")); - assertTrue(mermaidContent.contains("start --> left : check() and cond")); - assertTrue(mermaidContent.contains("start --> right : check() and !cond")); + String expected = """ + --- + state_definition_order: [start, left, right] + --- + stateDiagram-v2 + direction LR + + start --> left : check() and cond + start --> right : check() and !cond + """; + assertEquals(mermaidContent.strip(), expected.strip()); + } + + @Test + void testEarlyReturnInLambdaTransition() throws IOException { + String content = """ + package frc.robot; + import org.wpilib.gradlerio.graphgen.MakeStateMachineGraph; + + public class Robot { + @MakeStateMachineGraph + public StateMachine doubleSM() { + StateMachine sm1 = new StateMachine("Early Return Auto"); + State a = sm1.addState(cmd("a")); + State b = sm1.addState(cmd("b")); + State c = sm1.addState(cmd("c")); + + a.switchTo(() -> { + switch (number) { + case 0, 1: + return b; + } + if (conditionC) return c; + return a; + }).whenComplete(); + return sm1; + } + } + """; + Files.writeString(javaRoot.resolve("Robot.java"), content); + + taskProvider.get().run(); + + String earlyReturnContent = Files.readString(outputDir.resolve("Early Return Auto.mermaid")); + String expected = """ + --- + state_definition_order: [a, b, c] + --- + stateDiagram-v2 + direction LR + + a --> b : when complete and (number == 0 || number == 1) + a --> c : when complete and number != 0 && number != 1 && conditionC + a --> a : when complete and number != 0 && number != 1 && !conditionC + """; + assertEquals(earlyReturnContent.strip(), expected.strip()); } @Test @@ -190,10 +261,18 @@ private Command cmd(String name) { // Should not throw exception despite '_' being invalid in some Java versions for identifiers assertDoesNotThrow(() -> taskProvider.get().run()); - + String mermaidContent = Files.readString(outputDir.resolve("Unnamed Auto.mermaid")); - System.out.println(mermaidContent); - assertTrue(mermaidContent.contains("s1 --> s2")); + String expected = """ + --- + state_definition_order: [s1, s2] + --- + stateDiagram-v2 + direction LR + + s1 --> s2 : true + """; + assertEquals(mermaidContent.strip(), expected.strip()); } @Test From 5080e724e7eb449c2c6066336660fb7e99b5da15 Mon Sep 17 00:00:00 2001 From: Daniel1464 Date: Wed, 3 Jun 2026 16:26:29 -0400 Subject: [PATCH 08/14] Return condition analysis overhaul --- .../gradlerio/graphgen/CodeBlockAnalyzer.java | 126 ++++++++ .../graphgen/ExpressionAnalyzer.java | 75 +++++ .../graphgen/MakeStateMachineGraphsTask.java | 29 +- .../graphgen/ReturnConditionAnalyzer.java | 294 ------------------ .../org/wpilib/gradlerio/graphgen/Utils.java | 59 +++- .../MakeStateMachineGraphsTaskTest.java | 86 ++++- 6 files changed, 346 insertions(+), 323 deletions(-) create mode 100644 src/main/java/org/wpilib/gradlerio/graphgen/CodeBlockAnalyzer.java create mode 100644 src/main/java/org/wpilib/gradlerio/graphgen/ExpressionAnalyzer.java delete mode 100644 src/main/java/org/wpilib/gradlerio/graphgen/ReturnConditionAnalyzer.java diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/CodeBlockAnalyzer.java b/src/main/java/org/wpilib/gradlerio/graphgen/CodeBlockAnalyzer.java new file mode 100644 index 00000000..bfb43ba8 --- /dev/null +++ b/src/main/java/org/wpilib/gradlerio/graphgen/CodeBlockAnalyzer.java @@ -0,0 +1,126 @@ +package org.wpilib.gradlerio.graphgen; + +import com.github.javaparser.ast.Node; +import com.github.javaparser.ast.stmt.BlockStmt; +import com.github.javaparser.ast.stmt.Statement; + +import java.util.*; +import java.util.stream.Collectors; + +class CodeBlockAnalyzer { + static Map analyze(BlockStmt block) { + var rawReturns = new HashMap>(); + analyzeStatement(block, "true", rawReturns); + return rawReturns.entrySet().stream() + .collect(Collectors.toMap( + Map.Entry::getKey, + entry -> Utils.joinWithOr(entry.getValue()) + )); + } + + private static List analyzeStatement( + Statement stmt, + String pathCondition, + Map> returnsMap + ) { + var fallThroughConditions = new ArrayList(); + + if (stmt.isBlockStmt()) { + var currentPaths = Collections.singletonList(pathCondition); + for (var subStmt: stmt.asBlockStmt().getStatements()) { + var nextPaths = new ArrayList(); + for (var path: currentPaths) { + nextPaths.addAll(analyzeStatement(subStmt, path, returnsMap)); + } + currentPaths = nextPaths; + if (currentPaths.isEmpty()) { + break; // Block always returns before hitting subsequent lines + } + } + fallThroughConditions.addAll(currentPaths); + } else if (stmt.isReturnStmt()) { + var returnExpr = stmt.asReturnStmt().getExpression().orElseThrow(); + ExpressionAnalyzer.analyze(returnExpr, pathCondition, returnsMap); + } else if (stmt.isIfStmt()) { + var ifStmt = stmt.asIfStmt(); + var conditionStr = ifStmt.getCondition().toString(); + var thenCondition = Utils.joinWithAnd(pathCondition, conditionStr); + var elseCondition = Utils.joinWithAnd(pathCondition, Utils.negate(conditionStr)); + + fallThroughConditions.addAll(analyzeStatement(ifStmt.getThenStmt(), thenCondition, returnsMap)); + if (ifStmt.getElseStmt().isPresent()) { + fallThroughConditions.addAll(analyzeStatement(ifStmt.getElseStmt().get(), elseCondition, returnsMap)); + } else { + fallThroughConditions.add(elseCondition); + } + } else if (stmt.isSwitchStmt()) { + var switchStmt = stmt.asSwitchStmt(); + var selector = switchStmt.getSelector().toString(); + var switchExitPaths = new ArrayList(); + var pendingFallThrough = new ArrayList(); + var selectorMatches = new ArrayList(); + + for (var entry: switchStmt.getEntries()) { + List labels = entry.getLabels().stream().map(Node::toString).toList(); + String entryDirectCondition; + + if (labels.isEmpty()) { + // 'default' case handles whatever hasn't matched prior explicit conditions + if (selectorMatches.isEmpty()) { + entryDirectCondition = pathCondition; + } else { + var combinedMatches = Utils.joinWithOr(selectorMatches); + entryDirectCondition = Utils.joinWithAnd(pathCondition, Utils.negate(combinedMatches)); + } + } else { + var matchCondition = Utils.joinWithOr( + labels.stream().map(label -> selector + " == " + label).toList() + ); + entryDirectCondition = Utils.joinWithAnd(pathCondition, matchCondition); + selectorMatches.add(matchCondition); + } + + // Current case entry path context = Direct match path + pending fall-through paths from previous case + var caseIncomingPaths = new ArrayList(); + caseIncomingPaths.add(entryDirectCondition); + caseIncomingPaths.addAll(pendingFallThrough); + pendingFallThrough.clear(); // Consumed by this case entry + + // Process internal statement sequences within this switch branch + var currentCasePaths = caseIncomingPaths; + for (var subStmt: entry.getStatements()) { + if (subStmt.isBreakStmt()) { + // 'break' safely escapes the switch completely + switchExitPaths.addAll(currentCasePaths); + currentCasePaths = new ArrayList<>(); + break; + } + var nextCasePaths = new ArrayList(); + for (var path: currentCasePaths) { + nextCasePaths.addAll(analyzeStatement(subStmt, path, returnsMap)); + } + currentCasePaths = nextCasePaths; + if (currentCasePaths.isEmpty()) break; + } + + // Anything still active cascades directly to the next case block + pendingFallThrough.addAll(currentCasePaths); + } + + // Clean up left over paths + switchExitPaths.addAll(pendingFallThrough); + + // If no default branch existed, values not explicit in cases fall through past the switch block + boolean hasDefault = switchStmt.getEntries().stream().anyMatch(e -> e.getLabels().isEmpty()); + if (!hasDefault && !selectorMatches.isEmpty()) { + var combinedMatches = Utils.joinWithOr(selectorMatches); + switchExitPaths.add(Utils.joinWithAnd(pathCondition, Utils.negate(combinedMatches))); + } + fallThroughConditions.addAll(switchExitPaths); + } else { + fallThroughConditions.add(pathCondition); + } + + return fallThroughConditions; + } +} \ No newline at end of file diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/ExpressionAnalyzer.java b/src/main/java/org/wpilib/gradlerio/graphgen/ExpressionAnalyzer.java new file mode 100644 index 00000000..0887a116 --- /dev/null +++ b/src/main/java/org/wpilib/gradlerio/graphgen/ExpressionAnalyzer.java @@ -0,0 +1,75 @@ +package org.wpilib.gradlerio.graphgen; + +import com.github.javaparser.ast.expr.Expression; +import com.github.javaparser.ast.expr.SwitchExpr; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +class ExpressionAnalyzer { + static Map analyze(Expression expr) { + var rawReturns = new HashMap>(); + ExpressionAnalyzer.analyze(expr, "true", rawReturns); + return rawReturns.entrySet().stream() + .collect(Collectors.toMap( + Map.Entry::getKey, + entry -> Utils.joinWithOr(entry.getValue()) + )); + } + + static void analyze( + Expression expr, + String pathCondition, + Map> returnsMap + ) { + if (expr.isSwitchExpr()) { + handleSwitch(expr.asSwitchExpr(), pathCondition, returnsMap); + } else if (expr.isConditionalExpr()) { + var conditional = expr.asConditionalExpr(); + var condition = conditional.getCondition().toString(); + analyze(conditional.getThenExpr(), Utils.joinWithAnd(pathCondition, condition), returnsMap); + analyze(conditional.getElseExpr(), Utils.joinWithAnd(pathCondition, Utils.negate(condition)), returnsMap); + } else { + returnsMap.computeIfAbsent(expr.toString(), k -> new ArrayList<>()).add(pathCondition); + } + } + + private static void handleSwitch( + SwitchExpr sw, + String pathCondition, + Map> returnsMap + ) { + var selector = sw.getSelector().toString(); + var entries = sw.getEntries(); + var defaultClauses = new ArrayList(); + for (var entry: entries) { + String matchCondition; + if (entry.getLabels().isEmpty()) { + matchCondition = Utils.joinWithAnd(defaultClauses); + } else { + // One or more comma-separated labels (arrow style) or stacked case: lines + matchCondition = Utils.joinWithOr( + entry.getLabels().stream().map(label -> selector + " == " + label).toList() + ); + defaultClauses.add(Utils.negate(matchCondition)); + } + + var stmt = entry.getStatements().get(0); + if (stmt.isBlockStmt()) { + throw new RuntimeException("Yield Statements within switch blocks are not supported for @MakeStateMachineGraph."); + } + var expression = stmt.asExpressionStmt().getExpression(); + if (expression.isConditionalExpr()) { + analyze(expression.asConditionalExpr(), Utils.joinWithAnd(pathCondition, matchCondition), returnsMap); + } else if (expression.isSwitchExpr()) { + handleSwitch(expression.asSwitchExpr(), Utils.joinWithAnd(pathCondition, matchCondition), returnsMap); + } else { + var varName = expression.toString().trim(); + returnsMap.computeIfAbsent(varName, k -> new ArrayList<>()).add(pathCondition); + } + } + } +} diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java index 04d3f36e..da7bca4d 100644 --- a/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java +++ b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java @@ -176,14 +176,10 @@ private void extractFromFile( if (isWhenComplete) { transitionCond = "when complete"; } else { - var prefix = isWhen ? "" : "when complete and "; var expr = transitionCondExpr.orElseThrow().toString(); expr = expr.replace("() -> ", ""); expr = expr.replace(".getAsBoolean()", ""); - if (expr.contains("||") || expr.contains(".or(")) { - expr = Utils.addParentheses(expr); - } - transitionCond = prefix + expr; + transitionCond = isWhen ? expr : Utils.joinWithAnd("when complete", expr); } var toStateCall = toStateCallOpt.get(); @@ -227,13 +223,14 @@ private List transitionsFromLambdaExpr( String transitionCond ) { var transitions = new ArrayList(); - for (var entry: ReturnConditionAnalyzer.analyze(toState.asLambdaExpr()).entrySet()) { + var body = toState.asLambdaExpr().getBody(); + var returnConditions = body.isBlockStmt() + ? CodeBlockAnalyzer.analyze(body.asBlockStmt()) + : ExpressionAnalyzer.analyze(body.asExpressionStmt().getExpression()); + for (var entry: returnConditions.entrySet()) { var innerToState = entry.getKey(); var additionalCond = entry.getValue(); - if (additionalCond.contains("||") || additionalCond.contains(".or(")) { - additionalCond = Utils.addParentheses(additionalCond); - } - var fullCond = transitionCond + (additionalCond.isEmpty() ? "" : (" and " + additionalCond)); + var fullCond = Utils.joinWithAnd(transitionCond, additionalCond); transitions.add(new Transition(fromState, innerToState, fullCond)); } return transitions; @@ -253,7 +250,7 @@ private String generateMermaid(StateMachineGraph graph) { .append(" --> ") .append(sanitize(t.toState())) .append(" : ") - .append(sanitize(t.transitionCond())) + .append(sanitizeTransitionCond(t.transitionCond())) .append("\n"); } if (graph.initialState != null) { @@ -265,6 +262,16 @@ private String generateMermaid(StateMachineGraph graph) { return sb.toString(); } + private String sanitizeTransitionCond(String condition) { + condition = sanitize(condition); + if (condition.isEmpty() || condition.equals("true")) { + condition = "instant"; + } else if (condition.equals("false")) { + condition = "never"; + } + return condition; + } + private String sanitize(String s) { return s.replaceAll("[\"'\\n\\r]", "").trim(); } diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/ReturnConditionAnalyzer.java b/src/main/java/org/wpilib/gradlerio/graphgen/ReturnConditionAnalyzer.java deleted file mode 100644 index e1396a0b..00000000 --- a/src/main/java/org/wpilib/gradlerio/graphgen/ReturnConditionAnalyzer.java +++ /dev/null @@ -1,294 +0,0 @@ -package org.wpilib.gradlerio.graphgen; - -import com.github.javaparser.ast.expr.Expression; -import com.github.javaparser.ast.expr.LambdaExpr; -import com.github.javaparser.ast.expr.SwitchExpr; -import com.github.javaparser.ast.stmt.*; - -import java.util.*; -import java.util.stream.Collectors; - -/** - * Statically analyzes a lambda body and builds a map of: returnedVariableName -> condition string - * that must be true for that return to execute. - */ -class ReturnConditionAnalyzer { - /** - * Analyzes a lambda expression to determine - * - * @param expr full source of a single lambda - * @return LinkedHashMap preserving insertion order, keyed by returned variable name - */ - public static Map analyze(LambdaExpr expr) { - Map result = new LinkedHashMap<>(); - var body = expr.getBody(); - if (body.isBlockStmt()) { - walkBlock(body.asBlockStmt().getStatements(), Collections.emptyList(), result); - } else { - walkStatement(body, Collections.emptyList(), result, true); - } - return result; - } - - private static void walkBlock( - List stmts, - List pathConds, - Map result - ) { - List currentConds = pathConds; - for (Statement stmt : stmts) { - walkStatement(stmt, currentConds, result); - // If this statement only covers some paths (e.g. if-without-else, switch-without-default), - // subsequent statements are only reached when none of those paths fired. - String fallThrough = fallThroughCondition(stmt); - if (fallThrough != null) { - currentConds = append(currentConds, fallThrough); - } - } - } - - private static String fallThroughCondition(Statement stmt) { - if (stmt.isIfStmt()) { - // Walk the if/else-if chain collecting every condition; if there is a - // final else the chain is exhaustive and nothing falls through. - List branchConds = new ArrayList<>(); - Statement cur = stmt; - while (cur.isIfStmt()) { - IfStmt ifStmt = cur.asIfStmt(); - branchConds.add(ifStmt.getCondition().toString()); - if (ifStmt.getElseStmt().isEmpty()) { - // No final else — fall-through is possible when all conditions are false. - String combined = branchConds.size() == 1 - ? branchConds.get(0) - : Utils.addParentheses(String.join(" || ", branchConds)); - return negate(combined); - } - cur = ifStmt.getElseStmt().get(); - } - // Ended with a concrete else block — all paths covered, no fall-through. - return null; - } - - if (stmt.isSwitchStmt()) { - SwitchStmt sw = stmt.asSwitchStmt(); - boolean hasDefault = sw.getEntries().stream().anyMatch(e -> e.getLabels().isEmpty()); - if (hasDefault) return null; - // No default: falls through when no case label matches. - String selector = sw.getSelector().toString(); - List nonDefaultLabels = sw.getEntries().stream() - .flatMap(e -> e.getLabels().stream()) - .map(Expression::toString) - .toList(); - List negs = nonDefaultLabels.stream() - .map(l -> selector + " != " + l) - .collect(Collectors.toList()); - return joinConditions(negs); - } - - return null; - } - - private static void walkStatement( - Statement stmt, - List pathConds, - Map result - ) { - walkStatement(stmt, pathConds, result, false); - } - - private static void walkStatement( - Statement stmt, - List pathConds, - Map result, - boolean isLambdaInline - ) { - if (stmt.isReturnStmt()) { - handleReturn(stmt.asReturnStmt(), pathConds, result); - } else if (stmt.isIfStmt()) { - handleIf(stmt.asIfStmt(), pathConds, result); - } else if (stmt.isSwitchStmt()) { - handleSwitchStmt(stmt.asSwitchStmt(), pathConds, result); - } else if (stmt.isYieldStmt()) { - handleYield(stmt.asYieldStmt(), pathConds, result); - } else if (stmt.isBlockStmt()) { - walkBlock(stmt.asBlockStmt().getStatements(), pathConds, result); - } else if (stmt.isExpressionStmt()) { - var expr = stmt.asExpressionStmt().getExpression(); - if (expr.isConditionalExpr()) { - handleInlineConditional(expr, pathConds, result); - } else if (isLambdaInline) { - result.put(expr.toString(), ""); - } - } else if (!stmt.isEmptyStmt()) { - // For any other statement type, descend into direct child statements - stmt.findAll(Statement.class, s -> s != stmt && s.getParentNode().map(p -> p == stmt).orElse(false)) - .forEach(child -> walkStatement(child, pathConds, result)); - } - } - - private static void handleReturn( - ReturnStmt ret, - List pathConds, - Map result - ) { - ret.getExpression().ifPresent(expr -> { - if (expr.isSwitchExpr()) { - // return switch (x) { case 1 -> a; ... } - handleSwitchExpr(expr.asSwitchExpr(), pathConds, result); - } else { - String varName = expr.toString().trim(); - String condition = joinConditions(pathConds); - mergeCondition(result, varName, condition.isEmpty() ? "true" : condition); - } - }); - } - - private static void handleInlineConditional( - Expression expr, - List pathConds, - Map result - ) { - if (!expr.isConditionalExpr()) { - mergeCondition(result, expr.toString().trim(), joinConditions(pathConds)); - return; - } - var conditional = expr.asConditionalExpr(); - String cond = conditional.getCondition().toString(); - - handleInlineConditional(conditional.getThenExpr(), append(pathConds, cond), result); - handleInlineConditional(conditional.getElseExpr(), append(pathConds, negate(cond)), result); - } - - private static void handleYield( - YieldStmt yield, - List pathConds, - Map result - ) { - String varName = yield.getExpression().toString().trim(); - String condition = joinConditions(pathConds); - mergeCondition(result, varName, condition.isEmpty() ? "true" : condition); - } - - private static void handleSwitchStmt( - SwitchStmt sw, - List pathConds, - Map result - ) { - walkSwitchEntries(sw.getSelector().toString(), sw.getEntries(), pathConds, result); - } - - private static void handleSwitchExpr( - SwitchExpr sw, - List pathConds, - Map result - ) { - walkSwitchEntries(sw.getSelector().toString(), sw.getEntries(), pathConds, result); - } - - private static void walkSwitchEntries( - String selector, - List entries, - List pathConds, - Map result - ) { - // Collect every non-default label up front for building default negations - List nonDefaultLabels = entries.stream() - .flatMap(e -> e.getLabels().stream()) - .map(Expression::toString) - .toList(); - - for (SwitchEntry entry : entries) { - List entryConds; - - if (entry.getLabels().isEmpty()) { - // "default" / "default ->" clauses - List negs = nonDefaultLabels.stream() - .map(l -> selector + " != " + l) - .collect(Collectors.toList()); - entryConds = append(pathConds, joinConditions(negs)); - } else { - // One or more comma-separated labels (arrow style) or stacked case: lines - List labelConds = entry.getLabels().stream() - .map(l -> selector + " == " + l) - .collect(Collectors.toList()); - String labelCond = labelConds.size() == 1 - ? labelConds.get(0) - : Utils.addParentheses(String.join(" || ", labelConds)); - entryConds = append(pathConds, labelCond); - } - - // Arrow entries with a direct expression (no braces) appear as an - // ExpressionStmt wrapping the value — treat that as an implicit yield. - if (isArrowExpressionEntry(entry)) { - String varName = entry.getStatements().get(0) - .asExpressionStmt().getExpression().toString().trim(); - String condition = joinConditions(entryConds); - mergeCondition(result, varName, condition.isEmpty() ? "true" : condition); - } else { - walkBlock(entry.getStatements(), entryConds, result); - } - } - } - - // Returns true when this is an arrow-case entry whose body is a single - // bare expression (not a block and not a throw/yield statement). - // e.g. case 1 -> x; - private static boolean isArrowExpressionEntry(SwitchEntry entry) { - if (entry.getType() != SwitchEntry.Type.EXPRESSION) return false; - return entry.getStatements().size() == 1 - && entry.getStatements().get(0).isExpressionStmt(); - } - - private static void handleIf( - IfStmt ifStmt, - List pathConds, - Map result - ) { - String cond = ifStmt.getCondition().toString(); - - List thenConds = append(pathConds, cond); - walkStatement(ifStmt.getThenStmt(), thenConds, result); - - ifStmt.getElseStmt().ifPresent(elseStmt -> { - List elseConds = append(pathConds, negate(cond)); - walkStatement(elseStmt, elseConds, result); - }); - } - - // parses a boolean, Trigger, or BooleanSupplier, and returns the negation of that value. - private static String negate(String cond) { - cond = cond.trim(); - if (cond.startsWith("!(") && cond.endsWith(")")) return cond.substring(2, cond.length() - 1); - if (cond.endsWith(".negate()")) return cond.substring(0, cond.length() - 9); - if (cond.contains(" == ")) return cond.replace(" == ", " != "); - if (cond.contains(" != ")) return cond.replace(" != ", " == "); - if (cond.contains(" >= ")) return cond.replace(" >= ", " < "); - if (cond.contains(" <= ")) return cond.replace(" <= ", " > "); - if (cond.contains(" > ")) return cond.replace(" > ", " <= "); - if (cond.contains(" < ")) return cond.replace(" < ", " >= "); - - boolean hasLogicalOps = cond.contains("&&") || cond.contains("||") || cond.contains(".and(") || cond.contains(".or("); - if (hasLogicalOps) return "!" + Utils.addParentheses(cond); - if (cond.startsWith("!") && !cond.startsWith("!(")) { - return cond.substring(1); - } - return "!" + cond; - } - - private static List append(List conds, String newCond) { - if (newCond == null || newCond.isEmpty()) return conds; - List copy = new ArrayList<>(conds); - copy.add(newCond); - return copy; - } - - private static String joinConditions(List conds) { - if (conds.isEmpty()) return ""; - return String.join(" && ", conds); - } - - private static void mergeCondition(Map result, String varName, String newCond) { - result.merge(varName, newCond, (existing, incoming) -> - Utils.addParentheses(existing) + " || " + Utils.addParentheses(incoming)); - } -} \ No newline at end of file diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/Utils.java b/src/main/java/org/wpilib/gradlerio/graphgen/Utils.java index 18594e12..aaada2f5 100644 --- a/src/main/java/org/wpilib/gradlerio/graphgen/Utils.java +++ b/src/main/java/org/wpilib/gradlerio/graphgen/Utils.java @@ -1,10 +1,67 @@ package org.wpilib.gradlerio.graphgen; +import java.util.List; +import java.util.stream.Collectors; + class Utils { - public static String addParentheses(String condition) { + static String addParentheses(String condition) { if ((condition.startsWith("!(") || condition.startsWith("(")) && condition.endsWith(")")) { return condition; } return "(" + condition + ")"; } + + static String joinWithOr(List conditions) { + if (conditions.contains("true")) { + return "true"; + } + return conditions.stream() + .filter(cond -> !cond.equals("false") && !cond.isEmpty()) + .distinct() + .map(condition -> hasAndStmt(condition) ? addParentheses(condition) : condition) + .collect(Collectors.joining(" || ")); + } + + static String joinWithAnd(List conditions) { + if (conditions.contains("false")) { + return "false"; + } + return conditions.stream() + .filter(cond -> !cond.equals("true") && !cond.isEmpty()) + .distinct() + .map(condition -> hasOrStmt(condition) ? addParentheses(condition) : condition) + .collect(Collectors.joining(" && ")); + } + + static String joinWithAnd(String... conditions) { + return joinWithAnd(List.of(conditions)); + } + + static String negate(String condition) { + condition = condition.trim(); + + if (condition.startsWith("!(") && condition.endsWith(")")) return condition.substring(2, condition.length() - 1); + if (condition.endsWith(".negate()")) return condition.substring(0, condition.length() - 9); + if (hasAndStmt(condition) || hasOrStmt(condition)) return "!" + addParentheses(condition); + + if (condition.contains(" == ")) return condition.replace(" == ", " != "); + if (condition.contains(" != ")) return condition.replace(" != ", " == "); + if (condition.contains(" >= ")) return condition.replace(" >= ", " < "); + if (condition.contains(" <= ")) return condition.replace(" <= ", " > "); + if (condition.contains(" > ")) return condition.replace(" > ", " <= "); + if (condition.contains(" < ")) return condition.replace(" < ", " >= "); + if (condition.startsWith("!") && !condition.startsWith("!(")) { + return condition.substring(1); + } + + return "!" + condition; + } + + private static boolean hasAndStmt(String condition) { + return condition.contains("&&") || condition.contains(".and("); + } + + private static boolean hasOrStmt(String condition) { + return condition.contains("||") || condition.contains(".or("); + } } diff --git a/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java b/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java index 7e8b3520..35191bed 100644 --- a/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java +++ b/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java @@ -13,7 +13,6 @@ import static org.junit.jupiter.api.Assertions.*; class MakeStateMachineGraphsTaskTest { - @TempDir Path tempDir; @@ -83,13 +82,13 @@ private Command cmd(String name) { stateDiagram-v2 direction LR - state1 --> state2 : true + state1 --> state2 : instant state2 --> Exit_State_Machine : when complete classDef initialState color: #00FF00 class state1 initialState """; - assertEquals(mermaidContent.strip(), expected.strip()); + assertEquals(expected.strip(), mermaidContent.strip()); } @Test @@ -132,12 +131,12 @@ private Command cmd(String name) { stateDiagram-v2 direction LR - s1 --> s2 : when complete and someCondition() + s1 --> s2 : when complete && someCondition() s1 --> s3 : emergency() s2 --> s3 : emergency() s3 --> Exit_State_Machine : done() """; - assertEquals(mermaidContent.strip(), expected.strip()); + assertEquals(expected.strip(), mermaidContent.strip()); } @Test @@ -181,14 +180,63 @@ private Command cmd(String name) { stateDiagram-v2 direction LR - start --> left : check() and cond - start --> right : check() and !cond + start --> left : check() && cond + start --> right : check() && !cond + """; + assertEquals(expected.strip(), mermaidContent.strip()); + } + + @Test + void testNestedLogicStatements() throws IOException { + String content = """ + package frc.robot; + import org.wpilib.gradlerio.graphgen.MakeStateMachineGraph; + + public class Robot { + @MakeStateMachineGraph + public StateMachine doubleSM() { + StateMachine sm1 = new StateMachine("Early Return Auto"); + State a = sm1.addState(cmd("a")); + State b = sm1.addState(cmd("b")); + State c = sm1.addState(cmd("c")); + + a.switchTo(() -> { + switch (number) { + case 0, 1: + if (condition1) { + if (condition2) { + return c; + } + } + return b; + } + return a; + }).whenComplete(); + return sm1; + } + } + """; + Files.writeString(javaRoot.resolve("Robot.java"), content); + + taskProvider.get().run(); + + String mermaidContent = Files.readString(outputDir.resolve("Early Return Auto.mermaid")); + String expected = """ + --- + state_definition_order: [a, b, c] + --- + stateDiagram-v2 + direction LR + + a --> a : when complete && !(number == 0 || number == 1) + a --> b : when complete && (((number == 0 || number == 1) && condition1) && !condition2) || ((number == 0 || number == 1) && !condition1) + a --> c : when complete && (((number == 0 || number == 1) && condition1) && condition2) """; - assertEquals(mermaidContent.strip(), expected.strip()); + assertEquals(expected.strip(), mermaidContent.strip()); } @Test - void testEarlyReturnInLambdaTransition() throws IOException { + void sdf() throws IOException { String content = """ package frc.robot; import org.wpilib.gradlerio.graphgen.MakeStateMachineGraph; @@ -206,7 +254,11 @@ public StateMachine doubleSM() { case 0, 1: return b; } - if (conditionC) return c; + if (conditionC) { + return c; + } else { + System.out.println("No Return Here!"); + } return a; }).whenComplete(); return sm1; @@ -217,7 +269,7 @@ public StateMachine doubleSM() { taskProvider.get().run(); - String earlyReturnContent = Files.readString(outputDir.resolve("Early Return Auto.mermaid")); + String mermaidContent = Files.readString(outputDir.resolve("Early Return Auto.mermaid")); String expected = """ --- state_definition_order: [a, b, c] @@ -225,11 +277,11 @@ public StateMachine doubleSM() { stateDiagram-v2 direction LR - a --> b : when complete and (number == 0 || number == 1) - a --> c : when complete and number != 0 && number != 1 && conditionC - a --> a : when complete and number != 0 && number != 1 && !conditionC + a --> a : when complete && (!(number == 0 || number == 1) && !conditionC) + a --> b : when complete && (number == 0 || number == 1) + a --> c : when complete && (!(number == 0 || number == 1) && conditionC) """; - assertEquals(earlyReturnContent.strip(), expected.strip()); + assertEquals(expected.strip(), mermaidContent.strip()); } @Test @@ -270,9 +322,9 @@ private Command cmd(String name) { stateDiagram-v2 direction LR - s1 --> s2 : true + s1 --> s2 : instant """; - assertEquals(mermaidContent.strip(), expected.strip()); + assertEquals(expected.strip(), mermaidContent.strip()); } @Test From 1a102234c483064eee092d26c041968c5b08c514 Mon Sep 17 00:00:00 2001 From: Daniel1464 Date: Wed, 3 Jun 2026 18:04:40 -0400 Subject: [PATCH 09/14] Cleanup tests and support custom state machine types --- .../graphgen/MakeStateMachineGraphsTask.java | 38 +++-- .../MakeStateMachineGraphsTaskTest.java | 152 +++++++++++------- 2 files changed, 119 insertions(+), 71 deletions(-) diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java index da7bca4d..60c84ee1 100644 --- a/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java +++ b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java @@ -98,16 +98,36 @@ private void extractFromFile( var cu = parser.parse(content).getResult().orElseThrow(); cu.findAll(MethodDeclaration.class).forEach(method -> { var annotationOpt = method.getAnnotationByName("MakeStateMachineGraph"); - if (!method.getTypeAsString().equals("StateMachine") || annotationOpt.isEmpty()) return; + if (annotationOpt.isEmpty()) return; + var extractedType = "StateMachine"; + if (annotationOpt.get().isNormalAnnotationExpr()) { + var data = annotationOpt.get() + .asNormalAnnotationExpr() + .getPairs() + .stream() + .filter(pair -> pair.getNameAsString().equals("stateMachineType")) + .findFirst(); + if (data.isPresent()) { + var valueExpr = data.get().getValue(); + if (valueExpr.isStringLiteralExpr()) { + extractedType = valueExpr.asStringLiteralExpr().getValue(); + } + } + } + final var stateMachineType = extractedType; + + if (!method.getTypeAsString().equals(stateMachineType)) { + throw new RuntimeException("Method " + method.getNameAsString() + " is annotated with @MakeStateMachineGraph, but it does not return " + stateMachineType); + } var variableDefs = method.findAll(VariableDeclarationExpr.class); - var stateSupplierTypes = List.of("Supplier", "Supplier"); + var stateSupplierTypes = List.of("Supplier<" + stateMachineType + ".State>", "Supplier"); boolean hasStateSupplierVar = variableDefs.stream() .filter(dec -> dec.getVariables().size() == 1) .anyMatch(dec -> stateSupplierTypes.contains(dec.getVariable(0).getTypeAsString())); if (hasStateSupplierVar) { throw new RuntimeException( - "Methods annotated with @MakeStateMachineGraph cannot contain variables of type Supplier." + "Methods annotated with @MakeStateMachineGraph cannot contain variables of type Supplier<" + stateMachineType + ".State>." ); } var stateMachineDefs = variableDefs @@ -118,15 +138,15 @@ private void extractFromFile( var def = v.getVariable(0).getInitializer(); return def.isPresent() && def.get().isObjectCreationExpr() && - def.get().asObjectCreationExpr().getType().asString().endsWith("StateMachine"); + def.get().asObjectCreationExpr().getType().asString().endsWith(stateMachineType); }) .toList(); if (stateMachineDefs.isEmpty()) { - throw new RuntimeException("No StateMachine declaration found in " + method.getNameAsString()); + throw new RuntimeException("No " + stateMachineType + " declaration found in " + method.getNameAsString()); } else if (stateMachineDefs.size() > 1) { throw new RuntimeException( - "Multiple StateMachine declarations found in " + method.getNameAsString() - + ". Currently, the @MakeStateMachineGraph annotation doesn't support multiple StateMachine declarations in the same method." + "Multiple " + stateMachineType + " declarations found in " + method.getNameAsString() + + ". Currently, the @MakeStateMachineGraph annotation doesn't support multiple state machine declarations in the same method." ); } @@ -136,12 +156,12 @@ private void extractFromFile( throw new RuntimeException( "The graph generator requires the name " + "of the state machine in '" + method.getNameAsString() + - "' to be a simple string literal (e.g. new StateMachine(\"My Auto\"))." + "' to be a simple string literal (e.g. new " + stateMachineType + "(\"My Auto\"))." ); } var graphName = rawGraphName.toString().substring(1, rawGraphName.toString().length() - 1); if (graphs.containsKey(graphName)) { - throw new RuntimeException("You already have a State Machine named '" + graphName + "'"); + throw new RuntimeException("2 state machines are named '" + graphName + "'"); } var graph = new StateMachineGraph(); graphs.put(graphName, graph); diff --git a/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java b/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java index 35191bed..ceff1691 100644 --- a/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java +++ b/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java @@ -49,8 +49,7 @@ void testBasicStateMachine() throws IOException { public class Robot { @MakeStateMachineGraph public StateMachine simpleAuto() { - System.out.println("Processing simpleAuto"); - StateMachine sm = new StateMachine("Simple Auto"); + var sm = new StateMachine("Simple Auto"); var state1 = sm.addState(cmd("State1")); var state2 = sm.addState(cmd("State2")); @@ -102,7 +101,7 @@ void testComplexTransitions() throws IOException { public class Robot { @MakeStateMachineGraph public StateMachine complexAuto() { - StateMachine sm = new StateMachine("Complex Auto"); + var sm = new StateMachine("Complex Auto"); var s1 = sm.addState(cmd("S1")); var s2 = sm.addState(cmd("S2")); var s3 = sm.addState(cmd("S3")); @@ -150,7 +149,7 @@ void testIfElseInLambdaTransition() throws IOException { public class Robot { @MakeStateMachineGraph public StateMachine lambdaAuto() { - StateMachine sm = new StateMachine("Lambda Auto"); + var sm = new StateMachine("Lambda Auto"); var start = sm.addState(cmd("Start")); var left = sm.addState(cmd("Left")); var right = sm.addState(cmd("Right")); @@ -195,10 +194,10 @@ void testNestedLogicStatements() throws IOException { public class Robot { @MakeStateMachineGraph public StateMachine doubleSM() { - StateMachine sm1 = new StateMachine("Early Return Auto"); - State a = sm1.addState(cmd("a")); - State b = sm1.addState(cmd("b")); - State c = sm1.addState(cmd("c")); + var sm1 = new StateMachine("Early Return Auto"); + var a = sm1.addState(cmd("a")); + var b = sm1.addState(cmd("b")); + var c = sm1.addState(cmd("c")); a.switchTo(() -> { switch (number) { @@ -235,55 +234,6 @@ public StateMachine doubleSM() { assertEquals(expected.strip(), mermaidContent.strip()); } - @Test - void sdf() throws IOException { - String content = """ - package frc.robot; - import org.wpilib.gradlerio.graphgen.MakeStateMachineGraph; - - public class Robot { - @MakeStateMachineGraph - public StateMachine doubleSM() { - StateMachine sm1 = new StateMachine("Early Return Auto"); - State a = sm1.addState(cmd("a")); - State b = sm1.addState(cmd("b")); - State c = sm1.addState(cmd("c")); - - a.switchTo(() -> { - switch (number) { - case 0, 1: - return b; - } - if (conditionC) { - return c; - } else { - System.out.println("No Return Here!"); - } - return a; - }).whenComplete(); - return sm1; - } - } - """; - Files.writeString(javaRoot.resolve("Robot.java"), content); - - taskProvider.get().run(); - - String mermaidContent = Files.readString(outputDir.resolve("Early Return Auto.mermaid")); - String expected = """ - --- - state_definition_order: [a, b, c] - --- - stateDiagram-v2 - direction LR - - a --> a : when complete && (!(number == 0 || number == 1) && !conditionC) - a --> b : when complete && (number == 0 || number == 1) - a --> c : when complete && (!(number == 0 || number == 1) && conditionC) - """; - assertEquals(expected.strip(), mermaidContent.strip()); - } - @Test void testUnnamedVariables() throws IOException { String content = """ @@ -295,7 +245,7 @@ void testUnnamedVariables() throws IOException { public class Robot { @MakeStateMachineGraph public StateMachine unnamedAuto() { - StateMachine sm = new StateMachine("Unnamed Auto"); + var sm = new StateMachine("Unnamed Auto"); var s1 = sm.addState(cmd("S1")); var s2 = sm.addState(cmd("S2")); @@ -336,8 +286,8 @@ void testMultipleStateMachinesInMethodThrows() throws IOException { public class Robot { @MakeStateMachineGraph public StateMachine doubleSM() { - StateMachine sm1 = new StateMachine("SM1"); - StateMachine sm2 = new StateMachine("SM2"); + var sm1 = new StateMachine("SM1"); + var sm2 = new StateMachine("SM2"); return sm1; } } @@ -356,8 +306,8 @@ void stateSupplierInStateMachineThrows() throws IOException { public class Robot { @MakeStateMachineGraph public StateMachine doubleSM() { - StateMachine sm1 = new StateMachine("SM1"); - State state1 = sm1.addState(cmd("state1")); + var sm1 = new StateMachine("SM1"); + var state1 = sm1.addState(cmd("state1")); Supplier derived = () -> state1; return sm1; } @@ -371,4 +321,82 @@ private Command cmd(String name) { assertThrows(RuntimeException.class, () -> taskProvider.get().run()); } + + @Test + void testCustomStateMachineType() throws IOException { + String content = """ + package frc.robot; + + import org.wpilib.command3.Command; + import org.wpilib.gradlerio.graphgen.MakeStateMachineGraph; + + public class Robot { + @MakeStateMachineGraph(stateMachineType = "CustomStateMachine") + public CustomStateMachine customAuto() { + var sm = new CustomStateMachine("Custom Auto"); + var state1 = sm.addState(cmd("State1")); + var state2 = sm.addState(cmd("State2")); + + sm.setInitialState(state1); + + state1.switchTo(state2).when(() -> true); + state2.exitStateMachine().whenComplete(); + + return sm; + } + + private Command cmd(String name) { + return Command.noRequirements(coro -> while(true) coro.yield()).named(name); + } + } + """; + Files.writeString(javaRoot.resolve("Robot.java"), content); + + taskProvider.get().run(); + + Path mermaidFile = outputDir.resolve("Custom Auto.mermaid"); + assertTrue(Files.exists(mermaidFile), "Mermaid file should be generated"); + + String mermaidContent = Files.readString(mermaidFile); + String expected = """ + --- + state_definition_order: [state1, state2] + --- + stateDiagram-v2 + direction LR + + state1 --> state2 : instant + state2 --> Exit_State_Machine : when complete + + classDef initialState color: #00FF00 + class state1 initialState + """; + assertEquals(expected.strip(), mermaidContent.strip()); + } + + @Test + void stateSupplierInCustomStateMachineThrows() throws IOException { + String content = """ + package frc.robot; + import org.wpilib.gradlerio.graphgen.MakeStateMachineGraph; + + public class Robot { + @MakeStateMachineGraph(stateMachineType = "CustomStateMachine") + public CustomStateMachine doubleSM() { + var sm1 = new CustomStateMachine("SM1"); + var state1 = sm1.addState(cmd("state1")); + Supplier derived = () -> state1; + return sm1; + } + + private Command cmd(String name) { + return Command.noRequirements(coro -> while(true) coro.yield()).named(name); + } + } + """; + Files.writeString(javaRoot.resolve("Robot.java"), content); + + var ex = assertThrows(RuntimeException.class, () -> taskProvider.get().run()); + assertTrue(ex.getMessage().contains("Supplier"), "Error message should mention Supplier"); + } } \ No newline at end of file From 9fb92cdfe5cea7001c459b0115bf64255692a462 Mon Sep 17 00:00:00 2001 From: Daniel1464 Date: Wed, 3 Jun 2026 20:28:41 -0400 Subject: [PATCH 10/14] Fix docs and inline switch statements --- .../gradlerio/graphgen/CodeBlockAnalyzer.java | 6 + .../graphgen/ExpressionAnalyzer.java | 16 +- .../graphgen/MakeStateMachineGraphsTask.java | 2 +- .../MakeStateMachineGraphsTaskTest.java | 158 ++++++++++++++++++ 4 files changed, 176 insertions(+), 6 deletions(-) diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/CodeBlockAnalyzer.java b/src/main/java/org/wpilib/gradlerio/graphgen/CodeBlockAnalyzer.java index bfb43ba8..f620baa7 100644 --- a/src/main/java/org/wpilib/gradlerio/graphgen/CodeBlockAnalyzer.java +++ b/src/main/java/org/wpilib/gradlerio/graphgen/CodeBlockAnalyzer.java @@ -8,6 +8,12 @@ import java.util.stream.Collectors; class CodeBlockAnalyzer { + /** + * Analyzes a code block for returned state machine states, + * and the conditions needed for the state to be returned. + * @param block The code block to analyze + * @return A map of state names to the conditions needed for that state to be returned + */ static Map analyze(BlockStmt block) { var rawReturns = new HashMap>(); analyzeStatement(block, "true", rawReturns); diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/ExpressionAnalyzer.java b/src/main/java/org/wpilib/gradlerio/graphgen/ExpressionAnalyzer.java index 0887a116..8af93f6d 100644 --- a/src/main/java/org/wpilib/gradlerio/graphgen/ExpressionAnalyzer.java +++ b/src/main/java/org/wpilib/gradlerio/graphgen/ExpressionAnalyzer.java @@ -10,6 +10,12 @@ import java.util.stream.Collectors; class ExpressionAnalyzer { + /** + * Analyzes an inline expression for returned state machine states, + * and the conditions needed for the state to be returned. + * @param expr The expression to analyze + * @return A map of state names to the conditions needed for that state to be returned + */ static Map analyze(Expression expr) { var rawReturns = new HashMap>(); ExpressionAnalyzer.analyze(expr, "true", rawReturns); @@ -20,6 +26,7 @@ static Map analyze(Expression expr) { )); } + // A mutable variant of analyze() used internally by CodeBlockAnalyzer. static void analyze( Expression expr, String pathCondition, @@ -62,13 +69,12 @@ private static void handleSwitch( throw new RuntimeException("Yield Statements within switch blocks are not supported for @MakeStateMachineGraph."); } var expression = stmt.asExpressionStmt().getExpression(); - if (expression.isConditionalExpr()) { - analyze(expression.asConditionalExpr(), Utils.joinWithAnd(pathCondition, matchCondition), returnsMap); - } else if (expression.isSwitchExpr()) { - handleSwitch(expression.asSwitchExpr(), Utils.joinWithAnd(pathCondition, matchCondition), returnsMap); + var condition = Utils.joinWithAnd(pathCondition, matchCondition); + if (expression.isConditionalExpr() || expression.isSwitchExpr()) { + analyze(expression, condition, returnsMap); } else { var varName = expression.toString().trim(); - returnsMap.computeIfAbsent(varName, k -> new ArrayList<>()).add(pathCondition); + returnsMap.computeIfAbsent(varName, k -> new ArrayList<>()).add(condition); } } } diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java index 60c84ee1..af59535a 100644 --- a/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java +++ b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java @@ -23,7 +23,7 @@ import java.util.regex.Pattern; /** - * A gradle task that generates state machine names + * A gradle task that generates diagrams of state machines written in Commands V3. */ @DisableCachingByDefault public abstract class MakeStateMachineGraphsTask extends DefaultTask { diff --git a/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java b/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java index ceff1691..5d2299fa 100644 --- a/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java +++ b/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java @@ -399,4 +399,162 @@ private Command cmd(String name) { var ex = assertThrows(RuntimeException.class, () -> taskProvider.get().run()); assertTrue(ex.getMessage().contains("Supplier"), "Error message should mention Supplier"); } + + @Test + void testInlineLogicStatementsInLambda() throws IOException { + String content = """ + package frc.robot; + + import org.wpilib.command3.Command; + import org.wpilib.gradlerio.graphgen.MakeStateMachineGraph; + + public class Robot { + @MakeStateMachineGraph + public StateMachine inlineConditional() { + var sm = new StateMachine("Inline Conditional"); + var start = sm.addState(cmd("Start")); + var stateA = sm.addState(cmd("StateA")); + var stateB = sm.addState(cmd("StateB")); + + start.switchTo(() -> condition ? stateA : stateB).when(() -> check()); + + return sm; + } + + @MakeStateMachineGraph + public StateMachine inlineSwitch() { + var sm = new StateMachine("Inline Switch"); + var start = sm.addState(cmd("Start")); + var state1 = sm.addState(cmd("State1")); + var state2 = sm.addState(cmd("State2")); + + start.switchTo(() -> switch (value) { + case 0 -> state1; + default -> state2; + }).whenComplete(); + + return sm; + } + + private Command cmd(String name) { + return Command.noRequirements(coro -> while(true) coro.yield()).named(name); + } + } + """; + Files.writeString(javaRoot.resolve("Robot.java"), content); + + taskProvider.get().run(); + + Path condMermaid = outputDir.resolve("Inline Conditional.mermaid"); + assertTrue(Files.exists(condMermaid), "Conditional mermaid file should be generated"); + String condContent = Files.readString(condMermaid); + String expectedCond = """ + --- + state_definition_order: [start, stateA, stateB] + --- + stateDiagram-v2 + direction LR + + start --> stateB : check() && !condition + start --> stateA : check() && condition + """; + assertEquals(expectedCond.strip(), condContent.strip()); + + Path switchMermaid = outputDir.resolve("Inline Switch.mermaid"); + assertTrue(Files.exists(switchMermaid), "Switch mermaid file should be generated"); + String switchContent = Files.readString(switchMermaid); + String expectedSwitch = """ + --- + state_definition_order: [start, state1, state2] + --- + stateDiagram-v2 + direction LR + + start --> state2 : when complete && value != 0 + start --> state1 : when complete && value == 0 + """; + assertEquals(expectedSwitch.strip(), switchContent.strip()); + } + + @Test + void testReturnInlineLogicStatementsInLambda() throws IOException { + String content = """ + package frc.robot; + + import org.wpilib.command3.Command; + import org.wpilib.gradlerio.graphgen.MakeStateMachineGraph; + + public class Robot { + @MakeStateMachineGraph + public StateMachine inlineConditional() { + var sm = new StateMachine("Inline Conditional"); + var start = sm.addState(cmd("Start")); + var stateA = sm.addState(cmd("StateA")); + var stateB = sm.addState(cmd("StateB")); + + start.switchTo(() -> { + return condition ? stateA : stateB; + }).when(() -> check()); + + return sm; + } + + @MakeStateMachineGraph + public StateMachine inlineSwitch() { + var sm = new StateMachine("Inline Switch"); + var start = sm.addState(cmd("Start")); + var state1 = sm.addState(cmd("State1")); + var state2 = sm.addState(cmd("State2")); + + start.switchTo(() -> { + if (!condition) return start; + return switch (value) { + case 0 -> state1; + default -> state2; + }; + }).whenComplete(); + + return sm; + } + + private Command cmd(String name) { + return Command.noRequirements(coro -> while(true) coro.yield()).named(name); + } + } + """; + Files.writeString(javaRoot.resolve("Robot.java"), content); + + taskProvider.get().run(); + + Path condMermaid = outputDir.resolve("Inline Conditional.mermaid"); + assertTrue(Files.exists(condMermaid), "Conditional mermaid file should be generated"); + String condContent = Files.readString(condMermaid); + String expectedCond = """ + --- + state_definition_order: [start, stateA, stateB] + --- + stateDiagram-v2 + direction LR + + start --> stateB : check() && !condition + start --> stateA : check() && condition + """; + assertEquals(expectedCond.strip(), condContent.strip()); + + Path switchMermaid = outputDir.resolve("Inline Switch.mermaid"); + assertTrue(Files.exists(switchMermaid), "Switch mermaid file should be generated"); + String switchContent = Files.readString(switchMermaid); + String expectedSwitch = """ + --- + state_definition_order: [start, state1, state2] + --- + stateDiagram-v2 + direction LR + + start --> start : when complete && !condition + start --> state2 : when complete && (condition && value != 0) + start --> state1 : when complete && (condition && value == 0) + """; + assertEquals(expectedSwitch.strip(), switchContent.strip()); + } } \ No newline at end of file From 2e3b5a06afdb718c596da80ef92ce9e81c83a537 Mon Sep 17 00:00:00 2001 From: Daniel1464 Date: Thu, 4 Jun 2026 16:08:11 -0400 Subject: [PATCH 11/14] Improve error messages --- .../gradlerio/graphgen/CodeBlockAnalyzer.java | 5 +- .../gradlerio/graphgen/ErrorLogger.java | 27 ++++++++++ .../graphgen/ExpressionAnalyzer.java | 2 +- .../graphgen/MakeStateMachineGraphsTask.java | 54 ++++++++++--------- .../MakeStateMachineGraphsTaskTest.java | 2 + 5 files changed, 63 insertions(+), 27 deletions(-) create mode 100644 src/main/java/org/wpilib/gradlerio/graphgen/ErrorLogger.java diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/CodeBlockAnalyzer.java b/src/main/java/org/wpilib/gradlerio/graphgen/CodeBlockAnalyzer.java index f620baa7..fcba44c5 100644 --- a/src/main/java/org/wpilib/gradlerio/graphgen/CodeBlockAnalyzer.java +++ b/src/main/java/org/wpilib/gradlerio/graphgen/CodeBlockAnalyzer.java @@ -45,8 +45,9 @@ private static List analyzeStatement( } fallThroughConditions.addAll(currentPaths); } else if (stmt.isReturnStmt()) { - var returnExpr = stmt.asReturnStmt().getExpression().orElseThrow(); - ExpressionAnalyzer.analyze(returnExpr, pathCondition, returnsMap); + var returnExpr = stmt.asReturnStmt().getExpression(); + if (returnExpr.isEmpty()) return List.of(); + ExpressionAnalyzer.analyze(returnExpr.get(), pathCondition, returnsMap); } else if (stmt.isIfStmt()) { var ifStmt = stmt.asIfStmt(); var conditionStr = ifStmt.getCondition().toString(); diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/ErrorLogger.java b/src/main/java/org/wpilib/gradlerio/graphgen/ErrorLogger.java new file mode 100644 index 00000000..a1ad106d --- /dev/null +++ b/src/main/java/org/wpilib/gradlerio/graphgen/ErrorLogger.java @@ -0,0 +1,27 @@ +package org.wpilib.gradlerio.graphgen; + +import com.github.javaparser.ast.Node; + +public class ErrorLogger { + private static String currentFilePath = "unknown file"; + + static void setFilePath(String path) { + currentFilePath = path; + } + + /** + * A utility method that logs the line number and the file path + * that an error originates from, in addition to the error message. + * This serves to supplement the lack of stacktrace support within JavaParser. + * @param message The error message to throw + * @param node The location (statement or expression) that the error originates from + */ + static void throwError(String message, Node node) { + var begin = node.getBegin(); + if (begin.isPresent()) { + throw new RuntimeException(currentFilePath + ": Line " + begin.get().line + ": " + message); + } else { + throw new RuntimeException(currentFilePath + ": Line " + message); + } + } +} diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/ExpressionAnalyzer.java b/src/main/java/org/wpilib/gradlerio/graphgen/ExpressionAnalyzer.java index 8af93f6d..ac090bfc 100644 --- a/src/main/java/org/wpilib/gradlerio/graphgen/ExpressionAnalyzer.java +++ b/src/main/java/org/wpilib/gradlerio/graphgen/ExpressionAnalyzer.java @@ -66,7 +66,7 @@ private static void handleSwitch( var stmt = entry.getStatements().get(0); if (stmt.isBlockStmt()) { - throw new RuntimeException("Yield Statements within switch blocks are not supported for @MakeStateMachineGraph."); + ErrorLogger.throwError("Yield Statements within switch blocks are not supported for @MakeStateMachineGraph.", stmt); } var expression = stmt.asExpressionStmt().getExpression(); var condition = Utils.joinWithAnd(pathCondition, matchCondition); diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java index af59535a..e7e359e6 100644 --- a/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java +++ b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java @@ -64,6 +64,7 @@ private void extractFromDirectory(String javaRoot, String deployDir) throws IOEx .filter(path -> path.toString().endsWith(".java")) .forEach(path -> { try { + ErrorLogger.setFilePath(path.toString().replace(javaRoot, "")); extractFromFile(parser, path.toFile(), graphs); } catch (IOException e) { System.err.println("Failed to parse: " + path + " — " + e.getMessage()); @@ -95,8 +96,9 @@ private void extractFromFile( .matcher(content) .replaceAll(res -> "unused_var_" + UUID.randomUUID().toString().substring(0, 8)); - var cu = parser.parse(content).getResult().orElseThrow(); - cu.findAll(MethodDeclaration.class).forEach(method -> { + var cu = parser.parse(content).getResult(); + if (cu.isEmpty()) return; + cu.get().findAll(MethodDeclaration.class).forEach(method -> { var annotationOpt = method.getAnnotationByName("MakeStateMachineGraph"); if (annotationOpt.isEmpty()) return; var extractedType = "StateMachine"; @@ -117,19 +119,19 @@ private void extractFromFile( final var stateMachineType = extractedType; if (!method.getTypeAsString().equals(stateMachineType)) { - throw new RuntimeException("Method " + method.getNameAsString() + " is annotated with @MakeStateMachineGraph, but it does not return " + stateMachineType); + ErrorLogger.throwError("Method " + method.getNameAsString() + " is annotated with @MakeStateMachineGraph, but it does not return " + stateMachineType, method); } var variableDefs = method.findAll(VariableDeclarationExpr.class); var stateSupplierTypes = List.of("Supplier<" + stateMachineType + ".State>", "Supplier"); - boolean hasStateSupplierVar = variableDefs.stream() + variableDefs.stream() .filter(dec -> dec.getVariables().size() == 1) - .anyMatch(dec -> stateSupplierTypes.contains(dec.getVariable(0).getTypeAsString())); - if (hasStateSupplierVar) { - throw new RuntimeException( - "Methods annotated with @MakeStateMachineGraph cannot contain variables of type Supplier<" + stateMachineType + ".State>." - ); - } + .filter(dec -> stateSupplierTypes.contains(dec.getVariable(0).getTypeAsString())) + .findFirst() + .ifPresent(illegalSupplierVar -> ErrorLogger.throwError( + "Methods annotated with @MakeStateMachineGraph cannot contain variables of type Supplier.", + illegalSupplierVar + )); var stateMachineDefs = variableDefs .stream() .filter(v -> { @@ -142,30 +144,31 @@ private void extractFromFile( }) .toList(); if (stateMachineDefs.isEmpty()) { - throw new RuntimeException("No " + stateMachineType + " declaration found in " + method.getNameAsString()); + ErrorLogger.throwError("No " + stateMachineType + " declaration found in " + method.getNameAsString(), method); } else if (stateMachineDefs.size() > 1) { - throw new RuntimeException( + ErrorLogger.throwError( "Multiple " + stateMachineType + " declarations found in " + method.getNameAsString() - + ". Currently, the @MakeStateMachineGraph annotation doesn't support multiple state machine declarations in the same method." + + ". Currently, the @MakeStateMachineGraph annotation doesn't support multiple state machine declarations in the same method.", + stateMachineDefs.get(1) ); } var stateMachineDef = stateMachineDefs.getFirst().getVariable(0).getInitializer().orElseThrow(); var rawGraphName = stateMachineDef.asObjectCreationExpr().getArgument(0); if (!rawGraphName.isStringLiteralExpr()) { - throw new RuntimeException( + ErrorLogger.throwError( "The graph generator requires the name " + "of the state machine in '" + method.getNameAsString() + - "' to be a simple string literal (e.g. new " + stateMachineType + "(\"My Auto\"))." + "' to be a simple string literal (e.g. new " + stateMachineType + "(\"My Auto\")).", + stateMachineDef ); } var graphName = rawGraphName.toString().substring(1, rawGraphName.toString().length() - 1); if (graphs.containsKey(graphName)) { - throw new RuntimeException("2 state machines are named '" + graphName + "'"); + ErrorLogger.throwError("There is already a state machine named '" + graphName + "'", stateMachineDef); } var graph = new StateMachineGraph(); graphs.put(graphName, graph); - variableDefs.stream() .filter(v -> { var initializer = v.getVariable(0).getInitializer(); @@ -177,7 +180,7 @@ private void extractFromFile( method.findAll(MethodCallExpr.class).forEach(call -> { if (call.getNameAsString().equals("setInitialState")) { - graph.initialState = call.getArguments().getFirst().orElseThrow().toString(); + call.getArguments().getFirst().ifPresent(state -> graph.initialState = state.toString()); return; } @@ -196,7 +199,8 @@ private void extractFromFile( if (isWhenComplete) { transitionCond = "when complete"; } else { - var expr = transitionCondExpr.orElseThrow().toString(); + if (transitionCondExpr.isEmpty()) return; + var expr = transitionCondExpr.get().toString(); expr = expr.replace("() -> ", ""); expr = expr.replace(".getAsBoolean()", ""); transitionCond = isWhen ? expr : Utils.joinWithAnd("when complete", expr); @@ -206,8 +210,10 @@ private void extractFromFile( var toStateExpr = toStateCall.getArguments().getFirst(); var toState = toStateExpr.map(Expression::toString).orElse("Exit_State_Machine"); boolean toStateIsLambda = toStateExpr.isPresent() && toStateExpr.get().isLambdaExpr(); + + if (toStateCall.getScope().isEmpty()) return; if (List.of("switchTo", "exitStateMachine").contains(toStateCall.getNameAsString())) { - var fromState = toStateCall.getScope().orElseThrow().toString(); + var fromState = toStateCall.getScope().get().toString(); if (toStateIsLambda) { graph.transitions.addAll( transitionsFromLambdaExpr(toStateExpr.get(), fromState, transitionCond) @@ -218,10 +224,10 @@ private void extractFromFile( } else if (List.of("to", "toExitStateMachine").contains(toStateCall.getNameAsString())) { // If it's just a regular "to", there must be a switchFromAny before it. var argList = toStateCall.getScope() - .filter(s -> s instanceof MethodCallExpr) - .map(s -> ((MethodCallExpr) s)) - .filter(s -> s.getNameAsString().equals("switchFromAny")) - .map(MethodCallExpr::getArguments); + .filter(s -> s instanceof MethodCallExpr) + .map(s -> ((MethodCallExpr) s)) + .filter(s -> s.getNameAsString().equals("switchFromAny")) + .map(MethodCallExpr::getArguments); if (argList.isEmpty()) return; for (var fromState: argList.get()) { if (toStateIsLambda) { diff --git a/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java b/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java index 5d2299fa..68fa0f0c 100644 --- a/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java +++ b/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java @@ -398,6 +398,8 @@ private Command cmd(String name) { var ex = assertThrows(RuntimeException.class, () -> taskProvider.get().run()); assertTrue(ex.getMessage().contains("Supplier"), "Error message should mention Supplier"); + assertTrue(ex.getMessage().contains("Robot.java"), "Error message should contain file name"); + assertTrue(ex.getMessage().contains("Line 9:"), "Error message should contain line number 9"); } @Test From a7c875721e3b59219a31d49ffb490c8f75684fed Mon Sep 17 00:00:00 2001 From: Daniel Chen <108989218+Daniel1464@users.noreply.github.com> Date: Sat, 6 Jun 2026 01:21:10 -0400 Subject: [PATCH 12/14] Remove redundant toAbsolutePath() --- .../gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java b/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java index 68fa0f0c..7c1f6155 100644 --- a/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java +++ b/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java @@ -34,7 +34,7 @@ void setup() throws IOException { // Use the same path format as default to avoid any normalization issues in the task taskProvider.configure(task -> { task.getJavaRoot().set(javaRoot.toString()); - task.getDeployDirectory().set(deployDir.toAbsolutePath().toString()); + task.getDeployDirectory().set(deployDir.toString()); }); } @@ -559,4 +559,4 @@ private Command cmd(String name) { """; assertEquals(expectedSwitch.strip(), switchContent.strip()); } -} \ No newline at end of file +} From 753412c8b7089a23e0edc231bce2054502aef04a Mon Sep 17 00:00:00 2001 From: Daniel Chen <108989218+Daniel1464@users.noreply.github.com> Date: Sat, 6 Jun 2026 01:22:58 -0400 Subject: [PATCH 13/14] Fix small mistake --- .../gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java b/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java index 7c1f6155..6a652e6e 100644 --- a/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java +++ b/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java @@ -154,7 +154,7 @@ public StateMachine lambdaAuto() { var left = sm.addState(cmd("Left")); var right = sm.addState(cmd("Right")); - start.switchTo(cond -> { + start.switchTo(() -> { if (cond) return left; else return right; }).when(() -> check()); From ab4265a943f57455c706063d1841bb51c42a4661 Mon Sep 17 00:00:00 2001 From: Daniel1464 Date: Sun, 7 Jun 2026 23:06:16 -0400 Subject: [PATCH 14/14] output JSON metadata files as well --- .../graphgen/MakeStateMachineGraphsTask.java | 46 +++++++++++-------- .../MakeStateMachineGraphsTaskTest.java | 4 +- 2 files changed, 29 insertions(+), 21 deletions(-) diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java index e7e359e6..a2025343 100644 --- a/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java +++ b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java @@ -1,4 +1,3 @@ - package org.wpilib.gradlerio.graphgen; import com.github.javaparser.JavaParser; @@ -7,6 +6,7 @@ import com.github.javaparser.ast.expr.Expression; import com.github.javaparser.ast.expr.MethodCallExpr; import com.github.javaparser.ast.expr.VariableDeclarationExpr; +import com.google.gson.Gson; import org.gradle.api.DefaultTask; import org.gradle.api.provider.Property; import org.gradle.api.tasks.Input; @@ -15,7 +15,6 @@ import org.gradle.work.DisableCachingByDefault; import java.io.File; -import java.io.FileWriter; import java.io.IOException; import java.nio.file.Files; import java.nio.file.Paths; @@ -27,14 +26,15 @@ */ @DisableCachingByDefault public abstract class MakeStateMachineGraphsTask extends DefaultTask { - private record Transition(String fromState, String toState, String transitionCond) {} + private record Transition(String origin, String target, String condition) {} private static class StateMachineGraph { - List transitions = new ArrayList<>(); - String initialState; - List stateDefinitionOrder = new ArrayList<>(); + final List transitions = new ArrayList<>(); + final List stateDefinitionOrder = new ArrayList<>(); + String initialState = ""; } private final Pattern unusedVarPattern = Pattern.compile("(?<=[(,\\s])_(?=[\\s,)->])"); + private final Gson gson = new Gson(); @Input @Optional @@ -44,15 +44,20 @@ private static class StateMachineGraph { @Optional public abstract Property getDeployDirectory(); + @Input + @Optional + public abstract Property getDiagramGenerationDirectory(); + @TaskAction public void run() throws IOException { extractFromDirectory( getJavaRoot().getOrElse("src/main/java/"), - getDeployDirectory().getOrElse("src/main/deploy/") + getDeployDirectory().getOrElse("src/main/deploy/"), + getDiagramGenerationDirectory().getOrElse("stateMachineGraphs/") ); } - private void extractFromDirectory(String javaRoot, String deployDir) throws IOException { + private void extractFromDirectory(String javaRoot, String deployDir, String diagramGenDir) throws IOException { var graphs = new LinkedHashMap(); // Init configuration for javaparser var config = new ParserConfiguration(); @@ -71,16 +76,19 @@ private void extractFromDirectory(String javaRoot, String deployDir) throws IOEx } }); + new File(deployDir, "/stateMachineGraphData").mkdirs(); + new File(diagramGenDir).mkdirs(); for (var entry: graphs.entrySet()) { - var methodName = entry.getKey(); + var name = entry.getKey(); var graph = entry.getValue(); - var file = new File(deployDir, "/stateMachineGraphs/" + methodName + ".mermaid"); - if (file.getParentFile() != null) { - file.getParentFile().mkdirs(); - } - try (var writer = new FileWriter(file)) { - writer.write(generateMermaid(graph)); - } + Files.writeString( + new File(diagramGenDir, name + ".mermaid").toPath(), + generateMermaid(graph) + ); + Files.writeString( + new File(deployDir, "/stateMachineGraphData/" + name + ".json").toPath(), + gson.toJson(graph) + ); } } @@ -272,11 +280,11 @@ private String generateMermaid(StateMachineGraph graph) { sb.append(" direction LR\n\n"); for (var t : graph.transitions) { sb.append(" ") - .append(sanitize(t.fromState())) + .append(sanitize(t.origin())) .append(" --> ") - .append(sanitize(t.toState())) + .append(sanitize(t.target())) .append(" : ") - .append(sanitizeTransitionCond(t.transitionCond())) + .append(sanitizeTransitionCond(t.condition())) .append("\n"); } if (graph.initialState != null) { diff --git a/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java b/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java index 6a652e6e..155f6c5c 100644 --- a/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java +++ b/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java @@ -27,14 +27,14 @@ void setup() throws IOException { javaRoot = tempDir.resolve("src/main/java").toAbsolutePath(); Files.createDirectories(javaRoot); - + outputDir = tempDir.resolve("stateMachineGraphs").toAbsolutePath(); var deployDir = tempDir.resolve("src/main/deploy").toAbsolutePath(); - outputDir = deployDir.resolve("stateMachineGraphs"); // Use the same path format as default to avoid any normalization issues in the task taskProvider.configure(task -> { task.getJavaRoot().set(javaRoot.toString()); task.getDeployDirectory().set(deployDir.toString()); + task.getDiagramGenerationDirectory().set(outputDir.toString()); }); }