[cmd] Add Command.repeatedly(int count)

wpilibNewCommands/src/main/java/edu/wpi/first/wpilibj2/command/Command.java

@@ -401,6 +401,26 @@ public RepeatCommand repeatedly() {
     return new RepeatCommand(this);
   }
 
+  /**
+   * Decorates this command to run repeatedly, restarting it when it ends, until this command is run
+   * the specified number of times or is interrupted. The decorated command can still be canceled.
+   *
+   * <p>Note: This decorator works by adding this command to a composition. The command the
+   * decorator was called on cannot be scheduled independently or be added to a different
+   * composition (namely, decorators), unless it is manually cleared from the list of composed
+   * commands with {@link CommandScheduler#removeComposedCommand(Command)}. The command composition
+   * returned from this method can be further decorated without issue.
+   *
+   * @param repetitions the number of times to run the command.
+   * @return the decorated command
+   */
+  public ParallelRaceGroup repeatedly(int repetitions) {
+    // Use an array so that it stays in the heap instead of stack.
+    // We use an array with a size of 1 instead of `AtomicInteger` because of performance difference
+    int[] counter = {0};
+    return this.finallyDo(() -> counter[0]++).repeatedly().until(() -> counter[0] >= repetitions);
+  }
+
   /**
    * Decorates this command to run "by proxy" by wrapping it in a {@link ProxyCommand}. Use this for
    * "forking off" from command compositions when the user does not wish to extend the command's

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

@@ -109,6 +109,10 @@ CommandPtr Command::Repeatedly() && {
   return std::move(*this).ToPtr().Repeatedly();
 }
 
+CommandPtr Command::Repeatedly(int times) && {
+  return std::move(*this).ToPtr().Repeatedly(times);
+}
+
 CommandPtr Command::AsProxy() && {
   return std::move(*this).ToPtr().AsProxy();
 }

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

@@ -50,6 +50,15 @@ CommandPtr CommandPtr::Repeatedly() && {
   return std::move(*this);
 }
 
+CommandPtr CommandPtr::Repeatedly(int times) && {
+  AssertValid();
+  std::shared_ptr<int> countPtr = std::make_shared<int>(0);
+  return std::move(*this)
+      .FinallyDo([countPtr] { (*countPtr)++; })
+      .Repeatedly()
+      .Until([countPtr, times] { return ((*countPtr) >= times); });
+}
+
 CommandPtr CommandPtr::AsProxy() && {
   AssertValid();
   m_ptr = std::make_unique<ProxyCommand>(std::move(m_ptr));

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

@@ -269,6 +269,17 @@ class Command : public wpi::Sendable, public wpi::SendableHelper<Command> {
   [[nodiscard]]
   CommandPtr Repeatedly() &&;
 
+  /**
+   * Decorates this command to run repeatedly, restarting it when it ends,
+   * until this command is run the specified number of times or is interrupted.
+   * The decorated command can still be canceled.
+   *
+   * @param repetitions the number of times to run the command
+   * @return the decorated command
+   */
+  [[nodiscard]]
+  CommandPtr Repeatedly(int repetitions) &&;
+
   /**
    * Decorates this command to run "by proxy" by wrapping it in a ProxyCommand.
    * Use this for "forking off" from command compositions when the user does not

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

@@ -49,6 +49,16 @@ class CommandPtr final {
   [[nodiscard]]
   CommandPtr Repeatedly() &&;
 
+  /**
+   * Decorates this command to run repeatedly, restarting until the command runs
+   * for the given number of times. The decorated command can still be canceled.
+   *
+   * @param times the number of times to run the command
+   * @return the decorated command
+   */
+  [[nodiscard]]
+  CommandPtr Repeatedly(int times) &&;
+
   /**
    * Decorates this command to run "by proxy" by wrapping it in a ProxyCommand.
    * Use this for "forking off" from command compositions when the user does not

wpilibNewCommands/src/test/java/edu/wpi/first/wpilibj2/command/CommandDecoratorTest.java

@@ -420,6 +420,48 @@ void raceWithOrderTest() {
     }
   }
 
+  @Test
+  void repeatedlyTest() {
+    try (CommandScheduler scheduler = new CommandScheduler()) {
+      AtomicInteger counter = new AtomicInteger(0);
+
+      Command command = new InstantCommand(counter::incrementAndGet);
+
+      Command group = command.repeatedly();
+
+      scheduler.schedule(group);
+      for (int i = 1; i <= 50; i++) {
+        scheduler.run();
+        assertEquals(i, counter.get());
+      }
+
+      // Should still be scheduled
+      assertTrue(scheduler.isScheduled(group));
+    }
+  }
+
+  @Test
+  void repeatedlyCountTest() {
+    try (CommandScheduler scheduler = new CommandScheduler()) {
+      AtomicInteger counter = new AtomicInteger(0);
+
+      Command command = new InstantCommand(counter::incrementAndGet);
+
+      Command group = command.repeatedly(3);
+
+      scheduler.schedule(group);
+      assertEquals(1, counter.get());
+      for (int i = 1; i < 3; i++) {
+        scheduler.run();
+        assertEquals(i, counter.get());
+        assertTrue(scheduler.isScheduled(group), "Expected group to be scheduled with i = " + i);
+      }
+      scheduler.run();
+      assertEquals(3, counter.get(), "Loop should have run 3 times something went wrong");
+      assertFalse(scheduler.isScheduled(group), "This command should have gotten unscheduled");
+    }
+  }
+
   @Test
   void unlessTest() {
     try (CommandScheduler scheduler = new CommandScheduler()) {

wpilibNewCommands/src/test/native/cpp/frc2/command/CommandDecoratorTest.cpp

@@ -383,6 +383,45 @@ TEST_F(CommandDecoratorTest, RaceWithOrder) {
   EXPECT_TRUE(firstWasPolled);
 }
 
+TEST_F(CommandDecoratorTest, Repeatedly) {
+  CommandScheduler scheduler = GetScheduler();
+
+  int counter = 0;
+
+  auto command = InstantCommand([&counter] { counter++; }, {}).Repeatedly();
+
+  scheduler.Schedule(command);
+
+  for (int i = 1; i <= 50; i++) {
+    scheduler.Run();
+    EXPECT_EQ(i, counter);
+  }
+
+  EXPECT_TRUE(scheduler.IsScheduled(command));
+}
+
+TEST_F(CommandDecoratorTest, RepeatedlyCount) {
+  CommandScheduler scheduler = GetScheduler();
+
+  int counter = 0;
+
+  auto command = InstantCommand([&counter] { counter++; }, {}).Repeatedly(3);
+
+  scheduler.Schedule(command);
+  EXPECT_EQ(1, counter);
+  for (int i = 1; i < 3; i++) {
+    scheduler.Run();
+    EXPECT_EQ(i, counter);
+    EXPECT_TRUE(scheduler.IsScheduled(command))
+        << "Expected group to be scheduled with i = " << i;
+  }
+
+  scheduler.Run();
+  EXPECT_EQ(3, counter) << "Loop should have run 3 times something went wrong";
+  EXPECT_FALSE(scheduler.IsScheduled(command))
+      << "This command should have gotten unscheduled";
+}
+
 TEST_F(CommandDecoratorTest, Unless) {
   CommandScheduler scheduler = GetScheduler();