Coroutine.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.command3.CoroutineState.CANCELED;
import static org.wpilib.command3.CoroutineState.Canceled;
import static org.wpilib.command3.CoroutineState.FROZEN;
import static org.wpilib.command3.CoroutineState.ForkFailed;
import static org.wpilib.command3.CoroutineState.LIVE;
import static org.wpilib.command3.CoroutineState.forkFailed;
import static org.wpilib.command3.CoroutineState.parkedUntil;
import static org.wpilib.command3.CoroutineState.parkedWhile;
import static org.wpilib.units.Units.Nanoseconds;
import static org.wpilib.units.Units.Seconds;
import static org.wpilib.util.ErrorMessages.requireNonNullParam;
 
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.function.BooleanSupplier;
import java.util.function.Consumer;
import org.wpilib.command3.Scheduler.ScheduleResult;
import org.wpilib.system.RobotController;
import org.wpilib.units.measure.Time;

public final class Coroutine {
  private final Scheduler m_scheduler;
  private final Continuation m_backingContinuation;
 
  private boolean m_cancelOnForkFailure = true;
 
  private CoroutineState m_state = LIVE;
 
  boolean inUnmountableState() {
    return !m_state.shouldMount();
  }

  Coroutine(Scheduler scheduler, ContinuationScope scope, Consumer<Coroutine> callback) {
    m_scheduler = scheduler;
    m_backingContinuation = new Continuation(scope, () -> callback.accept(this));
  }

  public void setCancelOnForkFailure(boolean cancelOnForkFailure) {
    m_cancelOnForkFailure = cancelOnForkFailure;
  }

  public boolean isCancelOnForkFailure() {
    return m_cancelOnForkFailure;
  }
 
  // Package-private. User code can never access this because setting the flag happens immediately
  // before the coroutine yields itself.
  boolean isInterruptRequested() {
    // Note that we only get into a ForkFailed state if m_cancelOnForkFailure is true,
    // so there's no point in checking it here
    return m_state instanceof ForkFailed;
  }

  public void requestCancellation() {
    requireMounted();
 
    m_state = CANCELED;
    this.yield();
  }
 
  // package-private; setting the flag will immediately yield, so user code can never access this
  // in a state where it returns anything other than `false`
  boolean isCancellationRequested() {
    return m_state instanceof Canceled;
  }
 
  ForkResult getForkResult() {
    if (m_state instanceof ForkFailed(ForkResult failure)) {
      return failure;
    } else {
      return null;
    }
  }

  public boolean yield() {
    requireMounted();
 
    m_backingContinuation.yield();
 
    // If the coroutine was picked back up by the scheduler, it is definitionally live.
    // Reset the state to reflect this.
    m_state = LIVE;
    return true;
  }

  public void park() {
    requireMounted();
 
    m_state = FROZEN;
    this.yield();
  }

  public final class ForkResult {
    private final List<ScheduleResult.Successful> m_forkedCommands;
    private final List<ScheduleResult.Failure> m_failedCommands;
 
    // private - only the coroutine can construct these
    private ForkResult(
        Collection<ScheduleResult.Successful> successes, List<ScheduleResult.Failure> failures) {
      m_forkedCommands = List.copyOf(successes);
      m_failedCommands = List.copyOf(failures);
    }

    public boolean successful() {
      return m_failedCommands.isEmpty();
    }

    public boolean failed() {
      return !successful();
    }

    public boolean partialSuccess() {
      return !m_forkedCommands.isEmpty() && !m_failedCommands.isEmpty();
    }

    public void awaitCompletion() {
      if (m_forkedCommands.isEmpty()) {
        // Nothing to do; exit early
        return;
      }
      var tracker = CommandRunTracker.of(m_scheduler, m_forkedCommands);
      while (tracker.isAnyRunning()) {
        Coroutine.this.yield();
      }
    }

    public List<ScheduleResult.Failure> getFailedCommands() {
      return m_failedCommands;
    }

    public List<ScheduleResult.Successful> getForkedCommands() {
      return m_forkedCommands;
    }
  }

