feat: Add zero-copy SharedArray type to Frame Processor Plugins (#2383)

* feat: Create `TypedArray` class for Frame Processor Plugins

* Type

* feat: Pass `VisionCameraProxy` along (BREAKING)

* feat: Finish implementation

* Log a bit

* feat: Successfully convert JSI <> JNI buffers

* Wrap buffer

* fix: Fix using wrong Runtime

* feat: Add docs

* add zero copy example

* Format C++

* Create iOS base

* feat: Finish iOS implementation

* chore: Format

* fix: Use `NSData` instead of `NSMutableData`

* Format

* fix: Fix build when Frame Processors are disabled

* chore: Rename `TypedArray` to `SharedArray`

* fix: Fix Swift typings for Array

* Remove a few default inits

* fix: Fix Android build

* fix: Use `NSInteger`

* Update SharedArray.mm

* fix: Expose bytes directly on iOS (NSData was immutable)

package/android/CMakeLists.txt

@@ -34,6 +34,7 @@ add_library(
         src/main/cpp/frameprocessor/FrameProcessorPluginHostObject.cpp
         src/main/cpp/frameprocessor/JSIJNIConversion.cpp
         src/main/cpp/frameprocessor/VisionCameraProxy.cpp
+        src/main/cpp/frameprocessor/java-bindings/JSharedArray.cpp
         src/main/cpp/frameprocessor/java-bindings/JFrame.cpp
         src/main/cpp/frameprocessor/java-bindings/JFrameProcessor.cpp
         src/main/cpp/frameprocessor/java-bindings/JFrameProcessorPlugin.cpp

package/android/src/main/cpp/VisionCamera.cpp

@@ -1,4 +1,5 @@
 #include "JFrameProcessor.h"
+#include "JSharedArray.h"
 #include "JVisionCameraProxy.h"
 #include "JVisionCameraScheduler.h"
 #include "VideoPipeline.h"
@@ -14,6 +15,7 @@ JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM* vm, void*) {
     vision::VideoPipeline::registerNatives();
 #if VISION_CAMERA_ENABLE_FRAME_PROCESSORS
     vision::JFrameProcessor::registerNatives();
+    vision::JSharedArray::registerNatives();
 #endif
   });
 }

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

@@ -15,6 +15,8 @@
 
 #include "FrameHostObject.h"
 #include "JFrame.h"
+#include "JSITypedArray.h"
+#include "JSharedArray.h"
 
 namespace vision {
 
@@ -56,8 +58,24 @@ jni::local_ref<jobject> JSIJNIConversion::convertJSIValueToJNIObject(jsi::Runtim
         arrayList->add(jniItem);
       }
       return arrayList;
+    } else if (valueAsObject.isArrayBuffer(runtime)) {
+      // ArrayBuffer/TypedArray
+
+      TypedArrayBase array = getTypedArray(runtime, valueAsObject);
+      return JSharedArray::create(runtime, std::move(array));
+
     } else if (valueAsObject.isHostObject(runtime)) {
-      throw std::runtime_error("You can't pass HostObjects here.");
+
+      if (valueAsObject.isHostObject<FrameHostObject>(runtime)) {
+        // Frame
+
+        auto frame = valueAsObject.getHostObject<FrameHostObject>(runtime);
+        return jni::make_local(frame->frame);
+
+      } else {
+        throw std::runtime_error("The given HostObject is not supported by a Frame Processor Plugin.");
+      }
+
     } else {
       // Map<String, Object>
 
@@ -75,7 +93,7 @@ jni::local_ref<jobject> JSIJNIConversion::convertJSIValueToJNIObject(jsi::Runtim
     }
   } else {
     auto stringRepresentation = value.toString(runtime).utf8(runtime);
-    throw std::runtime_error("Failed to convert jsi::Value to JNI value - unsupported type!" + stringRepresentation);
+    throw std::runtime_error("Failed to convert jsi::Value to JNI value - unsupported type! " + stringRepresentation);
   }
 }
 
@@ -154,6 +172,12 @@ jsi::Value JSIJNIConversion::convertJNIObjectToJSIValue(jsi::Runtime& runtime, c
     // box into HostObject
     auto hostObject = std::make_shared<FrameHostObject>(frame);
     return jsi::Object::createFromHostObject(runtime, hostObject);
+  } else if (object->isInstanceOf(JSharedArray::javaClassStatic())) {
+    // SharedArray
+    auto sharedArray = static_ref_cast<JSharedArray::javaobject>(object);
+
+    std::shared_ptr<TypedArrayBase> array = sharedArray->cthis()->getTypedArray();
+    return array->getBuffer(runtime);
   }
 
   auto type = object->getClass()->toString();

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

