feat: Use JSI's ArrayBuffer instead of TypedArray (#2408)

* feat: Use JSI's `ArrayBuffer` instead of `TypedArray`

* fix: Fix move memory

* feat: Implement iOS

* Format

* Update JSIJNIConversion.cpp

* fix: Fix Android `toArrayBuffer` and other

* Catch FP call errors

* Update return type

* Use `CPU_READ_OFTEN` flag as well

* CPU flag

* Run destructors under `jni::ThreadScope`

* Update FrameProcessorPluginHostObject.cpp

* fix: Fix `toArrayBuffer()` crash

* Update Frame.ts

package/android/src/main/cpp/MutableJByteBuffer.cpp

@@ -0,0 +1,33 @@
+//
+// Created by Marc Rousavy on 17.01.24.
+//
+
+#include "MutableJByteBuffer.h"
+
+#include <fbjni/ByteBuffer.h>
+#include <fbjni/fbjni.h>
+#include <jsi/jsi.h>
+
+namespace vision {
+
+MutableJByteBuffer::MutableJByteBuffer(jni::alias_ref<jni::JByteBuffer> byteBuffer) {
+  _byteBuffer = jni::make_global(byteBuffer);
+}
+
+MutableJByteBuffer::~MutableJByteBuffer() noexcept {
+  jni::ThreadScope::WithClassLoader([&] { _byteBuffer.reset(); });
+}
+
+uint8_t* MutableJByteBuffer::data() {
+  return _byteBuffer->getDirectBytes();
+}
+
+size_t MutableJByteBuffer::size() const {
+  return _byteBuffer->getDirectSize();
+}
+
+jni::global_ref<jni::JByteBuffer> MutableJByteBuffer::getByteBuffer() {
+  return _byteBuffer;
+}
+
+} // namespace vision

package/android/src/main/cpp/MutableJByteBuffer.h

@@ -0,0 +1,33 @@
+//
+// Created by Marc Rousavy on 17.01.24.
+//
+
+#pragma once
+
+#include <fbjni/ByteBuffer.h>
+#include <fbjni/fbjni.h>
+#include <jsi/jsi.h>
+#include <memory>
+
+namespace vision {
+
+using namespace facebook;
+
+class MutableJByteBuffer : public jsi::MutableBuffer {
+public:
+  /**
+   * Wraps the given JByteBuffer in a MutableBuffer for use in JS.
+   */
+  explicit MutableJByteBuffer(jni::alias_ref<jni::JByteBuffer> byteBuffer);
+  ~MutableJByteBuffer();
+
+public:
+  uint8_t* data() override;
+  size_t size() const override;
+  jni::global_ref<jni::JByteBuffer> getByteBuffer();
+
+private:
+  jni::global_ref<jni::JByteBuffer> _byteBuffer;
+};
+
+} // namespace vision

package/android/src/main/cpp/frameprocessor/FrameHostObject.cpp

@@ -7,7 +7,7 @@
 #include <fbjni/fbjni.h>
 #include <jni.h>
 
-#include "JSITypedArray.h"
+#include "MutableRawBuffer.h"
 
 #include <string>
 #include <vector>