  private ForkResult checkAllForkable(Collection<? extends Command> commands) {
    // Check for user error; there's no reason to fork conflicting commands simultaneously
    // Because this is a bug in user code, throw an error instead of returning a failure result
    ConflictDetector.throwIfConflicts(commands);
 
    var successes = new ArrayList<ScheduleResult.Successful>();
    var failures = new ArrayList<ScheduleResult.Failure>();
    for (var command : commands) {
      var result = m_scheduler.isSchedulable(command);
      switch (result) {
        case ScheduleResult.Successful success -> successes.add(success);
        case ScheduleResult.Failure failure -> failures.add(failure);
      }
    }
 
    var result = new ForkResult(successes, failures);
    if (result.failed()) {
      if (m_cancelOnForkFailure) {
        // Canceling on fork failure means no coroutine or user code gets to run to handle the
        // failure result
        m_state = forkFailed(result);
        this.yield();
        // could throw an IllegalStateException, but probably shouldn't crash user code
        return result;
      } else {
        // always return the failure object if self-cancellation is disabled
        return result;
      }
    }
 
    return result;
  }

  private ForkResult doFork(Collection<? extends Command> commands) {
    var successes = new ArrayList<ScheduleResult.Successful>();
    var fails = new ArrayList<ScheduleResult.Failure>();
    for (var command : commands) {
      var result = m_scheduler.schedule(command);
      switch (result) {
        case ScheduleResult.Successful s -> successes.add(s);
        case ScheduleResult.Failure f -> fails.add(f);
      }
    }
    ForkResult result = new ForkResult(successes, fails);
 
    if (!fails.isEmpty() && m_cancelOnForkFailure) {
      // At least one of the commands couldn't be scheduled due to runtime behavior that can't be
      // checked ahead of time by `checkAllForkable`.
      //
      // We can't roll back the partial success; the commands that successfully started will have
      // already affected the state of the robot and may have interrupted previously running
      // commands, which we can't restore. The two options are to either cancel the entire command
      // composition - which is the default behavior - or to allow user code to handle the failure
      // and do something with it (eg trying something else, retrying, etc).
      m_state = forkFailed(result);
      this.yield();
      return result; // note: unreachable
    }
 
    return result;
  }

  public ForkResult fork(Command... commands) {
    requireMounted();
 
    requireNonNullParam(commands, "commands", "Coroutine.fork");
    for (int i = 0; i < commands.length; i++) {
      requireNonNullParam(commands[i], "commands[" + i + "]", "Coroutine.fork");
    }
 
    var forkableCheck = checkAllForkable(Arrays.asList(commands));
    if (forkableCheck.failed()) {
      return forkableCheck;
    }
 
    return doFork(Arrays.asList(commands));
  }

  public ForkResult fork(Collection<? extends Command> commands) {
    return fork(commands.toArray(Command[]::new));
  }

  public ForkResult await(Command command) {
    requireMounted();
 
    requireNonNullParam(command, "command", "Coroutine.await");
 
    // Don't need doFork() because there's no chance of sibling conflicts
    var scheduleResult = m_scheduler.schedule(command);
    switch (scheduleResult) {
      case ScheduleResult.Successful forked -> {
        int id = forked.runId();
        // If the command is a one-shot, then the schedule call will completely execute the command.
        // runId(command) will return 0 and the loop condition will never be met, so we'd return
        // immediately without yielding.
        var cond = parkedWhile(() -> m_scheduler.runId(forked.command()) == id);
        if (!cond.conditionMet()) {
          m_state = cond;
          this.yield();
        }
 
        return new ForkResult(List.of(forked), List.of());
      }
 
      case ScheduleResult.Failure failed -> {
        // Failed to fork
        var result = new ForkResult(List.of(), List.of(failed));
        if (m_cancelOnForkFailure) {
          m_state = forkFailed(result);
          this.yield();
          return result; // note: unreachable because the scheduler will never remount the coroutine
        }
        return result;
      }
    }
  }

  public ForkResult awaitAll(Collection<? extends Command> commands) {
    requireMounted();
 
    requireNonNullParam(commands, "commands", "Coroutine.awaitAll");
    int i = 0;
    for (Command command : commands) {
      requireNonNullParam(command, "commands[" + i + "]", "Coroutine.awaitAll");
      i++;
    }
 
    var forkableCheck = checkAllForkable(commands);
    if (forkableCheck.failed()) {
      return forkableCheck;
    }
 
    var forkResult = doFork(commands);
    if (forkResult.failed()) {
      return forkResult;
    }
 
    var tracker = CommandRunTracker.of(m_scheduler, forkResult.getForkedCommands());
    var cond = parkedWhile(tracker::isAnyRunning);
    if (!cond.conditionMet()) {
      m_state = cond;
      this.yield();
    }
 
    return forkResult;
  }

  public ForkResult awaitAll(Command... commands) {
    return awaitAll(Arrays.asList(commands));
  }

