From 40e70d488b0730fcd140ada61c221bd7690c7cc8 Mon Sep 17 00:00:00 2001
From: Jade Turner <spacey-sooty@proton.me>
Date: Thu, 10 Oct 2024 21:04:24 +0800
Subject: [PATCH] [wpiutil] Add a RawFrame JNI overload for byte[]

Allows avoiding two copies in https://github.com/wpilibsuite/allwpilib/pull/7176/

Signed-off-by: Jade Turner <spacey-sooty@proton.me>
---
 .../java/edu/wpi/first/util/RawFrame.java     | 22 ++++++++++++++
 .../java/edu/wpi/first/util/WPIUtilJNI.java   |  3 ++
 .../src/main/native/cpp/jni/WPIUtilJNI.cpp    | 29 +++++++++++++++++++
 3 files changed, 54 insertions(+)

diff --git a/wpiutil/src/main/java/edu/wpi/first/util/RawFrame.java b/wpiutil/src/main/java/edu/wpi/first/util/RawFrame.java
index dd074bf3409..9a194f95508 100644
--- a/wpiutil/src/main/java/edu/wpi/first/util/RawFrame.java
+++ b/wpiutil/src/main/java/edu/wpi/first/util/RawFrame.java
@@ -88,6 +88,28 @@ public void setData(ByteBuffer data, int width, int height, int stride, PixelFor
         m_nativeObj, data, data.limit(), width, height, stride, pixelFormat.getValue());
   }
 
+  /**
+   * Set frame data.
+   *
+   * @param data A Java byte[] pointing to the frame data.
+   * @param width The width of the frame, in pixels
+   * @param height The height of the frame, in pixels
+   * @param stride The number of bytes in each row of image data
+   * @param pixelFormat The PixelFormat of the frame
+   */
+  public void setData(byte[] data, int width, int height, int stride, PixelFormat pixelFormat) {
+    if (!data.isDirect()) {
+      throw new UnsupportedOperationException("ByteBuffer must be direct");
+    }
+    m_data = data;
+    m_width = width;
+    m_height = height;
+    m_stride = stride;
+    m_pixelFormat = pixelFormat;
+    WPIUtilJNI.setRawFrameData(
+        m_nativeObj, data, data.length, width, height, stride, pixelFormat.getValue());
+  }
+
   /**
    * Call to set frame information.
    *
diff --git a/wpiutil/src/main/java/edu/wpi/first/util/WPIUtilJNI.java b/wpiutil/src/main/java/edu/wpi/first/util/WPIUtilJNI.java
index 136e2c041d6..0f7fdbef5c6 100644
--- a/wpiutil/src/main/java/edu/wpi/first/util/WPIUtilJNI.java
+++ b/wpiutil/src/main/java/edu/wpi/first/util/WPIUtilJNI.java
@@ -169,6 +169,9 @@ public static synchronized void forceLoad() throws IOException {
   static native void setRawFrameData(
       long frame, ByteBuffer data, int size, int width, int height, int stride, int pixelFormat);
 
+  static native void setRawFrameData(
+      long frame, byte[] data, int size, int width, int height, int stride, int pixelFormat);
+
   static native void setRawFrameInfo(
       long frame, int size, int width, int height, int stride, int pixelFormat);
 
diff --git a/wpiutil/src/main/native/cpp/jni/WPIUtilJNI.cpp b/wpiutil/src/main/native/cpp/jni/WPIUtilJNI.cpp
index 15b78958acd..f1d32bdc835 100644
--- a/wpiutil/src/main/native/cpp/jni/WPIUtilJNI.cpp
+++ b/wpiutil/src/main/native/cpp/jni/WPIUtilJNI.cpp
@@ -395,6 +395,35 @@ Java_edu_wpi_first_util_WPIUtilJNI_setRawFrameData
   f->pixelFormat = pixelFormat;
 }
 
+/*
+ * Class:     edu_wpi_first_util_WPIUtilJNI
+ * Method:    setRawFrameData
+ * Signature: (J[BIIIII)V
+ */
+JNIEXPORT void JNICALL
+Java_edu_wpi_first_util_WPIUtilJNI_setRawFrameData
+  (JNIEnv* env, jclass, jlong frame, jbyteArray data, jint size, jint width,
+   jint height, jint stride, jint pixelFormat)
+{
+  auto* f = reinterpret_cast<wpi::RawFrame*>(frame);
+  if (!f) {
+    wpi::ThrowNullPointerException(env, "frame is null");
+    return;
+  }
+  auto buf = env->GetByteArrayElements(env, data, size);
+  if (!buf) {
+    wpi::ThrowNullPointerException(env, "data is null");
+    return;
+  }
+  // there's no way to free a passed-in direct byte buffer
+  f->SetData(buf, size, env->GetDirectBufferCapacity(data), nullptr,
+             [](void*, void*, size_t) {});
+  f->width = width;
+  f->height = height;
+  f->stride = stride;
+  f->pixelFormat = pixelFormat;
+}
+
 /*
  * Class:     edu_wpi_first_util_WPIUtilJNI
  * Method:    setRawFrameInfo