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..fcba44c5 --- /dev/null +++ b/src/main/java/org/wpilib/gradlerio/graphgen/CodeBlockAnalyzer.java @@ -0,0 +1,133 @@ +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 { + /** + * 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); + 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(); + 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(); + 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/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 new file mode 100644 index 00000000..ac090bfc --- /dev/null +++ b/src/main/java/org/wpilib/gradlerio/graphgen/ExpressionAnalyzer.java @@ -0,0 +1,81 @@ +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 { + /** + * 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); + return rawReturns.entrySet().stream() + .collect(Collectors.toMap( + Map.Entry::getKey, + entry -> Utils.joinWithOr(entry.getValue()) + )); + } + + // A mutable variant of analyze() used internally by CodeBlockAnalyzer. + 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()) { + ErrorLogger.throwError("Yield Statements within switch blocks are not supported for @MakeStateMachineGraph.", stmt); + } + var expression = stmt.asExpressionStmt().getExpression(); + 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(condition); + } + } + } +} diff --git a/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java new file mode 100644 index 00000000..a2025343 --- /dev/null +++ b/src/main/java/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTask.java @@ -0,0 +1,312 @@ +package org.wpilib.gradlerio.graphgen; + +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 com.google.gson.Gson; +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.IOException; +import java.nio.file.Files; +import java.nio.file.Paths; +import java.util.*; +import java.util.regex.Pattern; + +/** + * A gradle task that generates diagrams of state machines written in Commands V3. + */ +@DisableCachingByDefault +public abstract class MakeStateMachineGraphsTask extends DefaultTask { + private record Transition(String origin, String target, String condition) {} + private static class StateMachineGraph { + 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 + public abstract Property getJavaRoot(); + + @Input + @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/"), + getDiagramGenerationDirectory().getOrElse("stateMachineGraphs/") + ); + } + + private void extractFromDirectory(String javaRoot, String deployDir, String diagramGenDir) 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(javaRoot)) + .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()); + } + }); + + new File(deployDir, "/stateMachineGraphData").mkdirs(); + new File(diagramGenDir).mkdirs(); + for (var entry: graphs.entrySet()) { + var name = entry.getKey(); + var graph = entry.getValue(); + Files.writeString( + new File(diagramGenDir, name + ".mermaid").toPath(), + generateMermaid(graph) + ); + Files.writeString( + new File(deployDir, "/stateMachineGraphData/" + name + ".json").toPath(), + gson.toJson(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. + content = unusedVarPattern + .matcher(content) + .replaceAll(res -> "unused_var_" + UUID.randomUUID().toString().substring(0, 8)); + + 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"; + 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)) { + 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"); + variableDefs.stream() + .filter(dec -> dec.getVariables().size() == 1) + .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 -> { + // 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().endsWith(stateMachineType); + }) + .toList(); + if (stateMachineDefs.isEmpty()) { + ErrorLogger.throwError("No " + stateMachineType + " declaration found in " + method.getNameAsString(), method); + } else if (stateMachineDefs.size() > 1) { + ErrorLogger.throwError( + "Multiple " + stateMachineType + " declarations found in " + method.getNameAsString() + + ". 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()) { + 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\")).", + stateMachineDef + ); + } + var graphName = rawGraphName.toString().substring(1, rawGraphName.toString().length() - 1); + if (graphs.containsKey(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(); + return initializer.isPresent() && + initializer.get().isMethodCallExpr() && + initializer.get().asMethodCallExpr().getNameAsString().equals("addState"); + }) + .forEachOrdered(v -> graph.stateDefinitionOrder.add(v.getVariable(0).getNameAsString())); + + method.findAll(MethodCallExpr.class).forEach(call -> { + if (call.getNameAsString().equals("setInitialState")) { + call.getArguments().getFirst().ifPresent(state -> graph.initialState = state.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 { + if (transitionCondExpr.isEmpty()) return; + var expr = transitionCondExpr.get().toString(); + expr = expr.replace("() -> ", ""); + expr = expr.replace(".getAsBoolean()", ""); + transitionCond = isWhen ? expr : Utils.joinWithAnd("when complete", expr); + } + + var toStateCall = toStateCallOpt.get(); + 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().get().toString(); + if (toStateIsLambda) { + graph.transitions.addAll( + transitionsFromLambdaExpr(toStateExpr.get(), fromState, transitionCond) + ); + } else { + graph.transitions.add(new Transition(fromState, toState, transitionCond)); + } + } 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); + if (argList.isEmpty()) return; + for (var fromState: argList.get()) { + if (toStateIsLambda) { + graph.transitions.addAll( + transitionsFromLambdaExpr(toStateExpr.get(), fromState.toString(), transitionCond) + ); + } else { + graph.transitions.add(new Transition(fromState.toString(), toState, transitionCond)); + } + } + } + }); + }); + } + + private List transitionsFromLambdaExpr( + Expression toState, + String fromState, + String transitionCond + ) { + var transitions = new ArrayList(); + 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(); + var fullCond = Utils.joinWithAnd(transitionCond, additionalCond); + transitions.add(new Transition(fromState, innerToState, fullCond)); + } + return transitions; + } + + private String generateMermaid(StateMachineGraph graph) { + StringBuilder sb = new StringBuilder(); + sb.append("---\n"); + sb.append("state_definition_order: "); + sb.append(graph.stateDefinitionOrder); + sb.append("\n---\n"); + sb.append("stateDiagram-v2\n"); + sb.append(" direction LR\n\n"); + for (var t : graph.transitions) { + sb.append(" ") + .append(sanitize(t.origin())) + .append(" --> ") + .append(sanitize(t.target())) + .append(" : ") + .append(sanitizeTransitionCond(t.condition())) + .append("\n"); + } + if (graph.initialState != null) { + sb.append("\n classDef initialState color: #00FF00"); + sb.append("\n class "); + sb.append(graph.initialState); + sb.append(" initialState\n"); + } + 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/Utils.java b/src/main/java/org/wpilib/gradlerio/graphgen/Utils.java new file mode 100644 index 00000000..aaada2f5 --- /dev/null +++ b/src/main/java/org/wpilib/gradlerio/graphgen/Utils.java @@ -0,0 +1,67 @@ +package org.wpilib.gradlerio.graphgen; + +import java.util.List; +import java.util.stream.Collectors; + +class Utils { + 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/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."); + }); } } 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..155f6c5c --- /dev/null +++ b/src/test/groovy/org/wpilib/gradlerio/graphgen/MakeStateMachineGraphsTaskTest.java @@ -0,0 +1,562 @@ +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); + outputDir = tempDir.resolve("stateMachineGraphs").toAbsolutePath(); + var deployDir = tempDir.resolve("src/main/deploy").toAbsolutePath(); + + // 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()); + }); + } + + @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() { + var 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); + 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 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() { + var 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")); + String expected = """ + --- + state_definition_order: [s1, s2, s3] + --- + stateDiagram-v2 + direction LR + + s1 --> s2 : when complete && someCondition() + s1 --> s3 : emergency() + s2 --> s3 : emergency() + s3 --> Exit_State_Machine : done() + """; + assertEquals(expected.strip(), mermaidContent.strip()); + } + + @Test + void testIfElseInLambdaTransition() throws IOException { + String content = """ + package frc.robot; + + import org.wpilib.command3.Command; + import org.wpilib.gradlerio.graphgen.MakeStateMachineGraph; + + public class Robot { + @MakeStateMachineGraph + public StateMachine lambdaAuto() { + var 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(() -> { + 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")); + String expected = """ + --- + state_definition_order: [start, left, right] + --- + stateDiagram-v2 + direction LR + + 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() { + 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) { + 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(expected.strip(), mermaidContent.strip()); + } + + @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() { + var 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")); + String expected = """ + --- + state_definition_order: [s1, s2] + --- + stateDiagram-v2 + direction LR + + s1 --> s2 : instant + """; + assertEquals(expected.strip(), mermaidContent.strip()); + } + + @Test + void testMultipleStateMachinesInMethodThrows() throws IOException { + String content = """ + package frc.robot; + import org.wpilib.gradlerio.graphgen.MakeStateMachineGraph; + + public class Robot { + @MakeStateMachineGraph + public StateMachine doubleSM() { + var sm1 = new StateMachine("SM1"); + var sm2 = new StateMachine("SM2"); + return sm1; + } + } + """; + Files.writeString(javaRoot.resolve("Robot.java"), content); + + 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() { + var sm1 = new StateMachine("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); + + 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"); + 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 + 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()); + } +}