  public ForkResult awaitAny(Collection<? extends Command> commands) {
    requireMounted();
 
    requireNonNullParam(commands, "commands", "Coroutine.awaitAny");
    int i = 0;
    for (Command command : commands) {
      requireNonNullParam(command, "commands[" + i + "]", "Coroutine.awaitAny");
      i++;
    }
 
    var forkableCheck = checkAllForkable(commands);
    if (forkableCheck.failed()) {
      return forkableCheck;
    }
 
    var forkResult = doFork(commands);
    if (forkResult.failed()) {
      return forkResult;
    }
 
    var tracker = CommandRunTracker.of(m_scheduler, forkResult.getForkedCommands());
    var cond = parkedWhile(tracker::areAllRunning);
    if (!cond.conditionMet()) {
      m_state = cond;
      this.yield();
    }
 
    // At least one command exited; cancel the rest.
    // We need to use run IDs here because at least one command exited in the previous cycle.
    // If that command was immediately rescheduled, that would have happened outside this awaitAny
    // call and must be allowed to continue. Using a naive `commands.forEach(m_scheduler::cancel)`
    // loop would cause that rescheduled command to be canceled as well.
    for (var fork : forkResult.getForkedCommands()) {
      var command = fork.command();
      int runId = fork.runId();
 
      if (m_scheduler.runId(command) == runId) {
        // Only cancel an original invocation of the command. If it exited early and was later
        // rescheduled, that rescheduled run must be allowed to continue.
        m_scheduler.cancel(command);
      }
    }
 
    return forkResult;
  }

  public ForkResult awaitAny(Command... commands) {
    return awaitAny(Arrays.asList(commands));
  }

  public void wait(Time duration) {
    requireMounted();
 
    requireNonNullParam(duration, "duration", "Coroutine.wait");
 
    var cond = parkedUntil(RobotController.getTime() + (long) duration.in(Nanoseconds));
    if (duration.in(Seconds) > 0) {
      m_state = cond;
      this.yield();
    }
  }

  public enum WaitResult {
    CONDITION_MET {
      @Override
      public boolean timedOut() {
        return false;
      }
    },

    TIMED_OUT {
      @Override
      public boolean timedOut() {
        return true;
      }
    },
    ;

    public abstract boolean timedOut();
  }

  public WaitResult waitUntil(BooleanSupplier condition) {
    requireMounted();
 
    requireNonNullParam(condition, "condition", "Coroutine.waitUntil");
 
    var cond = parkedUntil(condition);
    if (!cond.conditionMet()) {
      m_state = cond;
      this.yield();
    }
 
    return WaitResult.CONDITION_MET;
  }

  public WaitResult waitUntil(BooleanSupplier condition, Time timeout) {
    requireMounted();
 
    requireNonNullParam(condition, "condition", "Coroutine.waitUntil");
    requireNonNullParam(timeout, "timeout", "Coroutine.waitUntil");
 
    if (condition.getAsBoolean()) {
      return WaitResult.CONDITION_MET;
    }
 
    long targetTimestampNanos = RobotController.getTime() + (long) timeout.in(Nanoseconds);
    m_state =
        parkedUntil(
            () -> condition.getAsBoolean() || RobotController.getTime() >= targetTimestampNanos);
    this.yield();
 
    // In case of a tiebreaker where the condition and the timeout are both reached in the same
    // cycle, let the condition win.
    if (condition.getAsBoolean()) {
      return WaitResult.CONDITION_MET;
    } else {
      return WaitResult.TIMED_OUT;
    }
  }

  public Scheduler scheduler() {
    requireMounted();
 
    return m_scheduler;
  }
 
  private boolean isMounted() {
    return m_backingContinuation.isMounted();
  }
 
  private void requireMounted() {
    // Note: attempting to yield() outside a command will already throw an error due to the
    // continuation being unmounted, but other actions like forking and awaiting should also
    // throw errors. For consistent messaging, we use this helper in all places, not just the
    // ones that interact with the backing continuation.
 
    if (isMounted()) {
      return;
    }
 
    throw new IllegalStateException("Coroutines can only be used by the command bound to them");
  }
 
  // Package-private for interaction with the scheduler.
  // These functions are not intended for team use.
 
  void runToYieldPoint() {
    m_backingContinuation.run();
  }
 
  void mount() {
    Continuation.mountContinuation(m_backingContinuation);
  }
 
  boolean isDone() {
    return m_backingContinuation.isDone();
  }
}