[commands] Add withDeadline modifier (#7299)

Co-authored-by: Ryan Blue <[email protected]>

wpilibNewCommands/src/main/java/edu/wpi/first/wpilibj2/command/Command.java

@@ -287,6 +287,24 @@ public SequentialCommandGroup andThen(Command... next) {
     return group;
   }
 
+  /**
+   * Creates a new command that runs this command and the deadline in parallel, finishing (and
+   * interrupting this command) when the deadline finishes.
+   *
+   * <p>Note: This decorator works by adding this command to a composition. The command the
+   * decorator was called on cannot be scheduled independently or be added to a different
+   * composition (namely, decorators), unless it is manually cleared from the list of composed
+   * commands with {@link CommandScheduler#removeComposedCommand(Command)}. The command composition
+   * returned from this method can be further decorated without issue.
+   *
+   * @param deadline the deadline of the command group
+   * @return the decorated command
+   * @see Command#deadlineFor
+   */
+  public ParallelDeadlineGroup withDeadline(Command deadline) {
+    return new ParallelDeadlineGroup(deadline, this);
+  }
+
   /**
    * Decorates this command with a set of commands to run parallel to it, ending when the calling
    * command ends and interrupting all the others. Often more convenient/less-verbose than
@@ -321,6 +339,7 @@ public ParallelDeadlineGroup deadlineWith(Command... parallel) {
    * @param parallel the commands to run in parallel. Note the parallel commands will be interrupted
    *     when the deadline command ends
    * @return the decorated command
+   * @see Command#withDeadline
    */
   public ParallelDeadlineGroup deadlineFor(Command... parallel) {
     return new ParallelDeadlineGroup(this, parallel);

wpilibNewCommands/src/main/native/cpp/frc2/command/Command.cpp

@@ -121,6 +121,10 @@ CommandPtr Command::OnlyIf(std::function<bool()> condition) && {
   return std::move(*this).ToPtr().OnlyIf(std::move(condition));
 }
 
+CommandPtr Command::WithDeadline(CommandPtr&& deadline) && {
+  return std::move(*this).ToPtr().WithDeadline(std::move(deadline));
+}
+
 CommandPtr Command::DeadlineFor(CommandPtr&& parallel) && {
   return std::move(*this).ToPtr().DeadlineFor(std::move(parallel));
 }

wpilibNewCommands/src/main/native/cpp/frc2/command/CommandPtr.cpp

@@ -168,6 +168,15 @@ CommandPtr CommandPtr::OnlyIf(std::function<bool()> condition) && {
   return std::move(*this).Unless(std::not_fn(std::move(condition)));
 }
 
+CommandPtr CommandPtr::WithDeadline(CommandPtr&& deadline) && {
+  AssertValid();
+  std::vector<std::unique_ptr<Command>> vec;
+  vec.emplace_back(std::move(m_ptr));
+  m_ptr = std::make_unique<ParallelDeadlineGroup>(std::move(deadline).Unwrap(),
+                                                  std::move(vec));
+  return std::move(*this);
+}
+
 CommandPtr CommandPtr::DeadlineWith(CommandPtr&& parallel) && {
   AssertValid();
   std::vector<std::unique_ptr<Command>> vec;

wpilibNewCommands/src/main/native/include/frc2/command/Command.h

@@ -309,6 +309,16 @@ class Command : public wpi::Sendable, public wpi::SendableHelper<Command> {
   [[nodiscard]]
   CommandPtr OnlyIf(std::function<bool()> condition) &&;
 
+  /**
+   * Creates a new command that runs this command and the deadline in parallel,
+   * finishing (and interrupting this command) when the deadline finishes.
+   *
+   * @param deadline the deadline of the command group
+   * @return the decorated command
+   * @see DeadlineFor
+   */
+  CommandPtr WithDeadline(CommandPtr&& deadline) &&;
+
   /**
    * Decorates this command with a set of commands to run parallel to it, ending
    * when the calling command ends and interrupting all the others. Often more
@@ -318,9 +328,11 @@ class Command : public wpi::Sendable, public wpi::SendableHelper<Command> {
    * @param parallel the commands to run in parallel. Note the parallel commands
    * will be interupted when the deadline command ends
    * @return the decorated command
+   * @see WithDeadline
    */
   [[nodiscard]]
   CommandPtr DeadlineFor(CommandPtr&& parallel) &&;
+
   /**
    * Decorates this command with a set of commands to run parallel to it, ending
    * when the last command ends. Often more convenient/less-verbose than

wpilibNewCommands/src/main/native/include/frc2/command/CommandPtr.h

@@ -182,6 +182,16 @@ class CommandPtr final {
   [[nodiscard]]
   CommandPtr OnlyIf(std::function<bool()> condition) &&;
 
+  /**
+   * Creates a new command that runs this command and the deadline in parallel,
+   * finishing (and interrupting this command) when the deadline finishes.
+   *
+   * @param deadline the deadline of the command group
+   * @return the decorated command
+   * @see DeadlineFor
+   */
+  CommandPtr WithDeadline(CommandPtr&& deadline) &&;
+
   /**
    * Decorates this command with a set of commands to run parallel to it, ending
    * when the calling command ends and interrupting all the others. Often more

wpilibNewCommands/src/test/java/edu/wpi/first/wpilibj2/command/CommandDecoratorTest.java

@@ -271,6 +271,57 @@ void deadlineForOrderTest() {
     }
   }
 
+  @Test
+  void withDeadlineTest() {
+    try (CommandScheduler scheduler = new CommandScheduler()) {
+      AtomicBoolean finish = new AtomicBoolean(false);
+
+      Command endsBeforeGroup = Commands.none().withDeadline(Commands.waitUntil(finish::get));
+      scheduler.schedule(endsBeforeGroup);
+      scheduler.run();
+      assertTrue(scheduler.isScheduled(endsBeforeGroup));
+      finish.set(true);
+      scheduler.run();
+      assertFalse(scheduler.isScheduled(endsBeforeGroup));
+      finish.set(false);
+
+      Command endsAfterGroup = Commands.idle().withDeadline(Commands.waitUntil(finish::get));
+      scheduler.schedule(endsAfterGroup);
+      scheduler.run();
+      assertTrue(scheduler.isScheduled(endsAfterGroup));
+      finish.set(true);
+      scheduler.run();
+      assertFalse(scheduler.isScheduled(endsAfterGroup));
+    }
+  }
+
+  @Test
+  void withDeadlineOrderTest() {
+    try (CommandScheduler scheduler = new CommandScheduler()) {
+      AtomicBoolean dictatorHasRun = new AtomicBoolean(false);
+      AtomicBoolean dictatorWasPolled = new AtomicBoolean(false);
+      Command dictator =
+          new FunctionalCommand(
+              () -> {},
+              () -> dictatorHasRun.set(true),
+              interrupted -> {},
+              () -> {
+                dictatorWasPolled.set(true);
+                return true;
+              });
+      Command other =
+          Commands.run(
+              () ->
+                  assertAll(
+                      () -> assertTrue(dictatorHasRun.get()),
+                      () -> assertTrue(dictatorWasPolled.get())));
+      Command group = other.withDeadline(dictator);
+      scheduler.schedule(group);
+      scheduler.run();
+      assertAll(() -> assertTrue(dictatorHasRun.get()), () -> assertTrue(dictatorWasPolled.get()));
+    }
+  }
+
   @Test
   void alongWithTest() {
     try (CommandScheduler scheduler = new CommandScheduler()) {

wpilibNewCommands/src/test/native/cpp/frc2/command/CommandDecoratorTest.cpp

@@ -221,6 +221,27 @@ TEST_F(CommandDecoratorTest, DeadlineFor) {
   EXPECT_FALSE(scheduler.IsScheduled(group));
 }
 
+TEST_F(CommandDecoratorTest, WithDeadline) {
+  CommandScheduler scheduler = GetScheduler();
+
+  bool finish = false;
+
+  auto dictator = WaitUntilCommand([&finish] { return finish; });
+  auto endsAfter = WaitUntilCommand([] { return false; });
+
+  auto group = std::move(endsAfter).WithDeadline(std::move(dictator).ToPtr());
+
+  scheduler.Schedule(group);
+  scheduler.Run();
+
+  EXPECT_TRUE(scheduler.IsScheduled(group));
+
+  finish = true;
+  scheduler.Run();
+
+  EXPECT_FALSE(scheduler.IsScheduled(group));
+}
+
 TEST_F(CommandDecoratorTest, AlongWith) {
   CommandScheduler scheduler = GetScheduler();
 
@@ -283,6 +304,33 @@ TEST_F(CommandDecoratorTest, DeadlineForOrder) {
   EXPECT_TRUE(dictatorWasPolled);
 }
 
+TEST_F(CommandDecoratorTest, WithDeadlineOrder) {
+  CommandScheduler scheduler = GetScheduler();
+
+  bool dictatorHasRun = false;
+  bool dictatorWasPolled = false;
+
+  auto dictator =
+      FunctionalCommand([] {}, [&dictatorHasRun] { dictatorHasRun = true; },
+                        [](bool interrupted) {},
+                        [&dictatorWasPolled] {
+                          dictatorWasPolled = true;
+                          return true;
+                        });
+  auto other = RunCommand([&dictatorHasRun, &dictatorWasPolled] {
+    EXPECT_TRUE(dictatorHasRun);
+    EXPECT_TRUE(dictatorWasPolled);
+  });
+
+  auto group = std::move(other).WithDeadline(std::move(dictator).ToPtr());
+
+  scheduler.Schedule(group);
+  scheduler.Run();
+
+  EXPECT_TRUE(dictatorHasRun);
+  EXPECT_TRUE(dictatorWasPolled);
+}
+
 TEST_F(CommandDecoratorTest, AlongWithOrder) {
   CommandScheduler scheduler = GetScheduler();