From 9ae38eaa7c72cd65055ba5d204545083dac251fc Mon Sep 17 00:00:00 2001
From: Starlight220 <53231611+Starlight220@users.noreply.github.com>
Date: Wed, 14 Sep 2022 04:33:19 +0300
Subject: [PATCH] [commands] Add owning overload to ProxyScheduleCommand
 (#4405)

---
 .../cpp/frc2/command/ProxyScheduleCommand.cpp |  7 ++++
 .../frc2/command/ProxyScheduleCommand.h       | 16 +++++++-
 .../frc2/command/ProxyScheduleCommandTest.cpp | 38 +++++++++++++++++--
 3 files changed, 56 insertions(+), 5 deletions(-)

diff --git a/wpilibNewCommands/src/main/native/cpp/frc2/command/ProxyScheduleCommand.cpp b/wpilibNewCommands/src/main/native/cpp/frc2/command/ProxyScheduleCommand.cpp
index 116ce3c112..8935b5ca67 100644
--- a/wpilibNewCommands/src/main/native/cpp/frc2/command/ProxyScheduleCommand.cpp
+++ b/wpilibNewCommands/src/main/native/cpp/frc2/command/ProxyScheduleCommand.cpp
@@ -15,6 +15,13 @@ ProxyScheduleCommand::ProxyScheduleCommand(Command* toSchedule) {
   SetInsert(m_toSchedule, {&toSchedule, 1});
 }
 
+ProxyScheduleCommand::ProxyScheduleCommand(
+    std::unique_ptr<Command>&& toSchedule)
+    : m_owning(std::move(toSchedule)) {
+  Command* ptr = m_owning.get();
+  SetInsert(m_toSchedule, {&ptr, 1});
+}
+
 void ProxyScheduleCommand::Initialize() {
   for (auto* command : m_toSchedule) {
     command->Schedule();
diff --git a/wpilibNewCommands/src/main/native/include/frc2/command/ProxyScheduleCommand.h b/wpilibNewCommands/src/main/native/include/frc2/command/ProxyScheduleCommand.h
index afc6f542e8..543e37eb2c 100644
--- a/wpilibNewCommands/src/main/native/include/frc2/command/ProxyScheduleCommand.h
+++ b/wpilibNewCommands/src/main/native/include/frc2/command/ProxyScheduleCommand.h
@@ -4,6 +4,8 @@
 
 #pragma once
 
+#include <memory>
+
 #include <wpi/SmallVector.h>
 #include <wpi/span.h>
 
@@ -33,9 +35,18 @@ class ProxyScheduleCommand
 
   explicit ProxyScheduleCommand(Command* toSchedule);
 
-  ProxyScheduleCommand(ProxyScheduleCommand&& other) = default;
+  /**
+   * Creates a new ProxyScheduleCommand that schedules the given command when
+   * initialized, and ends when it is no longer scheduled.
+   *
+   * <p>Note that this constructor passes ownership of the given command to the
+   * returned ProxyScheduleCommand.
+   *
+   * @param toSchedule the command to schedule
+   */
+  explicit ProxyScheduleCommand(std::unique_ptr<Command>&& toSchedule);
 
-  ProxyScheduleCommand(const ProxyScheduleCommand& other) = default;
+  ProxyScheduleCommand(ProxyScheduleCommand&& other) = default;
 
   void Initialize() override;
 
@@ -47,6 +58,7 @@ class ProxyScheduleCommand
 
  private:
   wpi::SmallVector<Command*, 4> m_toSchedule;
+  std::unique_ptr<Command> m_owning;
   bool m_finished{false};
 };
 }  // namespace frc2
diff --git a/wpilibNewCommands/src/test/native/cpp/frc2/command/ProxyScheduleCommandTest.cpp b/wpilibNewCommands/src/test/native/cpp/frc2/command/ProxyScheduleCommandTest.cpp
index def7be9ff6..a932eaa537 100644
--- a/wpilibNewCommands/src/test/native/cpp/frc2/command/ProxyScheduleCommandTest.cpp
+++ b/wpilibNewCommands/src/test/native/cpp/frc2/command/ProxyScheduleCommandTest.cpp
@@ -2,7 +2,7 @@
 // Open Source Software; you can modify and/or share it under the terms of
 // the WPILib BSD license file in the root directory of this project.
 
-#include <regex>
+#include <memory>
 
 #include "CommandTestBase.h"
 #include "frc2/command/InstantCommand.h"
@@ -12,7 +12,7 @@
 using namespace frc2;
 class ProxyScheduleCommandTest : public CommandTestBase {};
 
-TEST_F(ProxyScheduleCommandTest, ProxyScheduleCommandSchedule) {
+TEST_F(ProxyScheduleCommandTest, NonOwningCommandSchedule) {
   CommandScheduler& scheduler = CommandScheduler::GetInstance();
 
   bool scheduled = false;
@@ -27,7 +27,7 @@ TEST_F(ProxyScheduleCommandTest, ProxyScheduleCommandSchedule) {
   EXPECT_TRUE(scheduled);
 }
 
-TEST_F(ProxyScheduleCommandTest, ProxyScheduleCommandEnd) {
+TEST_F(ProxyScheduleCommandTest, NonOwningCommandEnd) {
   CommandScheduler& scheduler = CommandScheduler::GetInstance();
 
   bool finished = false;
@@ -45,3 +45,35 @@ TEST_F(ProxyScheduleCommandTest, ProxyScheduleCommandEnd) {
   scheduler.Run();
   EXPECT_FALSE(scheduler.IsScheduled(&command));
 }
+
+TEST_F(ProxyScheduleCommandTest, OwningCommandSchedule) {
+  CommandScheduler& scheduler = CommandScheduler::GetInstance();
+
+  bool scheduled = false;
+
+  ProxyScheduleCommand command(
+      std::make_unique<InstantCommand>([&scheduled] { scheduled = true; }));
+
+  scheduler.Schedule(&command);
+  scheduler.Run();
+
+  EXPECT_TRUE(scheduled);
+}
+
+TEST_F(ProxyScheduleCommandTest, OwningCommandEnd) {
+  CommandScheduler& scheduler = CommandScheduler::GetInstance();
+
+  bool finished = false;
+
+  ProxyScheduleCommand command(
+      std::make_unique<WaitUntilCommand>([&finished] { return finished; }));
+
+  scheduler.Schedule(&command);
+  scheduler.Run();
+
+  EXPECT_TRUE(scheduler.IsScheduled(&command));
+  finished = true;
+  scheduler.Run();
+  scheduler.Run();
+  EXPECT_FALSE(scheduler.IsScheduled(&command));
+}