@@ -0,0 +1,69 @@
+//
+// Created by Marc Rousavy on 12.01.24.
+//
+
+#include "JSharedArray.h"
+#include <android/log.h>
+
+namespace vision {
+
+using namespace facebook;
+
+TypedArrayKind getTypedArrayKind(int unsafeEnumValue) {
+  return static_cast<TypedArrayKind>(unsafeEnumValue);
+}
+
+jni::local_ref<JSharedArray::javaobject> JSharedArray::create(jsi::Runtime& runtime, TypedArrayBase array) {
+  return newObjectCxxArgs(runtime, std::make_shared<TypedArrayBase>(std::move(array)));
+}
+
+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);
+}
+
+JSharedArray::JSharedArray(const jni::alias_ref<JSharedArray::jhybridobject>& javaThis,
+                           const jni::alias_ref<JVisionCameraProxy::javaobject>& proxy, int dataType, int size) {
+  _javaPart = jni::make_global(javaThis);
+
+#if VISION_CAMERA_ENABLE_FRAME_PROCESSORS
+  jsi::Runtime& runtime = proxy->cthis()->getWorkletRuntime();
+#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);
+}
+
+void JSharedArray::registerNatives() {
+  registerHybrid({
+      makeNativeMethod("initHybrid", JSharedArray::initHybrid),
+      makeNativeMethod("getByteBuffer", JSharedArray::getByteBuffer),
+  });
+}
+
+jni::local_ref<jni::JByteBuffer> JSharedArray::getByteBuffer() {
+  return jni::make_local(_byteBuffer);
+}
+
+std::shared_ptr<TypedArrayBase> JSharedArray::getTypedArray() {
+  return _array;
+}
+
+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);
+}
+
+} // namespace vision

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

@@ -0,0 +1,47 @@
+//
+// Created by Marc Rousavy on 12.01.24.
+//
+
+#pragma once
+
+#include "JSITypedArray.h"
+#include "JVisionCameraProxy.h"
+#include <fbjni/ByteBuffer.h>
+#include <fbjni/fbjni.h>
+#include <jni.h>
+
+namespace vision {
+
+using namespace facebook;
+
+class JSharedArray : public jni::HybridClass<JSharedArray> {
+public:
+  static auto constexpr kJavaDescriptor = "Lcom/mrousavy/camera/frameprocessor/TypedArray;";
+  static void registerNatives();
+
+public:
+  static jni::local_ref<JSharedArray::javaobject> create(jsi::Runtime& runtime, TypedArrayBase array);
+
+public:
+  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);
+
+private:
+  static auto constexpr TAG = "TypedArray";
+  friend HybridBase;
+  jni::global_ref<javaobject> _javaPart;
+  jni::global_ref<jni::JByteBuffer> _byteBuffer;
+  std::shared_ptr<TypedArrayBase> _array;
+
+private:
+  explicit JSharedArray(jsi::Runtime& runtime, std::shared_ptr<TypedArrayBase> array);
+  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);
+};
+
+} // namespace vision

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

@@ -37,6 +37,12 @@ public:
     return _runtime;
   }
 
+#if VISION_CAMERA_ENABLE_FRAME_PROCESSORS
+  jsi::Runtime& getWorkletRuntime() {
+    return _workletContext->getWorkletRuntime();
+  }
+#endif
+
 private:
   friend HybridBase;
   jni::global_ref<JVisionCameraProxy::javaobject> _javaPart;

package/android/src/main/java/com/mrousavy/camera/frameprocessor/FrameProcessorPlugin.java

