Use generic index mapping function

KangarooKoala · KangarooKoala · commit 55957940da71 · 2024-11-11T11:56:58.000-08:00
diff --git a/wpilibc/src/main/native/cpp/LEDPattern.cpp b/wpilibc/src/main/native/cpp/LEDPattern.cpp
@@ -38,30 +38,26 @@ 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) {
+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[data.size() - 1 - i]; },
-                              data.size()},
-        [&](int i, Color color) { writer((data.size() - 1) - i, color); });
+        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](LEDPattern::LEDReader data, auto writer) {
-    self.ApplyTo(LEDPattern::LEDReader{[=, d = std::move(data)](auto i) {
-                                         int shiftedIndex = frc::FloorMod(
-                                             i + offset,
-                                             static_cast<int>(d.size()));
-                                         return d[shiftedIndex];
-                                       },
-                                       data.size()},
-                 [&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) {
@@ -70,30 +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(LEDPattern::LEDReader{[=, d = std::move(data)](auto i) {
-                                         int shiftedIndex = frc::FloorMod(
-                                             i + offset,
-                                             static_cast<int>(d.size()));
-                                         return d[shiftedIndex];
-                                       },
-                                       data.size()},
-                 [=](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(
@@ -103,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
@@ -113,22 +95,9 @@ LEDPattern LEDPattern::ScrollAtAbsoluteSpeed(
     // offset values for negative velocities
     auto offset = static_cast<int64_t>(now) / microsPerLed;
 
-    self.ApplyTo(LEDPattern::LEDReader{[=, d = std::move(data)](auto i) {
-                                         int shiftedIndex = frc::FloorMod(
-                                             i + offset,
-                                             static_cast<int>(d.size()));
-                                         return d[shiftedIndex];
-                                       },
-                                       data.size()},
-                 [=, &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
@@ -79,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
diff --git a/wpilibj/src/main/java/edu/wpi/first/wpilibj/LEDPattern.java b/wpilibj/src/main/java/edu/wpi/first/wpilibj/LEDPattern.java
@@ -93,6 +93,18 @@
  */
 @FunctionalInterface
 public interface LEDPattern {
+  /** A functional interface for index mapping functions. */
+  @FunctionalInterface
+  interface IndexMapper {
+    /**
+     * Maps the index.
+     *
+     * @param bufLen Length of the buffer
+     * @param index The index to map
+     */
+    int apply(int bufLen, int index);
+  }
+
   /**
    * Writes the pattern to an LED buffer. Dynamic animations should be called periodically (such as
    * with a command or with a periodic method) to refresh the buffer over time.
@@ -130,19 +142,12 @@ default <T extends LEDReader & LEDWriter> void applyTo(T readWriter) {
   }
 
   /**
-   * Creates a pattern that displays this one in reverse. Scrolling patterns will scroll in the
-   * opposite direction (but at the same speed). It will treat the end of an LED strip as the start,
-   * and the start of the strip as the end. This can be useful for making ping-pong patterns that
-   * travel from one end of an LED strip to the other, then reverse direction and move back to the
-   * start. This can also be useful when working with LED strips connected in a serpentine pattern
-   * (where the start of one strip is connected to the end of the previous one); however, consider
-   * using a {@link AddressableLEDBufferView#reversed() reversed view} of the overall buffer for
-   * that segment rather than reversing patterns.
+   * Creates a pattern with remapped indices.
    *
-   * @return the reverse pattern
-   * @see AddressableLEDBufferView#reversed()
+   * @param indexMapper the index mapper
+   * @return the mapped pattern
    */
-  default LEDPattern reversed() {
+  default LEDPattern mapIndex(IndexMapper indexMapper) {
     return (reader, writer) -> {
       int bufLen = reader.getLength();
       applyTo(
@@ -154,23 +159,40 @@ public int getLength() {
 
             @Override
             public int getRed(int index) {
-              return reader.getRed(getLength() - 1 - index);
+              return reader.getRed(indexMapper.apply(bufLen, index));
             }
 
             @Override
             public int getGreen(int index) {
-              return reader.getGreen(getLength() - 1 - index);
+              return reader.getGreen(indexMapper.apply(bufLen, index));
             }
 
             @Override
             public int getBlue(int index) {
-              return reader.getBlue(getLength() - 1 - index);
+              return reader.getBlue(indexMapper.apply(bufLen, index));
             }
           },
-          (i, r, g, b) -> writer.setRGB((bufLen - 1) - i, r, g, b));
+          (i, r, g, b) -> writer.setRGB(indexMapper.apply(bufLen, i), r, g, b));
     };
   }
 
+  /**
+   * Creates a pattern that displays this one in reverse. Scrolling patterns will scroll in the
+   * opposite direction (but at the same speed). It will treat the end of an LED strip as the start,
+   * and the start of the strip as the end. This can be useful for making ping-pong patterns that
+   * travel from one end of an LED strip to the other, then reverse direction and move back to the
+   * start. This can also be useful when working with LED strips connected in a serpentine pattern
+   * (where the start of one strip is connected to the end of the previous one); however, consider
+   * using a {@link AddressableLEDBufferView#reversed() reversed view} of the overall buffer for
+   * that segment rather than reversing patterns.
+   *
+   * @return the reverse pattern
+   * @see AddressableLEDBufferView#reversed()
+   */
+  default LEDPattern reversed() {
+    return mapIndex((length, index) -> length - 1 - index);
+  }
+
   /**
    * Creates a pattern that plays this one, but offset by a certain number of LEDs. The offset
    * pattern will wrap around, if necessary.
@@ -179,15 +201,7 @@ public int getBlue(int index) {
    * @return the offset pattern
    */
   default LEDPattern offsetBy(int offset) {
-    return (reader, writer) -> {
-      int bufLen = reader.getLength();
-      applyTo(
-          LEDReader.RemappedReader.offset(reader, offset),
-          (i, r, g, b) -> {
-            int shiftedIndex = Math.floorMod(i + offset, bufLen);
-            writer.setRGB(shiftedIndex, r, g, b);
-          });
-    };
+    return mapIndex((length, index) -> Math.floorMod(index + offset, length));
   }
 
   /**
@@ -211,23 +225,16 @@ default LEDPattern offsetBy(int offset) {
   default LEDPattern scrollAtRelativeSpeed(Frequency velocity) {
     final double periodMicros = velocity.asPeriod().in(Microseconds);
 
-    return (reader, writer) -> {
-      int bufLen = reader.getLength();
-      long now = WPIUtilJNI.now();
+    return mapIndex(
+        (bufLen, index) -> {
+          long now = WPIUtilJNI.now();
 
-      // index should move by (buf.length) / (period)
-      double t = (now % (long) periodMicros) / periodMicros;
-      int offset = (int) (t * bufLen);
+          // index should move by (buf.length) / (period)
+          double t = (now % (long) periodMicros) / periodMicros;
+          int offset = (int) (t * bufLen);
 
-      applyTo(
-          LEDReader.RemappedReader.offset(reader, offset),
-          (i, r, g, b) -> {
-            // floorMod so if the offset is negative, we still get positive outputs
-            int shiftedIndex = Math.floorMod(i + offset, bufLen);
-
-            writer.setRGB(shiftedIndex, r, g, b);
-          });
-    };
+          return Math.floorMod(index + offset, bufLen);
+        });
   }
 
   /**
@@ -262,22 +269,16 @@ default LEDPattern scrollAtAbsoluteSpeed(LinearVelocity velocity, Distance ledSp
     var metersPerMicro = velocity.in(Meters.per(Microsecond));
     var microsPerLED = (int) (ledSpacing.in(Meters) / metersPerMicro);
 
-    return (reader, writer) -> {
-      int bufLen = reader.getLength();
-      long now = WPIUtilJNI.now();
-
-      // every step in time that's a multiple of microsPerLED will increment the offset by 1
-      var offset = (int) (now / microsPerLED);
+    return mapIndex(
+        (bufLen, index) -> {
+          long now = WPIUtilJNI.now();
 
-      applyTo(
-          LEDReader.RemappedReader.offset(reader, offset),
-          (i, r, g, b) -> {
-            // floorMod so if the offset is negative, we still get positive outputs
-            int shiftedIndex = Math.floorMod(i + offset, bufLen);
+          // every step in time that's a multiple of microsPerLED will increment the offset by 1
+          var offset = (int) (now / microsPerLED);
 
-            writer.setRGB(shiftedIndex, r, g, b);
-          });
-    };
+          // floorMod so if the offset is negative, we still get positive outputs
+          return Math.floorMod(index + offset, bufLen);
+        });
   }
 
   /**
diff --git a/wpilibj/src/main/java/edu/wpi/first/wpilibj/LEDReader.java b/wpilibj/src/main/java/edu/wpi/first/wpilibj/LEDReader.java
@@ -4,10 +4,8 @@
 
 package edu.wpi.first.wpilibj;
 
-import edu.wpi.first.util.ErrorMessages;
 import edu.wpi.first.wpilibj.util.Color;
 import edu.wpi.first.wpilibj.util.Color8Bit;
-import java.util.function.IntUnaryOperator;
 
 /** Generic interface for reading data from an LED buffer. */
 public interface LEDReader {
@@ -98,69 +96,4 @@ default void forEach(IndexedColorIterator iterator) {
       iterator.accept(i, getRed(i), getGreen(i), getBlue(i));
     }
   }
-
-  /**
-   * An LED reader implementation that operates on remapped indices. Remapping can be done using
-   * arbitrary mapping functions; a static factory function is also provided for the common use case
-   * of offset readers for things like scrolling animations.
-   */
-  class RemappedReader implements LEDReader {
-    private final LEDReader m_reader;
-    private final IntUnaryOperator m_mapping;
-
-    /**
-     * Creates a new remapping reader for a backing reader and an arbitrary index remapping
-     * function.
-     *
-     * @param reader the backing reader to use
-     * @param mapping the mapping function to use
-     */
-    public RemappedReader(LEDReader reader, IntUnaryOperator mapping) {
-      m_reader = ErrorMessages.requireNonNullParam(reader, "reader", "RemappedReader");
-      m_mapping = ErrorMessages.requireNonNullParam(mapping, "mapping", "RemappedReader");
-    }
-
-    /**
-     * Creates a new offset LED reader based on an existing reader and integer offset. The offset
-     * reader will be circular and wrap around.
-     *
-     * @param reader the backing reader to use
-     * @param offset the offset to use when reading data
-     * @return the offset reader
-     */
-    public static RemappedReader offset(LEDReader reader, int offset) {
-      return new RemappedReader(
-          reader, original -> Math.floorMod(original + offset, reader.getLength()));
-    }
-
-    @Override
-    public int getLength() {
-      return m_reader.getLength();
-    }
-
-    @Override
-    public int getRed(int index) {
-      return m_reader.getRed(remapIndex(index));
-    }
-
-    @Override
-    public int getGreen(int index) {
-      return m_reader.getGreen(remapIndex(index));
-    }
-
-    @Override
-    public int getBlue(int index) {
-      return m_reader.getBlue(remapIndex(index));
-    }
-
-    /**
-     * Remaps an input index to the corresponding offset index.
-     *
-     * @param index the input index to remap
-     * @return the remapped index
-     */
-    public int remapIndex(int index) {
-      return m_mapping.applyAsInt(index);
-    }
-  }
 }
