[wpilibj, wpilibc] Fix LED patterns not offsetting reads (#6948)

Was causing bugs when combined with patterns that need to read back from the buffer (eg masks and overlays) Co-authored-by: Joseph Eng <s-engjo@bsd405.org>
wpilibsuite · Nov 29, 2024 · 5e1c6a8 · 5e1c6a8
1 parent a0af0fd
commit 5e1c6a8
Show file tree

Hide file tree

Showing 5 changed files with 529 additions and 89 deletions.
diff --git a/wpilibc/src/main/native/cpp/LEDPattern.cpp b/wpilibc/src/main/native/cpp/LEDPattern.cpp
@@ -18,33 +18,46 @@
 
 using namespace frc;
 
-LEDPattern::LEDPattern(LEDPatternFn impl) : m_impl(std::move(impl)) {}
+LEDPattern::LEDPattern(std::function<void(frc::LEDPattern::LEDReader,
+                                          std::function<void(int, frc::Color)>)>
+                           impl)
+    : m_impl(std::move(impl)) {}
+
+void LEDPattern::ApplyTo(LEDPattern::LEDReader reader,
+                         std::function<void(int, frc::Color)> writer) const {
+  m_impl(reader, writer);
+}
 
 void LEDPattern::ApplyTo(std::span<AddressableLED::LEDData> data,
-                         LEDWriterFn writer) const {
-  m_impl(data, writer);
+                         std::function<void(int, frc::Color)> writer) const {
+  ApplyTo(LEDPattern::LEDReader{[=](size_t i) { return data[i]; }, data.size()},
+          writer);
 }
 
 void LEDPattern::ApplyTo(std::span<AddressableLED::LEDData> data) const {
   ApplyTo(data, [&](int index, Color color) { data[index].SetLED(color); });
 }
 
-LEDPattern LEDPattern::Reversed() {
-  return LEDPattern{[self = *this](auto data, auto writer) {
-    self.ApplyTo(data, [&](int i, Color color) {
-      writer((data.size() - 1) - i, color);
-    });
+LEDPattern LEDPattern::MapIndex(
+    std::function<size_t(size_t, size_t)> indexMapper) {
+  return LEDPattern{[self = *this, indexMapper](auto data, auto writer) {
+    size_t bufLen = data.size();
+    self.ApplyTo(
+        LEDPattern::LEDReader{
+            [=](auto i) { return data[indexMapper(bufLen, i)]; }, bufLen},
+        [&](int i, Color color) { writer(indexMapper(bufLen, i), color); });
   }};
 }
 
+LEDPattern LEDPattern::Reversed() {
+  return MapIndex([](size_t bufLen, size_t i) { return bufLen - 1 - i; });
+}
+
 LEDPattern LEDPattern::OffsetBy(int offset) {
-  return LEDPattern{[=, self = *this](auto data, auto writer) {
-    self.ApplyTo(data, [&data, &writer, offset](int i, Color color) {
-      int shiftedIndex =
-          frc::FloorMod(i + offset, static_cast<int>(data.size()));
-      writer(shiftedIndex, color);
-    });
-  }};
+  return MapIndex([offset](size_t bufLen, size_t i) {
+    return frc::FloorMod(static_cast<int>(i) + offset,
+                         static_cast<int>(bufLen));
+  });
 }
 
 LEDPattern LEDPattern::ScrollAtRelativeSpeed(units::hertz_t velocity) {
@@ -53,21 +66,17 @@ LEDPattern LEDPattern::ScrollAtRelativeSpeed(units::hertz_t velocity) {
   // Invert and multiply by 1,000,000 to get microseconds
   double periodMicros = 1e6 / velocity.value();
 
-  return LEDPattern{[=, self = *this](auto data, auto writer) {
-    auto bufLen = data.size();
+  return MapIndex([=](size_t bufLen, size_t i) {
     auto now = wpi::Now();
 
     // index should move by (bufLen) / (period)
     double t =
         (now % static_cast<int64_t>(std::floor(periodMicros))) / periodMicros;
     int offset = static_cast<int>(std::floor(t * bufLen));
 
-    self.ApplyTo(data, [=](int i, Color color) {
-      // floorMod so if the offset is negative, we still get positive outputs
-      int shiftedIndex = frc::FloorMod(i + offset, static_cast<int>(bufLen));
-      writer(shiftedIndex, color);
-    });
-  }};
+    return frc::FloorMod(static_cast<int>(i) + offset,
+                         static_cast<int>(bufLen));
+  });
 }
 
 LEDPattern LEDPattern::ScrollAtAbsoluteSpeed(
@@ -77,8 +86,7 @@ LEDPattern LEDPattern::ScrollAtAbsoluteSpeed(
   auto microsPerLed =
       static_cast<int64_t>(std::floor((ledSpacing / velocity).value() * 1e6));
 
-  return LEDPattern{[=, self = *this](auto data, auto writer) {
-    auto bufLen = data.size();
+  return MapIndex([=](size_t bufLen, size_t i) {
     auto now = wpi::Now();
 
     // every step in time that's a multiple of microsPerLED will increment
@@ -87,13 +95,9 @@ LEDPattern LEDPattern::ScrollAtAbsoluteSpeed(
     // offset values for negative velocities
     auto offset = static_cast<int64_t>(now) / microsPerLed;
 
-    self.ApplyTo(data, [=, &writer](int i, Color color) {
-      // FloorMod so if the offset is negative, we still get positive outputs
-      int shiftedIndex = frc::FloorMod(i + offset, static_cast<int>(bufLen));
-
-      writer(shiftedIndex, color);
-    });
-  }};
+    return frc::FloorMod(static_cast<int>(i) + offset,
+                         static_cast<int>(bufLen));
+  });
 }
 
 LEDPattern LEDPattern::Blink(units::second_t onTime, units::second_t offTime) {

diff --git a/wpilibc/src/main/native/include/frc/LEDPattern.h b/wpilibc/src/main/native/include/frc/LEDPattern.h
@@ -18,21 +18,36 @@
 
 namespace frc {
 
-/**
- * Sets the LED at the given index to the given color.
- */
-using LEDWriterFn = std::function<void(int, frc::Color)>;
-
-/**
- * Accepts a data buffer (1st argument) and a callback (2nd argument) for
- * writing data.
- */
-using LEDPatternFn =
-    std::function<void(std::span<frc::AddressableLED::LEDData>, LEDWriterFn)>;
-
 class LEDPattern {
  public:
-  explicit LEDPattern(LEDPatternFn impl);
+  /**
+   * A wrapper around a length and an arbitrary reader function that accepts an
+   * LED index and returns data for the LED at that index. This configuration
+   * allows us to abstract over different container types without templating.
+   */
+  class LEDReader {
+   public:
+    LEDReader(std::function<frc::AddressableLED::LEDData(int)> impl,
+              size_t size)
+        : m_impl{std::move(impl)}, m_size{size} {}
+
+    frc::AddressableLED::LEDData operator[](size_t index) const {
+      return m_impl(index);
+    }
+
+    size_t size() const { return m_size; }
+
+   private:
+    std::function<frc::AddressableLED::LEDData(int)> m_impl;
+    size_t m_size;
+  };
+
+  explicit LEDPattern(std::function<void(frc::LEDPattern::LEDReader,
+                                         std::function<void(int, frc::Color)>)>
+                          impl);
+
+  void ApplyTo(LEDReader reader,
+               std::function<void(int, frc::Color)> writer) const;
 
   /**
    * Writes the pattern to an LED buffer. Dynamic animations should be called
@@ -48,7 +63,7 @@ class LEDPattern {
    * @param writer data writer for setting new LED colors on the LED strip
    */
   void ApplyTo(std::span<frc::AddressableLED::LEDData> data,
-               LEDWriterFn writer) const;
+               std::function<void(int, frc::Color)> writer) const;
 
   /**
    * Writes the pattern to an LED buffer. Dynamic animations should be called
@@ -64,6 +79,15 @@ class LEDPattern {
    */
   void ApplyTo(std::span<frc::AddressableLED::LEDData> data) const;
 
+  /**
+   * Creates a pattern with remapped indices.
+   *
+   * @param indexMapper the index mapper
+   * @return the mapped pattern
+   */
+  [[nodiscard]]
+  LEDPattern MapIndex(std::function<size_t(size_t, size_t)> indexMapper);
+
   /**
    * Creates a pattern that displays this one in reverse. Scrolling patterns
    * will scroll in the opposite direction (but at the same speed). It will
@@ -373,6 +397,8 @@ class LEDPattern {
   static LEDPattern Rainbow(int saturation, int value);
 
  private:
-  LEDPatternFn m_impl;
+  std::function<void(frc::LEDPattern::LEDReader,
+                     std::function<void(int, frc::Color)>)>
+      m_impl;
 };
 }  // namespace frc