@@ -102,15 +102,18 @@ jsi::Value FrameHostObject::get(jsi::Runtime& runtime, const jsi::PropNameID& pr
 
       static constexpr auto ARRAYBUFFER_CACHE_PROP_NAME = "__frameArrayBufferCache";
       if (!runtime.global().hasProperty(runtime, ARRAYBUFFER_CACHE_PROP_NAME)) {
-        vision::TypedArray<vision::TypedArrayKind::Uint8ClampedArray> arrayBuffer(runtime, size);
+        auto mutableBuffer = std::make_shared<vision::MutableRawBuffer>(size);
+        jsi::ArrayBuffer arrayBuffer(runtime, mutableBuffer);
         runtime.global().setProperty(runtime, ARRAYBUFFER_CACHE_PROP_NAME, arrayBuffer);
       }
 
       // Get from global JS cache
       auto arrayBufferCache = runtime.global().getPropertyAsObject(runtime, ARRAYBUFFER_CACHE_PROP_NAME);
-      auto arrayBuffer = vision::getTypedArray(runtime, arrayBufferCache).get<vision::TypedArrayKind::Uint8ClampedArray>(runtime);
+      auto arrayBuffer = arrayBufferCache.getArrayBuffer(runtime);
+
       if (arrayBuffer.size(runtime) != size) {
-        arrayBuffer = vision::TypedArray<vision::TypedArrayKind::Uint8ClampedArray>(runtime, size);
+        auto mutableBuffer = std::make_shared<vision::MutableRawBuffer>(size);
+        arrayBuffer = jsi::ArrayBuffer(runtime, mutableBuffer);
         runtime.global().setProperty(runtime, ARRAYBUFFER_CACHE_PROP_NAME, arrayBuffer);
       }
 
@@ -170,4 +173,6 @@ jsi::Value FrameHostObject::get(jsi::Runtime& runtime, const jsi::PropNameID& pr
   return HostObject::get(runtime, propName);
 }
 
+#undef JSI_FUNC
+
 } // namespace vision

package/android/src/main/cpp/frameprocessor/FrameProcessorPluginHostObject.cpp

