diff --git a/wpilibNewCommands/src/main/java/edu/wpi/first/wpilibj2/command/Command.java b/wpilibNewCommands/src/main/java/edu/wpi/first/wpilibj2/command/Command.java index f154817a0f..90d5f6bc37 100644 --- a/wpilibNewCommands/src/main/java/edu/wpi/first/wpilibj2/command/Command.java +++ b/wpilibNewCommands/src/main/java/edu/wpi/first/wpilibj2/command/Command.java @@ -4,6 +4,9 @@ package edu.wpi.first.wpilibj2.command; +import static edu.wpi.first.wpilibj.util.ErrorMessages.requireNonNullParam; + +import edu.wpi.first.util.function.BooleanConsumer; import java.util.Set; import java.util.function.BooleanSupplier; @@ -348,6 +351,42 @@ public interface Command { }; } + /** + * Decorates this command with a lambda to call on interrupt or end, following the command's + * inherent {@link #end(boolean)} method. + * + * @param end a lambda accepting a boolean parameter specifying whether the command was + * interrupted. + * @return the decorated command + */ + default WrapperCommand finallyDo(BooleanConsumer end) { + requireNonNullParam(end, "end", "Command.finallyDo()"); + return new WrapperCommand(this) { + @Override + public void end(boolean interrupted) { + super.end(interrupted); + end.accept(interrupted); + } + }; + } + + /** + * Decorates this command with a lambda to call on interrupt, following the command's inherent + * {@link #end(boolean)} method. + * + * @param handler a lambda to run when the command is interrupted + * @return the decorated command + */ + default WrapperCommand handleInterrupt(Runnable handler) { + requireNonNullParam(handler, "handler", "Command.handleInterrupt()"); + return finallyDo( + interrupted -> { + if (interrupted) { + handler.run(); + } + }); + } + /** Schedules this command. */ default void schedule() { CommandScheduler.getInstance().schedule(this); diff --git a/wpilibNewCommands/src/main/native/cpp/frc2/command/Command.cpp b/wpilibNewCommands/src/main/native/cpp/frc2/command/Command.cpp index 3a55ae82ce..f170e407b4 100644 --- a/wpilibNewCommands/src/main/native/cpp/frc2/command/Command.cpp +++ b/wpilibNewCommands/src/main/native/cpp/frc2/command/Command.cpp @@ -108,6 +108,16 @@ CommandPtr Command::Unless(std::function condition) && { .Unless(std::move(condition)); } +CommandPtr Command::FinallyDo(std::function end) && { + return CommandPtr(std::move(*this).TransferOwnership()) + .FinallyDo(std::move(end)); +} + +CommandPtr Command::HandleInterrupt(std::function handler) && { + return CommandPtr(std::move(*this).TransferOwnership()) + .HandleInterrupt(std::move(handler)); +} + void Command::Schedule() { CommandScheduler::GetInstance().Schedule(this); } diff --git a/wpilibNewCommands/src/main/native/cpp/frc2/command/CommandPtr.cpp b/wpilibNewCommands/src/main/native/cpp/frc2/command/CommandPtr.cpp index 1ca1284cf0..73ab14230c 100644 --- a/wpilibNewCommands/src/main/native/cpp/frc2/command/CommandPtr.cpp +++ b/wpilibNewCommands/src/main/native/cpp/frc2/command/CommandPtr.cpp @@ -167,6 +167,37 @@ CommandPtr CommandPtr::RaceWith(CommandPtr&& parallel) && { return std::move(*this); } +namespace { +class FinallyCommand : public WrapperCommand { + public: + FinallyCommand(std::unique_ptr&& command, + std::function end) + : WrapperCommand(std::move(command)), m_end(std::move(end)) {} + + void End(bool interrupted) override { + WrapperCommand::End(interrupted); + m_end(interrupted); + } + + private: + std::function m_end; +}; +} // namespace + +CommandPtr CommandPtr::FinallyDo(std::function end) && { + m_ptr = std::make_unique(std::move(m_ptr), std::move(end)); + return std::move(*this); +} + +CommandPtr CommandPtr::HandleInterrupt(std::function handler) && { + return std::move(*this).FinallyDo( + [handler = std::move(handler)](bool interrupted) { + if (interrupted) { + handler(); + } + }); +} + Command* CommandPtr::get() const { return m_ptr.get(); } diff --git a/wpilibNewCommands/src/main/native/include/frc2/command/Command.h b/wpilibNewCommands/src/main/native/include/frc2/command/Command.h index 72e983f0f8..cc1d83780d 100644 --- a/wpilibNewCommands/src/main/native/include/frc2/command/Command.h +++ b/wpilibNewCommands/src/main/native/include/frc2/command/Command.h @@ -261,6 +261,25 @@ class Command { [[nodiscard]] CommandPtr WithInterruptBehavior( Command::InterruptionBehavior interruptBehavior) &&; + /** + * Decorates this command with a lambda to call on interrupt or end, following + * the command's inherent Command::End(bool) method. + * + * @param end a lambda accepting a boolean parameter specifying whether the + * command was interrupted. + * @return the decorated command + */ + [[nodiscard]] CommandPtr FinallyDo(std::function end) &&; + + /** + * Decorates this command with a lambda to call on interrupt, following the + * command's inherent Command::End(bool) method. + * + * @param handler a lambda to run when the command is interrupted + * @return the decorated command + */ + [[nodiscard]] CommandPtr HandleInterrupt(std::function handler) &&; + /** * Schedules this command. */ diff --git a/wpilibNewCommands/src/main/native/include/frc2/command/CommandPtr.h b/wpilibNewCommands/src/main/native/include/frc2/command/CommandPtr.h index 321f014810..fbe4855124 100644 --- a/wpilibNewCommands/src/main/native/include/frc2/command/CommandPtr.h +++ b/wpilibNewCommands/src/main/native/include/frc2/command/CommandPtr.h @@ -207,6 +207,25 @@ class CommandPtr final { */ [[nodiscard]] CommandPtr RaceWith(CommandPtr&& parallel) &&; + /** + * Decorates this command with a lambda to call on interrupt or end, following + * the command's inherent Command::End(bool) method. + * + * @param end a lambda accepting a boolean parameter specifying whether the + * command was interrupted. + * @return the decorated command + */ + [[nodiscard]] CommandPtr FinallyDo(std::function end) &&; + + /** + * Decorates this command with a lambda to call on interrupt, following the + * command's inherent Command::End(bool) method. + * + * @param handler a lambda to run when the command is interrupted + * @return the decorated command + */ + [[nodiscard]] CommandPtr HandleInterrupt(std::function handler) &&; + /** * Get a raw pointer to the held command. */ diff --git a/wpilibNewCommands/src/test/java/edu/wpi/first/wpilibj2/command/CommandDecoratorTest.java b/wpilibNewCommands/src/test/java/edu/wpi/first/wpilibj2/command/CommandDecoratorTest.java index 3cba1e59a6..8cfc629200 100644 --- a/wpilibNewCommands/src/test/java/edu/wpi/first/wpilibj2/command/CommandDecoratorTest.java +++ b/wpilibNewCommands/src/test/java/edu/wpi/first/wpilibj2/command/CommandDecoratorTest.java @@ -4,12 +4,14 @@ package edu.wpi.first.wpilibj2.command; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; import edu.wpi.first.hal.HAL; import edu.wpi.first.wpilibj.simulation.SimHooks; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.parallel.ResourceLock; @@ -234,4 +236,75 @@ class CommandDecoratorTest extends CommandTestBase { assertTrue(hasRunCondition.get()); } } + + @Test + void finallyDoTest() { + try (CommandScheduler scheduler = new CommandScheduler()) { + AtomicInteger first = new AtomicInteger(0); + AtomicInteger second = new AtomicInteger(0); + + Command command = + new FunctionalCommand( + () -> {}, + () -> {}, + interrupted -> { + if (!interrupted) { + first.incrementAndGet(); + } + }, + () -> true) + .finallyDo( + interrupted -> { + if (!interrupted) { + // to differentiate between "didn't run" and "ran before command's `end()` + second.addAndGet(1 + first.get()); + } + }); + + scheduler.schedule(command); + assertEquals(0, first.get()); + assertEquals(0, second.get()); + scheduler.run(); + assertEquals(1, first.get()); + // if `second == 0`, neither of the lambdas ran. + // if `second == 1`, the second lambda ran before the first one + assertEquals(2, second.get()); + } + } + + // handleInterruptTest() implicitly tests the interrupt=true branch of finallyDo() + @Test + void handleInterruptTest() { + try (CommandScheduler scheduler = new CommandScheduler()) { + AtomicInteger first = new AtomicInteger(0); + AtomicInteger second = new AtomicInteger(0); + + Command command = + new FunctionalCommand( + () -> {}, + () -> {}, + interrupted -> { + if (interrupted) { + first.incrementAndGet(); + } + }, + () -> false) + .handleInterrupt( + () -> { + // to differentiate between "didn't run" and "ran before command's `end()` + second.addAndGet(1 + first.get()); + }); + + scheduler.schedule(command); + scheduler.run(); + assertEquals(0, first.get()); + assertEquals(0, second.get()); + + scheduler.cancel(command); + assertEquals(1, first.get()); + // if `second == 0`, neither of the lambdas ran. + // if `second == 1`, the second lambda ran before the first one + assertEquals(2, second.get()); + } + } } diff --git a/wpilibNewCommands/src/test/native/cpp/frc2/command/CommandDecoratorTest.cpp b/wpilibNewCommands/src/test/native/cpp/frc2/command/CommandDecoratorTest.cpp index 76e9cb39d7..d91ecb2e9f 100644 --- a/wpilibNewCommands/src/test/native/cpp/frc2/command/CommandDecoratorTest.cpp +++ b/wpilibNewCommands/src/test/native/cpp/frc2/command/CommandDecoratorTest.cpp @@ -7,6 +7,7 @@ #include "CommandTestBase.h" #include "frc2/command/ConditionalCommand.h" #include "frc2/command/EndlessCommand.h" +#include "frc2/command/FunctionalCommand.h" #include "frc2/command/InstantCommand.h" #include "frc2/command/ParallelRaceGroup.h" #include "frc2/command/PerpetualCommand.h" @@ -152,3 +153,62 @@ TEST_F(CommandDecoratorTest, Unless) { scheduler.Run(); EXPECT_TRUE(hasRun); } + +TEST_F(CommandDecoratorTest, FinallyDo) { + CommandScheduler scheduler = GetScheduler(); + int first = 0; + int second = 0; + CommandPtr command = FunctionalCommand([] {}, [] {}, + [&first](bool interrupted) { + if (!interrupted) { + first++; + } + }, + [] { return true; }) + .FinallyDo([&first, &second](bool interrupted) { + if (!interrupted) { + // to differentiate between "didn't run" and "ran + // before command's `end()` + second += 1 + first; + } + }); + + scheduler.Schedule(command); + EXPECT_EQ(0, first); + EXPECT_EQ(0, second); + scheduler.Run(); + EXPECT_EQ(1, first); + // if `second == 0`, neither of the lambdas ran. + // if `second == 1`, the second lambda ran before the first one + EXPECT_EQ(2, second); +} + +// handleInterruptTest() implicitly tests the interrupt=true branch of +// finallyDo() +TEST_F(CommandDecoratorTest, HandleInterrupt) { + CommandScheduler scheduler = GetScheduler(); + int first = 0; + int second = 0; + CommandPtr command = FunctionalCommand([] {}, [] {}, + [&first](bool interrupted) { + if (interrupted) { + first++; + } + }, + [] { return false; }) + .HandleInterrupt([&first, &second] { + // to differentiate between "didn't run" and "ran + // before command's `end()` + second += 1 + first; + }); + + scheduler.Schedule(command); + scheduler.Run(); + EXPECT_EQ(0, first); + EXPECT_EQ(0, second); + + scheduler.Cancel(command); + // if `second == 0`, neither of the lambdas ran. + // if `second == 1`, the second lambda ran before the first one + EXPECT_EQ(2, second); +}