From ec124bb662b893e6545e85ff650ca2ceb4aa080f Mon Sep 17 00:00:00 2001
From: Colin Wong <cwongmath@gmail.com>
Date: Mon, 28 Nov 2022 16:03:14 -0600
Subject: [PATCH] [commands] Allow unsetting a subsystem's default command
 (#4621)

---
 .../wpilibj2/command/CommandScheduler.java    | 20 +++++++++++++++++++
 .../cpp/frc2/command/CommandScheduler.cpp     |  9 +++++++++
 .../include/frc2/command/CommandScheduler.h   |  9 +++++++++
 3 files changed, 38 insertions(+)

diff --git a/wpilibNewCommands/src/main/java/edu/wpi/first/wpilibj2/command/CommandScheduler.java b/wpilibNewCommands/src/main/java/edu/wpi/first/wpilibj2/command/CommandScheduler.java
index 84850a03b4..4fb89ac361 100644
--- a/wpilibNewCommands/src/main/java/edu/wpi/first/wpilibj2/command/CommandScheduler.java
+++ b/wpilibNewCommands/src/main/java/edu/wpi/first/wpilibj2/command/CommandScheduler.java
@@ -367,6 +367,10 @@ public final class CommandScheduler implements NTSendable, AutoCloseable {
         DriverStation.reportWarning("Tried to register a null subsystem", true);
         continue;
       }
+      if (m_subsystems.containsKey(subsystem)) {
+        DriverStation.reportWarning("Tried to register an already-registered subsystem", true);
+        continue;
+      }
       m_subsystems.put(subsystem, null);
     }
   }
@@ -416,6 +420,22 @@ public final class CommandScheduler implements NTSendable, AutoCloseable {
     m_subsystems.put(subsystem, defaultCommand);
   }
 
+  /**
+   * Removes the default command for a subsystem. The current default command will run until another
+   * command is scheduled that requires the subsystem, at which point the current default command
+   * will not be re-scheduled.
+   *
+   * @param subsystem the subsystem whose default command will be removed
+   */
+  public void removeDefaultCommand(Subsystem subsystem) {
+    if (subsystem == null) {
+      DriverStation.reportWarning("Tried to remove a default command for a null subsystem", true);
+      return;
+    }
+
+    m_subsystems.put(subsystem, null);
+  }
+
   /**
    * Gets the default command associated with this subsystem. Null if this subsystem has no default
    * command associated with it.
diff --git a/wpilibNewCommands/src/main/native/cpp/frc2/command/CommandScheduler.cpp b/wpilibNewCommands/src/main/native/cpp/frc2/command/CommandScheduler.cpp
index 5f702efefd..f096e493bf 100644
--- a/wpilibNewCommands/src/main/native/cpp/frc2/command/CommandScheduler.cpp
+++ b/wpilibNewCommands/src/main/native/cpp/frc2/command/CommandScheduler.cpp
@@ -258,6 +258,11 @@ void CommandScheduler::Run() {
 }
 
 void CommandScheduler::RegisterSubsystem(Subsystem* subsystem) {
+  if (m_impl->subsystems.find(subsystem) != m_impl->subsystems.end()) {
+    std::puts("Tried to register an already-registered subsystem");
+    return;
+  }
+
   m_impl->subsystems[subsystem] = nullptr;
 }
 
@@ -306,6 +311,10 @@ void CommandScheduler::SetDefaultCommand(Subsystem* subsystem,
   SetDefaultCommandImpl(subsystem, std::move(defaultCommand).Unwrap());
 }
 
+void CommandScheduler::RemoveDefaultCommand(Subsystem* subsystem) {
+  m_impl->subsystems[subsystem] = nullptr;
+}
+
 Command* CommandScheduler::GetDefaultCommand(const Subsystem* subsystem) const {
   auto&& find = m_impl->subsystems.find(subsystem);
   if (find != m_impl->subsystems.end()) {
diff --git a/wpilibNewCommands/src/main/native/include/frc2/command/CommandScheduler.h b/wpilibNewCommands/src/main/native/include/frc2/command/CommandScheduler.h
index edce241d08..2e0b09b0d8 100644
--- a/wpilibNewCommands/src/main/native/include/frc2/command/CommandScheduler.h
+++ b/wpilibNewCommands/src/main/native/include/frc2/command/CommandScheduler.h
@@ -200,6 +200,15 @@ class CommandScheduler final : public nt::NTSendable,
    */
   void SetDefaultCommand(Subsystem* subsystem, CommandPtr&& defaultCommand);
 
+  /**
+   * Removes the default command for a subsystem. The current default command
+   * will run until another command is scheduled that requires the subsystem, at
+   * which point the current default command will not be re-scheduled.
+   *
+   * @param subsystem the subsystem whose default command will be removed
+   */
+  void RemoveDefaultCommand(Subsystem* subsystem);
+
   /**
    * Gets the default command associated with this subsystem.  Null if this
    * subsystem has no default command associated with it.