@@ -12,6 +12,13 @@ namespace vision {
 
 using namespace facebook;
 
+FrameProcessorPluginHostObject::FrameProcessorPluginHostObject(jni::alias_ref<JFrameProcessorPlugin::javaobject> plugin)
+    : _plugin(make_global(plugin)) {}
+
+FrameProcessorPluginHostObject::~FrameProcessorPluginHostObject() {
+  jni::ThreadScope::WithClassLoader([&] { _plugin.reset(); });
+}
+
 std::vector<jsi::PropNameID> FrameProcessorPluginHostObject::getPropertyNames(jsi::Runtime& runtime) {
   std::vector<jsi::PropNameID> result;
   result.push_back(jsi::PropNameID::forUtf8(runtime, std::string("call")));

package/android/src/main/cpp/frameprocessor/FrameProcessorPluginHostObject.h

@@ -16,8 +16,8 @@ using namespace facebook;
 
 class FrameProcessorPluginHostObject : public jsi::HostObject {
 public:
-  explicit FrameProcessorPluginHostObject(jni::alias_ref<JFrameProcessorPlugin::javaobject> plugin) : _plugin(make_global(plugin)) {}
-  ~FrameProcessorPluginHostObject() {}
+  explicit FrameProcessorPluginHostObject(jni::alias_ref<JFrameProcessorPlugin::javaobject> plugin);
+  ~FrameProcessorPluginHostObject();
 
 public:
   std::vector<jsi::PropNameID> getPropertyNames(jsi::Runtime& runtime) override;

package/android/src/main/cpp/frameprocessor/JSIJNIConversion.cpp

@@ -15,7 +15,6 @@
 
 #include "FrameHostObject.h"
 #include "JFrame.h"
-#include "JSITypedArray.h"
 #include "JSharedArray.h"
 
 namespace vision {
@@ -61,8 +60,8 @@ jni::local_ref<jobject> JSIJNIConversion::convertJSIValueToJNIObject(jsi::Runtim
     } else if (valueAsObject.isArrayBuffer(runtime)) {
       // ArrayBuffer/TypedArray
 
-      TypedArrayBase array = getTypedArray(runtime, valueAsObject);
-      return JSharedArray::create(runtime, std::move(array));
+      jsi::ArrayBuffer arrayBuffer = valueAsObject.getArrayBuffer(runtime);
+      return JSharedArray::create(runtime, std::move(arrayBuffer));
 
     } else if (valueAsObject.isHostObject(runtime)) {
 
@@ -176,8 +175,8 @@ jsi::Value JSIJNIConversion::convertJNIObjectToJSIValue(jsi::Runtime& runtime, c
     // SharedArray
     auto sharedArray = static_ref_cast<JSharedArray::javaobject>(object);
 
-    std::shared_ptr<TypedArrayBase> array = sharedArray->cthis()->getTypedArray();
-    return array->getBuffer(runtime);
+    std::shared_ptr<jsi::ArrayBuffer> array = sharedArray->cthis()->getArrayBuffer();
+    return array->getArrayBuffer(runtime);
   }
 
   auto type = object->getClass()->toString();

package/android/src/main/cpp/frameprocessor/VisionCameraProxy.cpp

@@ -13,7 +13,6 @@
 #include <fbjni/fbjni.h>
 
 #include "FrameProcessorPluginHostObject.h"
-#include "JSITypedArray.h"
 
 #include <memory>
 #include <string>

package/android/src/main/cpp/frameprocessor/java-bindings/JSharedArray.cpp

@@ -9,30 +9,23 @@ namespace vision {
 
 using namespace facebook;
 
-TypedArrayKind getTypedArrayKind(int unsafeEnumValue) {
-  return static_cast<TypedArrayKind>(unsafeEnumValue);
+jni::local_ref<JSharedArray::javaobject> JSharedArray::create(jsi::Runtime& runtime, jsi::ArrayBuffer arrayBuffer) {
+  jni::local_ref<JSharedArray::javaobject> instance = newObjectCxxArgs(runtime, std::make_shared<jsi::ArrayBuffer>(std::move(arrayBuffer)));
+  instance->cthis()->_javaPart = jni::make_global(instance);
+  return instance;
 }
 
-jni::local_ref<JSharedArray::javaobject> JSharedArray::create(jsi::Runtime& runtime, TypedArrayBase array) {
-  return newObjectCxxArgs(runtime, std::make_shared<TypedArrayBase>(std::move(array)));
+JSharedArray::JSharedArray(jsi::Runtime& runtime, std::shared_ptr<jsi::ArrayBuffer> arrayBuffer) {
+  size_t size = arrayBuffer->size(runtime);
+  jni::local_ref<JByteBuffer> byteBuffer = JByteBuffer::allocateDirect(size);
+
+  _arrayBuffer = arrayBuffer;
+  _byteBuffer = jni::make_global(byteBuffer);
+  _size = size;
 }
 
-jni::global_ref<jni::JByteBuffer> JSharedArray::wrapInByteBuffer(jsi::Runtime& runtime, std::shared_ptr<TypedArrayBase> typedArray) {
-  jsi::ArrayBuffer arrayBuffer = typedArray->getBuffer(runtime);
-  __android_log_print(ANDROID_LOG_INFO, TAG, "Wrapping ArrayBuffer in a JNI ByteBuffer...");
-  auto byteBuffer = jni::JByteBuffer::wrapBytes(arrayBuffer.data(runtime), arrayBuffer.size(runtime));
-  __android_log_print(ANDROID_LOG_INFO, TAG, "Successfully created TypedArray (JNI Size: %i)!", byteBuffer->getDirectSize());
-  return jni::make_global(byteBuffer);
-}
-
-JSharedArray::JSharedArray(jsi::Runtime& runtime, std::shared_ptr<TypedArrayBase> array) {
-  _array = array;
-  _byteBuffer = wrapInByteBuffer(runtime, _array);
-  _size = _array->size(runtime);
-}
-
-JSharedArray::JSharedArray(const jni::alias_ref<JSharedArray::jhybridobject>& javaThis,
-                           const jni::alias_ref<JVisionCameraProxy::javaobject>& proxy, int dataType, int size) {
+JSharedArray::JSharedArray(const jni::alias_ref<jhybridobject>& javaThis, const jni::alias_ref<JVisionCameraProxy::javaobject>& proxy,
+                           jni::alias_ref<JByteBuffer> byteBuffer) {
   _javaPart = jni::make_global(javaThis);
 
 #if VISION_CAMERA_ENABLE_FRAME_PROCESSORS
@@ -40,37 +33,48 @@ JSharedArray::JSharedArray(const jni::alias_ref<JSharedArray::jhybridobject>& ja
 #else
   jsi::Runtime& runtime = *proxy->cthis()->getJSRuntime();
 #endif
-  TypedArrayKind kind = getTypedArrayKind(dataType);
-  __android_log_print(ANDROID_LOG_INFO, TAG, "Allocating ArrayBuffer with size %i and type %i...", size, dataType);
-  _array = std::make_shared<TypedArrayBase>(runtime, size, kind);
-  _byteBuffer = wrapInByteBuffer(runtime, _array);
-  _size = size;
+  __android_log_print(ANDROID_LOG_INFO, TAG, "Allocating ArrayBuffer with size %i...", byteBuffer->getDirectSize());
+  _byteBuffer = jni::make_global(byteBuffer);
+  _size = _byteBuffer->getDirectSize();
+
+  auto mutableByteBuffer = std::make_shared<MutableJByteBuffer>(byteBuffer);
+  _arrayBuffer = std::make_shared<jsi::ArrayBuffer>(runtime, std::move(mutableByteBuffer));
 }
 
+JSharedArray::JSharedArray(const jni::alias_ref<JSharedArray::jhybridobject>& javaThis,
+                           const jni::alias_ref<JVisionCameraProxy::javaobject>& proxy, int size)
+    : JSharedArray(javaThis, proxy, JByteBuffer::allocateDirect(size)) {}
+
 void JSharedArray::registerNatives() {
   registerHybrid({
-      makeNativeMethod("initHybrid", JSharedArray::initHybrid),
+      makeNativeMethod("initHybrid", JSharedArray::initHybridAllocate),
+      makeNativeMethod("initHybrid", JSharedArray::initHybridWrap),
       makeNativeMethod("getByteBuffer", JSharedArray::getByteBuffer),
       makeNativeMethod("getSize", JSharedArray::getSize),
   });
 }
 
-jni::local_ref<jni::JByteBuffer> JSharedArray::getByteBuffer() {
-  return jni::make_local(_byteBuffer);
+jni::global_ref<jni::JByteBuffer> JSharedArray::getByteBuffer() {
+  return _byteBuffer;
+}
+
+std::shared_ptr<jsi::ArrayBuffer> JSharedArray::getArrayBuffer() {
+  return _arrayBuffer;
 }
 
 jint JSharedArray::getSize() {
   return _size;
 }
 
-std::shared_ptr<TypedArrayBase> JSharedArray::getTypedArray() {
-  return _array;
+jni::local_ref<JSharedArray::jhybriddata>
+JSharedArray::initHybridAllocate(jni::alias_ref<jhybridobject> javaThis, jni::alias_ref<JVisionCameraProxy::javaobject> proxy, jint size) {
+  return makeCxxInstance(javaThis, proxy, size);
 }
 
-jni::local_ref<JSharedArray::jhybriddata> JSharedArray::initHybrid(jni::alias_ref<jhybridobject> javaThis,
-                                                                   jni::alias_ref<JVisionCameraProxy::javaobject> proxy, jint type,
-                                                                   jint size) {
-  return makeCxxInstance(javaThis, proxy, type, size);
+jni::local_ref<JSharedArray::jhybriddata> JSharedArray::initHybridWrap(jni::alias_ref<jhybridobject> javaThis,
+                                                                       jni::alias_ref<JVisionCameraProxy::javaobject> proxy,
+                                                                       jni::alias_ref<JByteBuffer> byteBuffer) {
+  return makeCxxInstance(javaThis, proxy, byteBuffer);
 }
 
 } // namespace vision

package/android/src/main/cpp/frameprocessor/java-bindings/JSharedArray.h

@@ -4,8 +4,8 @@
 
 #pragma once
 
-#include "JSITypedArray.h"
 #include "JVisionCameraProxy.h"
+#include "MutableJByteBuffer.h"
 #include <fbjni/ByteBuffer.h>
 #include <fbjni/fbjni.h>
 #include <jni.h>
@@ -20,30 +20,32 @@ public:
   static void registerNatives();
 
 public:
-  static jni::local_ref<JSharedArray::javaobject> create(jsi::Runtime& runtime, TypedArrayBase array);
+  static jni::local_ref<JSharedArray::javaobject> create(jsi::Runtime& runtime, jsi::ArrayBuffer arrayBuffer);
 
 public:
   jint getSize();
-  jni::local_ref<jni::JByteBuffer> getByteBuffer();
-  std::shared_ptr<TypedArrayBase> getTypedArray();
-
-private:
-  jni::global_ref<jni::JByteBuffer> wrapInByteBuffer(jsi::Runtime& runtime, std::shared_ptr<TypedArrayBase> typedArray);
+  jni::global_ref<jni::JByteBuffer> getByteBuffer();
+  std::shared_ptr<jsi::ArrayBuffer> getArrayBuffer();
 
 private:
   static auto constexpr TAG = "SharedArray";
   friend HybridBase;
   jni::global_ref<javaobject> _javaPart;
   jni::global_ref<jni::JByteBuffer> _byteBuffer;
-  std::shared_ptr<TypedArrayBase> _array;
+  std::shared_ptr<jsi::ArrayBuffer> _arrayBuffer;
   int _size;
 
 private:
-  explicit JSharedArray(jsi::Runtime& runtime, std::shared_ptr<TypedArrayBase> array);
+  explicit JSharedArray(jsi::Runtime& runtime, std::shared_ptr<jsi::ArrayBuffer> arrayBuffer);
   explicit JSharedArray(const jni::alias_ref<jhybridobject>& javaThis, const jni::alias_ref<JVisionCameraProxy::javaobject>& proxy,
-                        int dataType, int size);
-  static jni::local_ref<jhybriddata> initHybrid(jni::alias_ref<jhybridobject> javaThis,
-                                                jni::alias_ref<JVisionCameraProxy::javaobject> proxy, jint dataType, jint size);
+                        int size);
+  explicit JSharedArray(const jni::alias_ref<jhybridobject>& javaThis, const jni::alias_ref<JVisionCameraProxy::javaobject>& proxy,
+                        jni::alias_ref<JByteBuffer> byteBuffer);
+  static jni::local_ref<jhybriddata> initHybridAllocate(jni::alias_ref<jhybridobject> javaThis,
+                                                        jni::alias_ref<JVisionCameraProxy::javaobject> proxy, jint size);
+  static jni::local_ref<jhybriddata> initHybridWrap(jni::alias_ref<jhybridobject> javaThis,
+                                                    jni::alias_ref<JVisionCameraProxy::javaobject> proxy,
+                                                    jni::alias_ref<JByteBuffer> byteBuffer);
 };
 
 } // namespace vision

package/android/src/main/cpp/frameprocessor/java-bindings/JVisionCameraProxy.cpp

@@ -11,7 +11,6 @@
 #include <jsi/jsi.h>
 
 #include "FrameProcessorPluginHostObject.h"
-#include "JSITypedArray.h"
 
 #if VISION_CAMERA_ENABLE_FRAME_PROCESSORS
 #include <react-native-worklets-core/WKTJsiWorklet.h>
@@ -51,10 +50,7 @@ JVisionCameraProxy::JVisionCameraProxy(const jni::alias_ref<JVisionCameraProxy::
 
 JVisionCameraProxy::~JVisionCameraProxy() {
 #if VISION_CAMERA_ENABLE_FRAME_PROCESSORS
-  __android_log_write(ANDROID_LOG_INFO, TAG, "Destroying Context...");
-  // Destroy ArrayBuffer cache for both the JS and the Worklet Runtime.
-  invalidateArrayBufferCache(*_workletContext->getJsRuntime());
-  invalidateArrayBufferCache(_workletContext->getWorkletRuntime());
+  __android_log_write(ANDROID_LOG_INFO, TAG, "Destroying JVisionCameraProxy...");
 #endif
 }