[wpimath] Use JNI for trajectory serialization (#3257)

wpimath/src/main/java/edu/wpi/first/math/WPIMathJNI.java

@@ -91,6 +91,38 @@ public final class WPIMathJNI {
    */
   public static native boolean isStabilizable(int states, int inputs, double[] A, double[] B);
 
+  /**
+   * Loads a Pathweaver JSON.
+   *
+   * @param path The path to the JSON.
+   * @return A double array with the trajectory states from the JSON.
+   */
+  public static native double[] fromPathweaverJson(String path) throws IOException;
+
+  /**
+   * Converts a trajectory into a Pathweaver JSON and saves it.
+   *
+   * @param elements The elements of the trajectory.
+   * @param path The location to save the JSON to.
+   */
+  public static native void toPathweaverJson(double[] elements, String path) throws IOException;
+
+  /**
+   * Deserializes a trajectory JSON into a double[] of trajectory elements.
+   *
+   * @param json The JSON containing the serialized trajectory.
+   * @return A double array with the trajectory states.
+   */
+  public static native double[] deserializeTrajectory(String json);
+
+  /**
+   * Serializes the trajectory into a JSON string.
+   *
+   * @param elements The elements of the trajectory.
+   * @return A JSON containing the serialized trajectory.
+   */
+  public static native String serializeTrajectory(double[] elements);
+
   public static class Helper {
     private static AtomicBoolean extractOnStaticLoad = new AtomicBoolean(true);
 

wpimath/src/main/java/edu/wpi/first/wpilibj/trajectory/TrajectoryUtil.java

@@ -4,73 +4,117 @@
 
 package edu.wpi.first.wpilibj.trajectory;
 
-import com.fasterxml.jackson.core.JsonProcessingException;
-import com.fasterxml.jackson.databind.ObjectMapper;
-import com.fasterxml.jackson.databind.ObjectReader;
-import com.fasterxml.jackson.databind.ObjectWriter;
-import java.io.BufferedReader;
-import java.io.BufferedWriter;
+import edu.wpi.first.math.WPIMathJNI;
+import edu.wpi.first.wpilibj.geometry.Pose2d;
+import edu.wpi.first.wpilibj.geometry.Rotation2d;
 import java.io.IOException;
-import java.nio.file.Files;
 import java.nio.file.Path;
-import java.util.Arrays;
+import java.util.ArrayList;
+import java.util.List;
 
 public final class TrajectoryUtil {
-  private static final ObjectReader READER = new ObjectMapper().readerFor(Trajectory.State[].class);
-  private static final ObjectWriter WRITER = new ObjectMapper().writerFor(Trajectory.State[].class);
-
   private TrajectoryUtil() {
     throw new UnsupportedOperationException("This is a utility class!");
   }
 
+  /**
+   * Creates a trajectory from a double[] of elements.
+   *
+   * @param elements A double[] containing the raw elements of the trajectory.
+   * @return A trajectory created from the elements.
+   */
+  @SuppressWarnings("PMD.AvoidInstantiatingObjectsInLoops")
+  private static Trajectory createTrajectoryFromElements(double[] elements) {
+    // Make sure that the elements have the correct length.
+    if (elements.length % 7 != 0) {
+      throw new TrajectorySerializationException(
+          "An error occurred when converting trajectory elements into a trajectory.");
+    }
+
+    // Create a list of states from the elements.
+    List<Trajectory.State> states = new ArrayList<>();
+    for (int i = 0; i < elements.length; i += 7) {
+      states.add(
+          new Trajectory.State(
+              elements[i],
+              elements[i + 1],
+              elements[i + 2],
+              new Pose2d(elements[i + 3], elements[i + 4], new Rotation2d(elements[i + 5])),
+              elements[i + 6]));
+    }
+    return new Trajectory(states);
+  }
+
+  /**
+   * Returns a double[] of elements from the given trajectory.
+   *
+   * @param trajectory The trajectory to retrieve raw elements from.
+   * @return A double[] of elements from the given trajectory.
+   */
+  private static double[] getElementsFromTrajectory(Trajectory trajectory) {
+    // Create a double[] of elements and fill it from the trajectory states.
+    double[] elements = new double[trajectory.getStates().size() * 7];
+
+    for (int i = 0; i < trajectory.getStates().size() * 7; i += 7) {
+      var state = trajectory.getStates().get(i / 7);
+      elements[i] = state.timeSeconds;
+      elements[i + 1] = state.velocityMetersPerSecond;
+      elements[i + 2] = state.accelerationMetersPerSecondSq;
+      elements[i + 3] = state.poseMeters.getX();
+      elements[i + 4] = state.poseMeters.getY();
+      elements[i + 5] = state.poseMeters.getRotation().getRadians();
+      elements[i + 6] = state.curvatureRadPerMeter;
+    }
+    return elements;
+  }
+
   /**
    * Imports a Trajectory from a PathWeaver-style JSON file.
    *
-   * @param path the path of the json file to import from
+   * @param path The path of the json file to import from
    * @return The trajectory represented by the file.
-   * @throws IOException if reading from the file fails
+   * @throws IOException if reading from the file fails.
    */
   public static Trajectory fromPathweaverJson(Path path) throws IOException {
-    try (BufferedReader reader = Files.newBufferedReader(path)) {
-      Trajectory.State[] state = READER.readValue(reader);
-      return new Trajectory(Arrays.asList(state));
-    }
+    return createTrajectoryFromElements(WPIMathJNI.fromPathweaverJson(path.toString()));
   }
 
   /**
    * Exports a Trajectory to a PathWeaver-style JSON file.
    *
-   * @param trajectory the trajectory to export
-   * @param path the path of the file to export to
-   * @throws IOException if writing to the file fails
+   * @param trajectory The trajectory to export
+   * @param path The path of the file to export to
+   * @throws IOException if writing to the file fails.
    */
   public static void toPathweaverJson(Trajectory trajectory, Path path) throws IOException {
-    Files.createDirectories(path.getParent());
-    try (BufferedWriter writer = Files.newBufferedWriter(path)) {
-      WRITER.writeValue(writer, trajectory.getStates().toArray(new Trajectory.State[0]));
-    }
+    WPIMathJNI.toPathweaverJson(getElementsFromTrajectory(trajectory), path.toString());
   }
 
   /**
    * Deserializes a Trajectory from PathWeaver-style JSON.
    *
-   * @param json the string containing the serialized JSON
+   * @param json The string containing the serialized JSON
    * @return the trajectory represented by the JSON
-   * @throws JsonProcessingException if deserializing the JSON fails
+   * @throws TrajectorySerializationException if deserialization of the string fails.
    */
-  public static Trajectory deserializeTrajectory(String json) throws JsonProcessingException {
-    Trajectory.State[] state = READER.readValue(json);
-    return new Trajectory(Arrays.asList(state));
+  public static Trajectory deserializeTrajectory(String json) {
+    return createTrajectoryFromElements(WPIMathJNI.deserializeTrajectory(json));
   }
 
   /**
    * Serializes a Trajectory to PathWeaver-style JSON.
    *
-   * @param trajectory the trajectory to export
-   * @return the string containing the serialized JSON
-   * @throws JsonProcessingException if serializing the Trajectory fails
+   * @param trajectory The trajectory to export
+   * @return The string containing the serialized JSON
+   * @throws TrajectorySerializationException if serialization of the trajectory fails.
    */
-  public static String serializeTrajectory(Trajectory trajectory) throws JsonProcessingException {
-    return WRITER.writeValueAsString(trajectory.getStates().toArray(new Trajectory.State[0]));
+  public static String serializeTrajectory(Trajectory trajectory) {
+    return WPIMathJNI.serializeTrajectory(getElementsFromTrajectory(trajectory));
+  }
+
+  public static class TrajectorySerializationException extends RuntimeException {
+    public TrajectorySerializationException(String message) {
+      super(message);
+    }
   }
 }

wpimath/src/main/native/cpp/jni/WPIMathJNI.cpp

@@ -13,6 +13,7 @@
 #include "Eigen/QR"
 #include "drake/math/discrete_algebraic_riccati_equation.h"
 #include "edu_wpi_first_math_WPIMathJNI.h"
+#include "frc/trajectory/TrajectoryUtil.h"
 #include "unsupported/Eigen/MatrixFunctions"
 
 using namespace wpi::java;
@@ -43,6 +44,46 @@ bool check_stabilizable(const Eigen::Ref<const Eigen::MatrixXd>& A,
   return true;
 }
 
+std::vector<double> GetElementsFromTrajectory(
+    const frc::Trajectory& trajectory) {
+  std::vector<double> elements;
+  elements.reserve(trajectory.States().size() * 7);
+
+  for (auto&& state : trajectory.States()) {
+    elements.push_back(state.t.to<double>());
+    elements.push_back(state.velocity.to<double>());
+    elements.push_back(state.acceleration.to<double>());
+    elements.push_back(state.pose.X().to<double>());
+    elements.push_back(state.pose.Y().to<double>());
+    elements.push_back(state.pose.Rotation().Radians().to<double>());
+    elements.push_back(state.curvature.to<double>());
+  }
+
+  return elements;
+}
+
+frc::Trajectory CreateTrajectoryFromElements(wpi::ArrayRef<double> elements) {
+  // Make sure that the elements have the correct length.
+  assert(elements.size() % 7 == 0);
+
+  // Create a vector of states from the elements.
+  std::vector<frc::Trajectory::State> states;
+  states.reserve(elements.size() / 7);
+
+  for (size_t i = 0; i < elements.size(); i += 7) {
+    states.emplace_back(frc::Trajectory::State{
+        units::second_t{elements[i]},
+        units::meters_per_second_t{elements[i + 1]},
+        units::meters_per_second_squared_t{elements[i + 2]},
+        frc::Pose2d{units::meter_t{elements[i + 3]},
+                    units::meter_t{elements[i + 4]},
+                    units::radian_t{elements[i + 5]}},
+        units::curvature_t{elements[i + 6]}});
+  }
+
+  return frc::Trajectory(states);
+}
+
 extern "C" {
 
 /*
@@ -163,4 +204,99 @@ Java_edu_wpi_first_math_WPIMathJNI_isStabilizable
   return isStabilizable;
 }
 
+/*
+ * Class:     edu_wpi_first_math_WPIMathJNI
+ * Method:    fromPathweaverJson
+ * Signature: (Ljava/lang/String;)[D
+ */
+JNIEXPORT jdoubleArray JNICALL
+Java_edu_wpi_first_math_WPIMathJNI_fromPathweaverJson
+  (JNIEnv* env, jclass, jstring path)
+{
+  try {
+    auto trajectory =
+        frc::TrajectoryUtil::FromPathweaverJson(JStringRef{env, path}.c_str());
+    std::vector<double> elements = GetElementsFromTrajectory(trajectory);
+    return MakeJDoubleArray(env, elements);
+  } catch (std::exception& e) {
+    jclass cls = env->FindClass("java/lang/IOException");
+    if (cls) {
+      env->ThrowNew(cls, e.what());
+    }
+    return nullptr;
+  }
+}
+
+/*
+ * Class:     edu_wpi_first_math_WPIMathJNI
+ * Method:    toPathweaverJson
+ * Signature: ([DLjava/lang/String;)V
+ */
+JNIEXPORT void JNICALL
+Java_edu_wpi_first_math_WPIMathJNI_toPathweaverJson
+  (JNIEnv* env, jclass, jdoubleArray elements, jstring path)
+{
+  try {
+    auto trajectory =
+        CreateTrajectoryFromElements(JDoubleArrayRef{env, elements});
+    frc::TrajectoryUtil::ToPathweaverJson(trajectory,
+                                          JStringRef{env, path}.c_str());
+  } catch (std::exception& e) {
+    jclass cls = env->FindClass("java/lang/IOException");
+    if (cls) {
+      env->ThrowNew(cls, e.what());
+    }
+  }
+}
+
+/*
+ * Class:     edu_wpi_first_math_WPIMathJNI
+ * Method:    deserializeTrajectory
+ * Signature: (Ljava/lang/String;)[D
+ */
+JNIEXPORT jdoubleArray JNICALL
+Java_edu_wpi_first_math_WPIMathJNI_deserializeTrajectory
+  (JNIEnv* env, jclass, jstring json)
+{
+  try {
+    auto trajectory = frc::TrajectoryUtil::DeserializeTrajectory(
+        JStringRef{env, json}.c_str());
+    std::vector<double> elements = GetElementsFromTrajectory(trajectory);
+    return MakeJDoubleArray(env, elements);
+  } catch (std::exception& e) {
+    jclass cls = env->FindClass(
+        "edu/wpi/first/wpilibj/trajectory/TrajectoryUtil$"
+        "TrajectorySerializationException");
+    if (cls) {
+      env->ThrowNew(cls, e.what());
+    }
+    return nullptr;
+  }
+}
+
+/*
+ * Class:     edu_wpi_first_math_WPIMathJNI
+ * Method:    serializeTrajectory
+ * Signature: ([D)Ljava/lang/String;
+ */
+JNIEXPORT jstring JNICALL
+Java_edu_wpi_first_math_WPIMathJNI_serializeTrajectory
+  (JNIEnv* env, jclass, jdoubleArray elements)
+{
+  try {
+    auto trajectory =
+        CreateTrajectoryFromElements(JDoubleArrayRef{env, elements});
+    return MakeJString(env,
+                       frc::TrajectoryUtil::SerializeTrajectory(trajectory));
+  } catch (std::exception& e) {
+    jclass cls = env->FindClass(
+        "edu/wpi/first/wpilibj/trajectory/TrajectoryUtil$"
+        "TrajectorySerializationException");
+    if (cls) {
+      env->ThrowNew(cls, e.what());
+    }
+    return nullptr;
+  }
+}
+
 }  // extern "C"