From db3e4bfdc8337b582dd780dc72074dc9d8cd5787 Mon Sep 17 00:00:00 2001
From: Ryan Blue <ryanzblue@gmail.com>
Date: Tue, 15 Oct 2024 16:11:36 -0400
Subject: [PATCH] Allow CommandPtrs to be owned by the scheduler

Previously users would have to keep track of dynamically created CommandPtrs. This adds an ownership-taking version of schedule which places the command in a temporary store in the scheduler. The command will be freed when the command's lifecycle ends.
---
 .../cpp/frc2/command/CommandScheduler.cpp     | 14 ++++++
 .../include/frc2/command/CommandScheduler.h   | 13 +++++
 .../native/cpp/frc2/command/SchedulerTest.cpp | 49 +++++++++++++++++++
 3 files changed, 76 insertions(+)

diff --git a/wpilibNewCommands/src/main/native/cpp/frc2/command/CommandScheduler.cpp b/wpilibNewCommands/src/main/native/cpp/frc2/command/CommandScheduler.cpp
index aa57bf00d08..6c09d0f6127 100644
--- a/wpilibNewCommands/src/main/native/cpp/frc2/command/CommandScheduler.cpp
+++ b/wpilibNewCommands/src/main/native/cpp/frc2/command/CommandScheduler.cpp
@@ -7,6 +7,7 @@
 #include <cstdio>
 #include <memory>
 #include <string>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
@@ -62,6 +63,11 @@ class CommandScheduler::Impl {
   wpi::SmallVector<Command*, 4> toCancelCommands;
   wpi::SmallVector<std::optional<Command*>, 4> toCancelInterruptors;
   wpi::SmallSet<Command*, 4> endingCommands;
+
+  // Map of Command* -> CommandPtr for CommandPtrs transferred to the scheduler
+  // via Schedule(CommandPtr&&). These are erased (destroyed) at the very end of
+  // the loop cycle when the command lifecycle is complete.
+  std::unordered_map<Command*, CommandPtr> ownedCommands;
 };
 
 template <typename TMap, typename TKey>
@@ -174,6 +180,12 @@ void CommandScheduler::Schedule(const CommandPtr& command) {
   Schedule(command.get());
 }
 
+void CommandScheduler::Schedule(CommandPtr&& command) {
+  auto ptr = command.get();
+  m_impl->ownedCommands.emplace(ptr, std::move(command));
+  Schedule(ptr);
+}
+
 void CommandScheduler::Run() {
   if (m_impl->disabled) {
     return;
@@ -226,6 +238,8 @@ void CommandScheduler::Run() {
       }
 
       m_watchdog.AddEpoch(command->GetName() + ".End(false)");
+      // remove owned commands after everything else is done
+      m_impl->ownedCommands.erase(command);
     }
   }
   m_impl->inRunLoop = false;
diff --git a/wpilibNewCommands/src/main/native/include/frc2/command/CommandScheduler.h b/wpilibNewCommands/src/main/native/include/frc2/command/CommandScheduler.h
index fde2c6d3e3f..2f43ff7bc2e 100644
--- a/wpilibNewCommands/src/main/native/include/frc2/command/CommandScheduler.h
+++ b/wpilibNewCommands/src/main/native/include/frc2/command/CommandScheduler.h
@@ -92,6 +92,17 @@ class CommandScheduler final : public wpi::Sendable,
    */
   void Schedule(const CommandPtr& command);
 
+  /**
+   * Schedules a command for execution. Does nothing if the command is already
+   * scheduled. If a command's requirements are not available, it will only be
+   * started if all the commands currently using those requirements are
+   * interruptible. If this is the case, they will be interrupted and the
+   * command will be scheduled.
+   *
+   * @param command the command to schedule
+   */
+  void Schedule(CommandPtr&& command);
+
   /**
    * Schedules a command for execution. Does nothing if the command is already
    * scheduled. If a command's requirements are not available, it will only be
@@ -99,6 +110,8 @@ class CommandScheduler final : public wpi::Sendable,
    * scheduled as interruptible. If this is the case, they will be interrupted
    * and the command will be scheduled.
    *
+   * The pointer must remain valid through the entire lifecycle of the command.
+   *
    * @param command the command to schedule
    */
   void Schedule(Command* command);
diff --git a/wpilibNewCommands/src/test/native/cpp/frc2/command/SchedulerTest.cpp b/wpilibNewCommands/src/test/native/cpp/frc2/command/SchedulerTest.cpp
index ef97bb07692..538d8df40c5 100644
--- a/wpilibNewCommands/src/test/native/cpp/frc2/command/SchedulerTest.cpp
+++ b/wpilibNewCommands/src/test/native/cpp/frc2/command/SchedulerTest.cpp
@@ -171,3 +171,52 @@ TEST_F(SchedulerTest, ScheduleScheduledNoOp) {
 
   EXPECT_EQ(counter, 1);
 }
+
+class TrackDestroyCommand
+    : public frc2::CommandHelper<Command, TrackDestroyCommand> {
+ public:
+  explicit TrackDestroyCommand(wpi::unique_function<void()> deleteFunc)
+      : m_deleteFunc{std::move(deleteFunc)} {}
+  TrackDestroyCommand(TrackDestroyCommand&& other)
+      : m_deleteFunc{std::exchange(other.m_deleteFunc, [] {})} {}
+  TrackDestroyCommand& operator=(TrackDestroyCommand&& other) {
+    m_deleteFunc = std::exchange(other.m_deleteFunc, [] {});
+    return *this;
+  }
+  ~TrackDestroyCommand() override { m_deleteFunc(); }
+
+ private:
+  wpi::unique_function<void()> m_deleteFunc;
+};
+
+TEST_F(SchedulerTest, ScheduleCommandPtr) {
+  CommandScheduler scheduler = GetScheduler();
+  int destructionCounter = 0;
+  int runCounter = 0;
+
+  bool finish = false;
+
+  {
+    auto commandPtr =
+        TrackDestroyCommand([&destructionCounter] { destructionCounter++; })
+            .AlongWith(std::move(
+                frc2::InstantCommand([&runCounter] { runCounter++; }).ToPtr()))
+            .Until([&finish] { return finish; });
+    EXPECT_EQ(destructionCounter, 0) << "Composition should not delete command";
+
+    scheduler.Schedule(std::move(commandPtr));
+    EXPECT_EQ(destructionCounter, 0)
+        << "Scheduling should not delete CommandPtr";
+  }
+  EXPECT_EQ(destructionCounter, 0)
+      << "Scheduler should own CommandPtr after scheduling";
+  scheduler.Run();
+  EXPECT_EQ(runCounter, 1);
+  EXPECT_EQ(destructionCounter, 0) << "Scheduler should not destroy CommandPtr "
+                                      "until command lifetime is complete";
+  finish = true;
+  scheduler.Run();
+  EXPECT_EQ(runCounter, 1);
+  EXPECT_EQ(destructionCounter, 1)
+      << "Scheduler should delete command after command completes";
+}