[wpiutil,cscore,apriltag] Fix RawFrame (#6098)

wpiutil/src/main/native/cpp/RawFrame.cpp

@@ -6,25 +6,44 @@
 
 #include <wpi/MemAlloc.h>
 
+#include <cstring>
+
 extern "C" {
-void WPI_AllocateRawFrameData(WPI_RawFrame* frame, int requestedSize) {
-  if (frame->dataLength >= requestedSize) {
-    return;
+int WPI_AllocateRawFrameData(WPI_RawFrame* frame, size_t requestedSize) {
+  if (frame->capacity >= requestedSize) {
+    return 0;
   }
-  if (frame->data) {
-    frame->data =
-        static_cast<char*>(wpi::safe_realloc(frame->data, requestedSize));
-  } else {
-    frame->data = static_cast<char*>(wpi::safe_malloc(requestedSize));
-  }
-  frame->dataLength = requestedSize;
+  WPI_FreeRawFrameData(frame);
+  frame->data = static_cast<uint8_t*>(wpi::safe_malloc(requestedSize));
+  frame->capacity = requestedSize;
+  frame->size = 0;
+  return 1;
 }
 
 void WPI_FreeRawFrameData(WPI_RawFrame* frame) {
   if (frame->data) {
-    std::free(frame->data);
+    if (frame->freeFunc) {
+      frame->freeFunc(frame->freeCbData, frame->data, frame->capacity);
+    } else {
+      std::free(frame->data);
+    }
     frame->data = nullptr;
-    frame->dataLength = 0;
+    frame->freeFunc = nullptr;
+    frame->freeCbData = nullptr;
+    frame->capacity = 0;
   }
 }
+
+void WPI_SetRawFrameData(WPI_RawFrame* frame, void* data, size_t size,
+                         size_t capacity, void* cbdata,
+                         void (*freeFunc)(void* cbdata, void* data,
+                                          size_t capacity)) {
+  WPI_FreeRawFrameData(frame);
+  frame->data = static_cast<uint8_t*>(data);
+  frame->freeFunc = freeFunc;
+  frame->freeCbData = cbdata;
+  frame->capacity = capacity;
+  frame->size = size;
+}
+
 }  // extern "C"

wpiutil/src/main/native/cpp/jni/WPIUtilJNI.cpp

@@ -327,9 +327,7 @@ JNIEXPORT jlong JNICALL
 Java_edu_wpi_first_util_WPIUtilJNI_allocateRawFrame
   (JNIEnv*, jclass)
 {
-  wpi::RawFrame* rawFrame = new wpi::RawFrame{};
-  intptr_t rawFrameIntPtr = reinterpret_cast<intptr_t>(rawFrame);
-  return static_cast<jlong>(rawFrameIntPtr);
+  return reinterpret_cast<jlong>(new wpi::RawFrame);
 }
 
 /*
@@ -339,11 +337,76 @@ Java_edu_wpi_first_util_WPIUtilJNI_allocateRawFrame
  */
 JNIEXPORT void JNICALL
 Java_edu_wpi_first_util_WPIUtilJNI_freeRawFrame
-  (JNIEnv*, jclass, jlong rawFrame)
+  (JNIEnv*, jclass, jlong frame)
 {
-  wpi::RawFrame* ptr =
-      reinterpret_cast<wpi::RawFrame*>(static_cast<intptr_t>(rawFrame));
-  delete ptr;
+  delete reinterpret_cast<wpi::RawFrame*>(frame);
+}
+
+/*
+ * Class:     edu_wpi_first_util_WPIUtilJNI
+ * Method:    getRawFrameDataPtr
+ * Signature: (J)J
+ */
+JNIEXPORT jlong JNICALL
+Java_edu_wpi_first_util_WPIUtilJNI_getRawFrameDataPtr
+  (JNIEnv* env, jclass, jlong frame)
+{
+  auto* f = reinterpret_cast<wpi::RawFrame*>(frame);
+  if (!f) {
+    wpi::ThrowNullPointerException(env, "frame is null");
+    return 0;
+  }
+  return reinterpret_cast<jlong>(f->data);
+}
+
+/*
+ * Class:     edu_wpi_first_util_WPIUtilJNI
+ * Method:    setRawFrameData
+ * Signature: (JLjava/lang/Object;IIIII)V
+ */
+JNIEXPORT void JNICALL
+Java_edu_wpi_first_util_WPIUtilJNI_setRawFrameData
+  (JNIEnv* env, jclass, jlong frame, jobject 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->GetDirectBufferAddress(data);
+  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
+ * Signature: (JIIIII)V
+ */
+JNIEXPORT void JNICALL
+Java_edu_wpi_first_util_WPIUtilJNI_setRawFrameInfo
+  (JNIEnv* env, jclass, jlong frame, 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;
+  }
+  f->width = width;
+  f->height = height;
+  f->stride = stride;
+  f->pixelFormat = pixelFormat;
 }
 
 }  // extern "C"

wpiutil/src/main/native/include/wpi/RawFrame.h

@@ -5,38 +5,65 @@
 #ifndef WPIUTIL_WPI_RAWFRAME_H_
 #define WPIUTIL_WPI_RAWFRAME_H_
 
-/**
- * Raw Frame
- */
-typedef struct WPI_RawFrame {  // NOLINT
-  char* data;
-  int dataLength;
-  int pixelFormat;
-  int width;
-  int height;
-  int totalData;
-} WPI_RawFrame;
+#include <stdint.h>
+
+#ifdef __cplusplus
+#include <concepts>
+#include <cstddef>
+#else
+
+#include <stddef.h>  // NOLINT
+
+#endif
+
+#ifdef WPI_RAWFRAME_JNI
+#include "jni_util.h"
+#endif
+
+// NOLINT
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
+/**
+ * Raw Frame
+ */
+typedef struct WPI_RawFrame {  // NOLINT
+  // image data
+  uint8_t* data;
+  // function to free image data (may be NULL)
+  void (*freeFunc)(void* cbdata, void* data, size_t capacity);
+  void* freeCbData;  // data passed to freeFunc
+  size_t capacity;   // data buffer capacity, in bytes
+  size_t size;       // actual size of data, in bytes
+  int pixelFormat;   // WPI_PixelFormat
+  int width;         // width of image, in pixels
+  int height;        // height of image, in pixels
+  int stride;        // size of each row of data, in bytes (may be 0)
+} WPI_RawFrame;
+
 /**
  * Pixel formats
  */
 enum WPI_PixelFormat {
-  WPI_PIXFMT_UNKNOWN = 0,
-  WPI_PIXFMT_MJPEG,
-  WPI_PIXFMT_YUYV,
-  WPI_PIXFMT_RGB565,
-  WPI_PIXFMT_BGR,
-  WPI_PIXFMT_GRAY,
-  WPI_PIXFMT_Y16,
-  WPI_PIXFMT_UYVY
+  WPI_PIXFMT_UNKNOWN = 0,  // unknown
+  WPI_PIXFMT_MJPEG,        // Motion-JPEG (compressed image data)
+  WPI_PIXFMT_YUYV,         // YUV 4:2:2, 16 bpp
+  WPI_PIXFMT_RGB565,       // RGB 5-6-5, 16 bpp
+  WPI_PIXFMT_BGR,          // BGR 8-8-8, 24 bpp
+  WPI_PIXFMT_GRAY,         // Grayscale, 8 bpp
+  WPI_PIXFMT_Y16,          // Grayscale, 16 bpp
+  WPI_PIXFMT_UYVY,         // YUV 4:2:2, 16 bpp
 };
 
-void WPI_AllocateRawFrameData(WPI_RawFrame* frame, int requestedSize);
+// Returns nonzero if the frame data was allocated/reallocated
+int WPI_AllocateRawFrameData(WPI_RawFrame* frame, size_t requestedSize);
 void WPI_FreeRawFrameData(WPI_RawFrame* frame);
+void WPI_SetRawFrameData(WPI_RawFrame* frame, void* data, size_t size,
+                         size_t capacity, void* cbdata,
+                         void (*freeFunc)(void* cbdata, void* data,
+                                          size_t capacity));
 
 #ifdef __cplusplus
 }  // extern "C"
@@ -47,17 +74,69 @@ namespace wpi {
 struct RawFrame : public WPI_RawFrame {
   RawFrame() {
     data = nullptr;
-    dataLength = 0;
+    freeFunc = nullptr;
+    freeCbData = nullptr;
+    capacity = 0;
+    size = 0;
     pixelFormat = WPI_PIXFMT_UNKNOWN;
     width = 0;
     height = 0;
-    totalData = 0;
+  }
+  RawFrame(const RawFrame&) = delete;
+  RawFrame& operator=(const RawFrame&) = delete;
+  RawFrame(RawFrame&& rhs) noexcept : WPI_RawFrame{rhs} {
+    rhs.data = nullptr;
+    rhs.freeFunc = nullptr;
+    rhs.freeCbData = nullptr;
+    rhs.capacity = 0;
+    rhs.size = 0;
+  }
+  RawFrame& operator=(RawFrame&& rhs) noexcept {
+    *static_cast<WPI_RawFrame*>(this) = rhs;
+    rhs.data = nullptr;
+    rhs.freeFunc = nullptr;
+    rhs.freeCbData = nullptr;
+    rhs.capacity = 0;
+    rhs.size = 0;
+    return *this;
+  }
+
+  void SetData(void* data, size_t size, size_t capacity, void* cbdata,
+               void (*freeFunc)(void* cbdata, void* data, size_t capacity)) {
+    WPI_SetRawFrameData(this, data, size, capacity, cbdata, freeFunc);
+  }
+
+  // returns true if the frame data was allocated/reallocated
+  bool Reserve(size_t size) {
+    return WPI_AllocateRawFrameData(this, size) != 0;
   }
 
   ~RawFrame() { WPI_FreeRawFrameData(this); }
 };
-}  // namespace wpi
 
+#ifdef WPI_RAWFRAME_JNI
+template <std::same_as<wpi::RawFrame> T>
+void SetFrameData(JNIEnv* env, jclass rawFrameCls, jobject jframe,
+                  const T& frame, bool newData) {
+  if (newData) {
+    static jmethodID setData = env->GetMethodID(rawFrameCls, "setDataJNI",
+                                                "(Ljava/nio/ByteBuffer;IIII)V");
+    env->CallVoidMethod(
+        jframe, setData, env->NewDirectByteBuffer(frame.data, frame.size),
+        static_cast<jint>(frame.width), static_cast<jint>(frame.height),
+        static_cast<jint>(frame.stride), static_cast<jint>(frame.pixelFormat));
+  } else {
+    static jmethodID setInfo =
+        env->GetMethodID(rawFrameCls, "setInfoJNI", "(IIII)V");
+    env->CallVoidMethod(jframe, setInfo, static_cast<jint>(frame.width),
+                        static_cast<jint>(frame.height),
+                        static_cast<jint>(frame.stride),
+                        static_cast<jint>(frame.pixelFormat));
+  }
+}
+#endif
+
+}  // namespace wpi
 #endif
 
 #endif  // WPIUTIL_WPI_RAWFRAME_H_