[commands] SelectCommand: Fix leakage and multiple composition bug (#…

…5571)
wpilibsuite · Aug 30, 2023 · cf19102 · cf19102
1 parent 171375f
commit cf19102
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 8 deletions.
diff --git a/wpilibNewCommands/src/main/java/edu/wpi/first/wpilibj2/command/SelectCommand.java b/wpilibNewCommands/src/main/java/edu/wpi/first/wpilibj2/command/SelectCommand.java
@@ -27,6 +27,9 @@ public class SelectCommand extends Command {
   private boolean m_runsWhenDisabled = true;
   private InterruptionBehavior m_interruptBehavior = InterruptionBehavior.kCancelIncoming;
 
+  private final Command m_defaultCommand =
+      new PrintCommand("SelectCommand selector value does not correspond to any command!");
+
   /**
    * Creates a new SelectCommand.
    *
@@ -37,6 +40,7 @@ public SelectCommand(Map<Object, Command> commands, Supplier<Object> selector) {
     m_commands = requireNonNullParam(commands, "commands", "SelectCommand");
     m_selector = requireNonNullParam(selector, "selector", "SelectCommand");
 
+    CommandScheduler.getInstance().registerComposedCommands(m_defaultCommand);
     CommandScheduler.getInstance()
         .registerComposedCommands(commands.values().toArray(new Command[] {}));
 
@@ -51,12 +55,7 @@ public SelectCommand(Map<Object, Command> commands, Supplier<Object> selector) {
 
   @Override
   public void initialize() {
-    if (!m_commands.containsKey(m_selector.get())) {
-      m_selectedCommand =
-          new PrintCommand("SelectCommand selector value does not correspond to any command!");
-    } else {
-      m_selectedCommand = m_commands.get(m_selector.get());
-    }
+    m_selectedCommand = m_commands.getOrDefault(m_selector.get(), m_defaultCommand);
     m_selectedCommand.initialize();
   }
 

diff --git a/wpilibNewCommands/src/main/native/include/frc2/command/SelectCommand.h b/wpilibNewCommands/src/main/native/include/frc2/command/SelectCommand.h
@@ -54,8 +54,10 @@ class SelectCommand : public CommandHelper<Command, SelectCommand<Key>> {
          std::make_unique<std::decay_t<Commands>>(std::move(commands.second))),
      ...);
 
+    m_defaultCommand.SetComposed(true);
     for (auto&& command : foo) {
       CommandScheduler::GetInstance().RequireUngrouped(command.second.get());
+      command.second.get()->SetComposed(true);
     }
 
     for (auto&& command : foo) {
@@ -73,8 +75,10 @@ class SelectCommand : public CommandHelper<Command, SelectCommand<Key>> {
       std::function<Key()> selector,
       std::vector<std::pair<Key, std::unique_ptr<Command>>>&& commands)
       : m_selector{std::move(selector)} {
+    m_defaultCommand.SetComposed(true);
     for (auto&& command : commands) {
       CommandScheduler::GetInstance().RequireUngrouped(command.second.get());
+      command.second.get()->SetComposed(true);
     }
 
     for (auto&& command : commands) {
@@ -139,14 +143,16 @@ class SelectCommand : public CommandHelper<Command, SelectCommand<Key>> {
   bool m_runsWhenDisabled = true;
   Command::InterruptionBehavior m_interruptBehavior{
       Command::InterruptionBehavior::kCancelIncoming};
+
+  PrintCommand m_defaultCommand{
+      "SelectCommand selector value does not correspond to any command!"};
 };
 
 template <typename T>
 void SelectCommand<T>::Initialize() {
   auto find = m_commands.find(m_selector());
   if (find == m_commands.end()) {
-    m_selectedCommand = new PrintCommand(
-        "SelectCommand selector value does not correspond to any command!");
+    m_selectedCommand = &m_defaultCommand;
   } else {
     m_selectedCommand = find->second.get();
   }