StateMachine.java

// 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 org.wpilib.command3;
 
import static org.wpilib.util.ErrorMessages.requireNonNullParam;
 
import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.function.BooleanSupplier;
import java.util.function.Supplier;
import org.wpilib.annotation.NoDiscard;
import org.wpilib.annotation.PostConstructionInitializer;

public final class StateMachine implements Command {
  private final String m_name;
  private int m_priority = DEFAULT_PRIORITY;
  private State m_initialState = null;
  private final List<State> m_states = new ArrayList<>();
 
  // Marks the state machine as no longer accepting new states. Used when a no-arg switchFromAny()
  // is called to prevent surprising behavior if users try to define new states afterward
  private boolean m_closedForNewStates = false;

  public StateMachine(String name) {
    requireNonNullParam(name, "name", "StateMachine");
    m_name = name;
  }

  public void setPriority(int priority) {
    m_priority = priority;
  }

  @Override
  public int priority() {
    return m_priority;
  }
 
  @Override
  public String name() {
    return m_name;
  }

  @Override
  public Set<Mechanism> requirements() {
    // The machine itself doesn't have any requirements. Commands bound to the various states that
    // the machine moves through may have requirements, however.
    return Set.of();
  }

  @NoDiscard
  public State addState(Command command) {
    requireNonNullParam(command, "command", "StateMachine.addState");
    if (m_closedForNewStates) {
      throw new IllegalStateException(
          "Cannot add new states to a state machine after switchFromAny() has been called with no"
              + " arguments");
    }
    var state = new State(this, command);
    m_states.add(state);
    return state;
  }

  public TransitionNeedsTargetStage switchFromAny(State... states) {
    if (states.length == 0) {
      m_closedForNewStates = true;
      return new TransitionNeedsTargetStage(List.copyOf(m_states));
    } else {
      return new TransitionNeedsTargetStage(List.of(states));
    }
  }

  @PostConstructionInitializer
  public void setInitialState(State initialState) {
    requireNonNullParam(initialState, "initialState", "StateMachine.setInitialState");
    if (!this.equals(initialState.m_stateMachine)) {
      throw new IllegalArgumentException("Cannot set initial state in a different state machine");
    }
    m_initialState = initialState;
  }
 
  @Override
  public void run(Coroutine coroutine) {
    if (m_initialState == null) {
      throw new IllegalStateException(
          m_name + " does not have an initial state. Use .setInitialState() to provide one.");
    }
 
    var currentState = m_initialState;
 
    outer_loop:
    while (currentState != null) {
      final var currentCommand = currentState.command();
      // NOTE: The state machine will exit if the child command fails to be forked. This happens
      //       if the command shares requirements with a higher-priority command.
      coroutine.fork(currentCommand);
      currentState.runEnterCallbacks();
      boolean didYield = false;
 
      while (coroutine.scheduler().isRunning(currentCommand)) {
        for (var transition : currentState.transitions()) {
          if (transition.shouldTransition()) {
            // Cancel the current state's command and move to the next state specified by the
            // transition. Break the state loop early to avoid an unnecessary yield() call and
            // allow the next state's command to start in the same loop iteration that the
            // previous state completed. If the next state is null, the state machine will exit
            // immediately.
            // Note: to prevent infinite loops when states transition to themselves, we require
            // the transition signal to be a rising edge on the user-supplied condition to ensure
            // that the transition is only triggered once per loop iteration.
            currentState.runExitCallbacks();
            coroutine.scheduler().cancel(currentCommand);
            currentState = verifyState(transition.nextState());
            continue outer_loop;
          }
        }
 
        // Yield after checking all transitions.
        // Note: this will be skipped if a transition is triggered.
        coroutine.yield();
        didYield = true;
      }
 
      // Move to the next configured state if no transition was hit before the command completed.
      // We need to be careful about states with oneshot commands; they will complete immediately
      // in the `fork()` call above and never enter the `while` loop and thus never yield.
      // Therefore, we inject a yield call at the end here to ensure that the state machine will
      // always yield once per state. This has a downside of adding extra loop cycles to states that
      // may not need them (and has slightly different behavior to SequentialCommandGroup, which
      // runs commands as fast as possible).
      currentState.runExitCallbacks();
      currentState = verifyState(currentState.nextState());
      if (!didYield && currentState != null) {
        // No need to yield if we're exiting the state machine
        coroutine.yield();
      }
    }
  }
 
  private State verifyState(State next) {
    if (next == null || this.equals(next.m_stateMachine)) {
      // OK
      return next;
    }
 
    // Bad user setup
    throw new IllegalStateException(
        "The next state does not belong to this state machine. Check the state for "
            + next.command().name());
  }

  public static final class State {
    private final StateMachine m_stateMachine;

    private final Command m_command;

    private final List<Completion> m_completions = new ArrayList<>();

    private Supplier<State> m_defaultNextState = () -> null;

    private final List<Transition> m_transitions = new ArrayList<>();
 
    private final List<Runnable> m_enterCallbacks = new ArrayList<>();
    private final List<Runnable> m_exitCallbacks = new ArrayList<>();
 
    private State(StateMachine stateMachine, Command command) {
      m_stateMachine = stateMachine;
      m_command = command;
    }
 
    private Command command() {
      return m_command;
    }
 
    private List<Transition> transitions() {
      return m_transitions;
    }
 
    private void addTransition(Transition transition) {
      m_transitions.add(transition);
    }

    private void setNextState(Supplier<State> nextState) {
      m_defaultNextState = nextState;
    }
 
