[commands] Add interruptor parameter to onCommandInterrupt callbacks (#5461)

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

@@ -48,7 +48,7 @@ class CommandScheduler::Impl {
   // every command.
   wpi::SmallVector<Action, 4> initActions;
   wpi::SmallVector<Action, 4> executeActions;
-  wpi::SmallVector<Action, 4> interruptActions;
+  wpi::SmallVector<InterruptAction, 4> interruptActions;
   wpi::SmallVector<Action, 4> finishActions;
 
   // Flag and queues for avoiding concurrent modification if commands are
@@ -56,7 +56,8 @@ class CommandScheduler::Impl {
 
   bool inRunLoop = false;
   wpi::SmallVector<Command*, 4> toSchedule;
-  wpi::SmallVector<Command*, 4> toCancel;
+  wpi::SmallVector<Command*, 4> toCancelCommands;
+  wpi::SmallVector<std::optional<Command*>, 4> toCancelInterruptors;
 };
 
 template <typename TMap, typename TKey>
@@ -138,7 +139,7 @@ void CommandScheduler::Schedule(Command* command) {
   if (isDisjoint || allInterruptible) {
     if (allInterruptible) {
       for (auto&& cmdToCancel : intersection) {
-        Cancel(cmdToCancel);
+        Cancel(cmdToCancel, std::make_optional(command));
       }
     }
     m_impl->scheduledCommands.insert(command);
@@ -196,7 +197,7 @@ void CommandScheduler::Run() {
   // Run scheduled commands, remove finished commands.
   for (Command* command : m_impl->scheduledCommands) {
     if (!command->RunsWhenDisabled() && frc::RobotState::IsDisabled()) {
-      Cancel(command);
+      Cancel(command, std::nullopt);
       continue;
     }
 
@@ -226,12 +227,13 @@ void CommandScheduler::Run() {
     Schedule(command);
   }
 
-  for (auto&& command : m_impl->toCancel) {
-    Cancel(command);
+  for (size_t i = 0; i < m_impl->toCancelCommands.size(); i++) {
+    Cancel(m_impl->toCancelCommands[i], m_impl->toCancelInterruptors[i]);
   }
 
   m_impl->toSchedule.clear();
-  m_impl->toCancel.clear();
+  m_impl->toCancelCommands.clear();
+  m_impl->toCancelInterruptors.clear();
 
   // Add default commands for un-required registered subsystems.
   for (auto&& subsystem : m_impl->subsystems) {
@@ -319,13 +321,15 @@ Command* CommandScheduler::GetDefaultCommand(const Subsystem* subsystem) const {
   }
 }
 
-void CommandScheduler::Cancel(Command* command) {
+void CommandScheduler::Cancel(Command* command,
+                              std::optional<Command*> interruptor) {
   if (!m_impl) {
     return;
   }
 
   if (m_impl->inRunLoop) {
-    m_impl->toCancel.emplace_back(command);
+    m_impl->toCancelCommands.emplace_back(command);
+    m_impl->toCancelInterruptors.emplace_back(interruptor);
     return;
   }
 
@@ -341,11 +345,15 @@ void CommandScheduler::Cancel(Command* command) {
   }
   command->End(true);
   for (auto&& action : m_impl->interruptActions) {
-    action(*command);
+    action(*command, interruptor);
   }
   m_watchdog.AddEpoch(command->GetName() + ".End(true)");
 }
 
+void CommandScheduler::Cancel(Command* command) {
+  Cancel(command, std::nullopt);
+}
+
 void CommandScheduler::Cancel(const CommandPtr& command) {
   Cancel(command.get());
 }
@@ -424,6 +432,14 @@ void CommandScheduler::OnCommandExecute(Action action) {
 }
 
 void CommandScheduler::OnCommandInterrupt(Action action) {
+  m_impl->interruptActions.emplace_back(
+      [action = std::move(action)](const Command& command,
+                                   const std::optional<Command*>& interruptor) {
+        action(command);
+      });
+}
+
+void CommandScheduler::OnCommandInterrupt(InterruptAction action) {
   m_impl->interruptActions.emplace_back(std::move(action));
 }
 

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

@@ -8,6 +8,7 @@
 #include <functional>
 #include <initializer_list>
 #include <memory>
+#include <optional>
 #include <span>
 #include <utility>
 
@@ -48,6 +49,8 @@ class CommandScheduler final : public wpi::Sendable,
   CommandScheduler& operator=(const CommandScheduler&) = delete;
 
   using Action = std::function<void(const Command&)>;
+  using InterruptAction =
+      std::function<void(const Command&, const std::optional<Command*>&)>;
 
   /**
    * Changes the period of the loop overrun watchdog. This should be kept in
@@ -353,6 +356,16 @@ class CommandScheduler final : public wpi::Sendable,
    */
   void OnCommandInterrupt(Action action);
 
+  /**
+   * Adds an action to perform on the interruption of any command by the
+   * scheduler. The action receives the interrupted command and an optional
+   * containing the interrupting command, or nullopt if it was not canceled by a
+   * command (e.g., by Cancel()).
+   *
+   * @param action the action to perform
+   */
+  void OnCommandInterrupt(InterruptAction action);
+
   /**
    * Adds an action to perform on the finishing of any command by the scheduler.
    *
@@ -397,6 +410,8 @@ class CommandScheduler final : public wpi::Sendable,
   void SetDefaultCommandImpl(Subsystem* subsystem,
                              std::unique_ptr<Command> command);
 
+  void Cancel(Command* command, std::optional<Command*> interruptor);
+
   class Impl;
   std::unique_ptr<Impl> m_impl;