From 26eeb500579e1bc64b9cc59bb46399b8518b7986 Mon Sep 17 00:00:00 2001
From: Jonah Snider <jonah@jonahsnider.com>
Date: Sun, 4 Aug 2024 23:19:04 -0700
Subject: [PATCH] Add StateMachine class

---
 .../util/scheduling/LifecycleSubsystem.java   |  6 +-
 .../util/state_machines/StateMachine.java     | 92 +++++++++++++++++++
 2 files changed, 96 insertions(+), 2 deletions(-)
 create mode 100644 src/main/java/frc/robot/util/state_machines/StateMachine.java

diff --git a/src/main/java/frc/robot/util/scheduling/LifecycleSubsystem.java b/src/main/java/frc/robot/util/scheduling/LifecycleSubsystem.java
index babe1f7..f3877fb 100644
--- a/src/main/java/frc/robot/util/scheduling/LifecycleSubsystem.java
+++ b/src/main/java/frc/robot/util/scheduling/LifecycleSubsystem.java
@@ -15,6 +15,8 @@ public class LifecycleSubsystem extends SubsystemBase {
   private final Stopwatch stopwatch = Stopwatch.getInstance();
   private final String loggerName;
 
+  protected final String subsystemName;
+
   private LifecycleStage previousStage = null;
 
   public LifecycleSubsystem(SubsystemPriority priority) {
@@ -23,8 +25,8 @@ public LifecycleSubsystem(SubsystemPriority priority) {
     LifecycleSubsystemManager.getInstance().registerSubsystem(this);
 
     String name = this.getClass().getSimpleName();
-    name = name.substring(name.lastIndexOf('.') + 1);
-    loggerName = "Scheduler/LifecycleSubsystem/" + name + ".periodic()";
+    subsystemName = name.substring(name.lastIndexOf('.') + 1);
+    loggerName = "Scheduler/LifecycleSubsystem/" + subsystemName + ".periodic()";
   }
 
   /** {@link IterativeRobotBase#robotPeriodic()} */
diff --git a/src/main/java/frc/robot/util/state_machines/StateMachine.java b/src/main/java/frc/robot/util/state_machines/StateMachine.java
new file mode 100644
index 0000000..948d885
--- /dev/null
+++ b/src/main/java/frc/robot/util/state_machines/StateMachine.java
@@ -0,0 +1,92 @@
+// Copyright (c) FIRST and other WPILib contributors.
+// Open Source Software; you can modify and/or share it under the terms of
+// the WPILib BSD license file in the root directory of this project.
+
+package frc.robot.util.state_machines;
+
+import dev.doglog.DogLog;
+import edu.wpi.first.wpilibj2.command.Command;
+import edu.wpi.first.wpilibj2.command.Commands;
+import frc.robot.util.scheduling.LifecycleSubsystem;
+import frc.robot.util.scheduling.SubsystemPriority;
+import java.util.Set;
+
+/** A state machine backed by {@link LifecycleSubsystem}. */
+public abstract class StateMachine<S extends Enum<S>> extends LifecycleSubsystem {
+  private S state;
+
+  protected StateMachine(SubsystemPriority priority, S initialState) {
+    super(priority);
+    state = initialState;
+  }
+
+  @Override
+  public void robotPeriodic() {
+    DogLog.log(subsystemName + "/State", state);
+
+    collectInputs();
+
+    var stateBeforeTransitions = state;
+    state = getNextState(state);
+
+    if (state != stateBeforeTransitions) {
+      DogLog.log(subsystemName + "/StateAfterTransition", state);
+      afterTransition(state);
+    } else {
+      DogLog.log(subsystemName + "/StateAfterTransition", "(no change)");
+    }
+  }
+
+  protected void collectInputs() {}
+
+  /**
+   * Process transitions from one state to another.
+   *
+   * @param currentState The current state.
+   * @return The new state after processing transitions.
+   */
+  protected S getNextState(S currentState) {
+    return currentState;
+  }
+
+  /**
+   * Runs once after entering a new state. This is where you should run state actions.
+   *
+   * @param newState The newly entered state.
+   */
+  protected void afterTransition(S newState) {}
+
+  /** Used to change to a new state when a request is made. */
+  protected void setStateFromRequest(S requestedState) {
+    this.state = requestedState;
+  }
+
+  /**
+   * Gets the current state.
+   *
+   * @return The current state.
+   */
+  public S getState() {
+    return state;
+  }
+
+  /**
+   * Creates a command that waits until this state machine is in the given state.
+   *
+   * @param goalState The state to wait for.
+   * @return A command that waits until the state is equal to the goal state.
+   */
+  public Command waitForState(S goalState) {
+    return Commands.waitUntil(() -> this.state == goalState);
+  }
+
+  /**
+   * Creates a command that waits until this state machine is in any of the given states.
+   *
+   * @param goalStates A set of the states to wait for.
+   * @return A command that waits until the state is equal to any of the goal states.
+   */
+  public Command waitForStates(Set<S> goalStates) {
+    return Commands.waitUntil(() -> goalStates.contains(this.state));
+  }
+}