    // Custom boolean supplier classes may override .equals to do boolean value comparisons,
    // particularly in Kotlin code. Check reference equality instead to just remove bindings to
    // the same condition object.
    @SuppressWarnings("PMD.CompareObjectsWithEquals")
    private void addCompletion(BooleanSupplier condition, Supplier<State> next) {
      // Remove any preexisting completion with the same condition
      m_completions.removeIf(c -> c.getCondition() == condition);
      m_completions.add(new Completion(next, condition));
    }
 
    private State nextState() {
      for (var completion : m_completions) {
        if (completion.shouldTransition()) {
          return completion.nextState();
        }
      }
 
      // No conditional transition has been met, use the default next state.
      // If this was never set or was set to be null, the state machine will exit.
      return m_defaultNextState.get();
    }
 
    private void runEnterCallbacks() {
      m_enterCallbacks.forEach(Runnable::run);
    }
 
    private void runExitCallbacks() {
      m_exitCallbacks.forEach(Runnable::run);
    }

    public void onEnter(Runnable callback) {
      requireNonNullParam(callback, "callback", "State.onEnter");
      m_enterCallbacks.add(callback);
    }

    public void onExit(Runnable callback) {
      requireNonNullParam(callback, "callback", "State.onExit");
      m_exitCallbacks.add(callback);
    }

    public TransitionNeedsConditionStage switchTo(State to) {
      requireNonNullParam(to, "to", "State.switchTo");
      if (!m_stateMachine.equals(to.m_stateMachine)) {
        throw new IllegalArgumentException(
            "Cannot transition to a state in a different state machine");
      }
      return new TransitionNeedsTargetStage(List.of(this)).to(to);
    }

    public TransitionNeedsConditionStage switchTo(Supplier<State> dynamic) {
      requireNonNullParam(dynamic, "dynamic", "State.switchTo");
      // Unfortunately, we can't check up front that the supplier will always return a state for
      // this state machine. The output will need to be checked when the supplier is called
      return new TransitionNeedsTargetStage(List.of(this)).to(dynamic);
    }

    public TransitionNeedsConditionStage exitStateMachine() {
      return new TransitionNeedsConditionStage(List.of(this), () -> null);
    }
  }

  @NoDiscard("Use .to() to specify the target state")
  public static final class TransitionNeedsTargetStage {
    private final List<State> m_from;
 
    private TransitionNeedsTargetStage(List<State> from) {
      m_from = from;
    }

    public TransitionNeedsConditionStage to(State to) {
      requireNonNullParam(to, "to", "NeedsTargetTransitionBuilder.to");
      for (var state : m_from) {
        if (!state.m_stateMachine.equals(to.m_stateMachine)) {
          throw new IllegalArgumentException(
              "Cannot transition to a state in a different state machine");
        }
      }
      return new TransitionNeedsConditionStage(m_from, () -> to);
    }

    public TransitionNeedsConditionStage to(Supplier<State> dynamic) {
      requireNonNullParam(dynamic, "dynamic", "NeedsTargetTransitionBuilder.to");
      return new TransitionNeedsConditionStage(m_from, dynamic);
    }

    public TransitionNeedsConditionStage toExitStateMachine() {
      return new TransitionNeedsConditionStage(m_from, () -> null);
    }
  }

  @NoDiscard("Use .when() or .whenComplete() to specify the transition condition")
  public static final class TransitionNeedsConditionStage {
    private final List<State> m_originatingStates;
 
    // Note: A null result from the supplier indicates that the transition will cause the state
    //       machine to exit
    private final Supplier<State> m_targetStateSupplier;
 
    private TransitionNeedsConditionStage(List<State> from, Supplier<State> to) {
      requireNonNullParam(from, "from", "TransitionNeedsConditionStage");
      requireNonNullParam(to, "to", "TransitionNeedsConditionStage");
 
      m_originatingStates = from;
      m_targetStateSupplier = to;
    }

    public void when(BooleanSupplier condition) {
      requireNonNullParam(condition, "condition", "NeedsConditionTransitionBuilder.when");
      var transition = new Transition(m_targetStateSupplier, condition);
      m_originatingStates.forEach(originatingState -> originatingState.addTransition(transition));
    }

    public void whenComplete() {
      m_originatingStates.forEach(state -> state.setNextState(m_targetStateSupplier));
    }

    public void whenCompleteAnd(BooleanSupplier condition) {
      requireNonNullParam(condition, "condition", "NeedsConditionTransitionBuilder.whenComplete");
      m_originatingStates.forEach(state -> state.addCompletion(condition, m_targetStateSupplier));
    }
  }

  private static final class Completion {
    private final Supplier<State> m_nextSupplier;
    private final BooleanSupplier m_condition;

    private Completion(Supplier<State> next, BooleanSupplier condition) {
      m_nextSupplier = next;
      m_condition = condition;
    }
 
    private boolean shouldTransition() {
      return m_condition.getAsBoolean();
    }
 
    public State nextState() {
      return m_nextSupplier.get();
    }
 
    public BooleanSupplier getCondition() {
      return m_condition;
    }
  }

  private static final class Transition {
    private final Supplier<State> m_nextSupplier;

    private final BooleanSupplier m_condition;
 
    private boolean m_previousSignal = false;
 
    private Transition(Supplier<State> next, BooleanSupplier condition) {
      m_nextSupplier = next;
      m_condition = condition;
    }

    private boolean shouldTransition() {
      // Wrap the condition in a rising edge detector so that it will only trigger a single time per
      // loop iteration. This prevents issues with a state transitioning to itself like so:
      // state1.switchTo(state1).when(() -> foo == true);
      // If the condition is itself a rising edge detector, this wrapping is redundant but harmless.
      boolean currentValue = m_condition.getAsBoolean();
      boolean isRisingEdge = currentValue && !m_previousSignal;
      m_previousSignal = currentValue;
      return isRisingEdge;
    }
 
    private State nextState() {
      return m_nextSupplier.get();
    }
  }
}