@@ -23,9 +23,8 @@ public abstract class FrameProcessorPlugin {
      * The initializer of this Frame Processor Plugin.
      * This is called everytime this Frame Processor Plugin is loaded from the JS side (`initFrameProcessorPlugin(..)`).
      * Optionally override this method to implement custom initialization logic.
-     * @param options An options dictionary passed from the JS side, or null if none.
      */
-    public FrameProcessorPlugin(@Nullable Map<String, Object> options) {}
+    public FrameProcessorPlugin() { }
 
     /**
      * The actual Frame Processor Plugin's implementation that runs when `plugin.call(..)` is called in the JS Frame Processor.

package/android/src/main/java/com/mrousavy/camera/frameprocessor/FrameProcessorPluginRegistry.java

@@ -2,6 +2,7 @@ package com.mrousavy.camera.frameprocessor;
 
 import android.util.Log;
 import androidx.annotation.Keep;
+import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import com.facebook.proguard.annotations.DoNotStrip;
 import java.util.Map;
@@ -24,7 +25,7 @@ public class FrameProcessorPluginRegistry {
 
     @DoNotStrip
     @Keep
-    public static FrameProcessorPlugin getPlugin(String name, Map<String, Object> options) {
+    public static FrameProcessorPlugin getPlugin(String name, VisionCameraProxy proxy, Map<String, Object> options) {
         Log.i(TAG, "Looking up Frame Processor Plugin \"" + name + "\"...");
         PluginInitializer initializer = Plugins.get(name);
         if (initializer == null) {
@@ -32,10 +33,10 @@ public class FrameProcessorPluginRegistry {
             return null;
         }
         Log.i(TAG, "Frame Processor Plugin \"" + name + "\" found! Initializing...");
-        return initializer.initializePlugin(options);
+        return initializer.initializePlugin(proxy, options);
     }
 
     public interface PluginInitializer {
-        FrameProcessorPlugin initializePlugin(@Nullable Map<String, Object> options);
+        FrameProcessorPlugin initializePlugin(@NonNull VisionCameraProxy proxy, @Nullable Map<String, Object> options);
     }
 }

package/android/src/main/java/com/mrousavy/camera/frameprocessor/SharedArray.java

@@ -0,0 +1,59 @@
+package com.mrousavy.camera.frameprocessor;
+
+import androidx.annotation.Keep;
+
+import com.facebook.jni.HybridData;
+import com.facebook.proguard.annotations.DoNotStrip;
+
+import java.nio.ByteBuffer;
+
+/**
+ * A JSI TypedArray/ArrayBuffer implementation used for passing buffers between JS and Native without copying data.
+ * ByteBuffers are used for efficient data transfer.
+ *
+ * @noinspection JavaJniMissingFunction
+ */
+public final class SharedArray {
+    @DoNotStrip
+    @Keep
+    private final HybridData mHybridData;
+
+    @DoNotStrip
+    @Keep
+    public SharedArray(HybridData hybridData) {
+        mHybridData = hybridData;
+    }
+
+    /**
+     * Allocate a new SharedArray. Use `getByteBuffer` to obtain a reference to the direct ByteBuffer for writing.
+     * @param proxy The VisionCamera Proxy from the Frame Processor Plugin's initializer.
+     * @param dataType The ArrayBuffer's data type. `Type.Int8Array` = `Int8Array` in JS
+     * @param size The size of the ArrayBuffer.
+     */
+    public SharedArray(VisionCameraProxy proxy, Type dataType, int size) {
+        mHybridData = initHybrid(proxy, dataType.ordinal(), size);
+    }
+
+    /**
+     * Gets the direct ByteBuffer that can be used to directly update the JSI ArrayBuffer.
+     */
+    public native ByteBuffer getByteBuffer();
+
+    private native HybridData initHybrid(VisionCameraProxy proxy, int dataType, int size);
+
+    /**
+     * The Type of the SharedArray.
+     */
+    public enum Type {
+        // Values start at 0 and need to match with JSITypedArray.h::TypedArrayKind
+        Int8Array,
+        Int16Array,
+        Int32Array,
+        Uint8Array,
+        Uint8ClampedArray,
+        Uint16Array,
+        Uint32Array,
+        Float32Array,
+        Float64Array,
+    }
+}

package/android/src/main/java/com/mrousavy/camera/frameprocessor/VisionCameraProxy.kt

@@ -72,7 +72,7 @@ class VisionCameraProxy(context: ReactApplicationContext) {
   @DoNotStrip
   @Keep
   fun initFrameProcessorPlugin(name: String, options: Map<String, Any>): FrameProcessorPlugin =
-    FrameProcessorPluginRegistry.getPlugin(name, options)
+    FrameProcessorPluginRegistry.getPlugin(name, this, options)
 
   // private C++ funcs
   private external fun initHybrid(jsContext: Long, jsCallInvokerHolder: CallInvokerHolderImpl, scheduler: VisionCameraScheduler): HybridData