Add transform methods to Trajectory (#2187)

wpilibc/src/main/native/cpp/trajectory/Trajectory.cpp

@@ -107,6 +107,33 @@ Trajectory::State Trajectory::Sample(units::second_t t) const {
                                 (t - prevSample.t) / (sample.t - prevSample.t));
 }
 
+Trajectory Trajectory::TransformBy(const Transform2d& transform) {
+  auto& firstState = m_states[0];
+  auto& firstPose = firstState.pose;
+
+  // Calculate the transformed first pose.
+  auto newFirstPose = firstPose + transform;
+  auto newStates = m_states;
+  newStates[0].pose = newFirstPose;
+
+  for (unsigned int i = 1; i < newStates.size(); i++) {
+    auto& state = newStates[i];
+    // We are transforming relative to the coordinate frame of the new initial
+    // pose.
+    state.pose = newFirstPose + (state.pose - firstPose);
+  }
+
+  return Trajectory(newStates);
+}
+
+Trajectory Trajectory::RelativeTo(const Pose2d& pose) {
+  auto newStates = m_states;
+  for (auto& state : newStates) {
+    state.pose = state.pose.RelativeTo(pose);
+  }
+  return Trajectory(newStates);
+}
+
 void frc::to_json(wpi::json& json, const Trajectory::State& state) {
   json = wpi::json{{"time", state.t.to<double>()},
                    {"velocity", state.velocity.to<double>()},

wpilibc/src/main/native/include/frc/trajectory/Trajectory.h

@@ -12,6 +12,7 @@
 #include <units/units.h>
 
 #include "frc/geometry/Pose2d.h"
+#include "frc/geometry/Transform2d.h"
 
 namespace wpi {
 class json;
@@ -105,6 +106,27 @@ class Trajectory {
    */
   State Sample(units::second_t t) const;
 
+  /**
+   * Transforms all poses in the trajectory by the given transform. This is
+   * useful for converting a robot-relative trajectory into a field-relative
+   * trajectory. This works with respect to the first pose in the trajectory.
+   *
+   * @param transform The transform to transform the trajectory by.
+   * @return The transformed trajectory.
+   */
+  Trajectory TransformBy(const Transform2d& transform);
+
+  /**
+   * Transforms all poses in the trajectory so that they are relative to the
+   * given pose. This is useful for converting a field-relative trajectory
+   * into a robot-relative trajectory.
+   *
+   * @param pose The pose that is the origin of the coordinate frame that
+   *             the current trajectory will be transformed into.
+   * @return The transformed trajectory.
+   */
+  Trajectory RelativeTo(const Pose2d& pose);
+
   /**
    * Returns the initial pose of the trajectory.
    *