react-native-webgpu 0.10.1 → 0.10.3

package/android/src/main/java/com/webgpu/WebGPUView.java

@@ -1,6 +1,7 @@
 package com.webgpu;
 
 import android.content.Context;
+import android.os.Build;
 import android.view.Surface;
 import android.view.View;
 
@@ -10,12 +11,16 @@ import com.facebook.react.views.view.ReactViewGroup;
 
 public class WebGPUView extends ReactViewGroup implements WebGPUAPI {
 
+  // Backing view kinds, see updateView().
+  private static final int KIND_SURFACE_VIEW = 0;
+  private static final int KIND_TEXTURE_VIEW = 1;
+  private static final int KIND_HARDWARE_BUFFER_VIEW = 2;
 
   private int mContextId;
   private boolean mOpaque = true;
   private String mSurfaceType = "auto";
   private boolean mZOrderOnTop = false;
-  private boolean mAppliedTextureView;
+  private int mAppliedKind = -1;
   private boolean mAppliedZOrderOnTop;
   private WebGPUModule mModule;
   private View mView = null;
@@ -46,35 +51,62 @@ public class WebGPUView extends ReactViewGroup implements WebGPUAPI {
     mZOrderOnTop = value;
   }
 
+  // Resolve the backing view from the props. "auto" picks SurfaceView for an
+  // opaque canvas and TextureView for a non-opaque one. WebGPUHardwareBufferView
+  // (a plain View drawing AHardwareBuffers inline, API 29+) is opt-in via
+  // surfaceType until it gets a runtime fallback to TextureView when the
+  // AHardwareBuffer import fails on a device; below API 29 the request itself
+  // falls back to TextureView.
+  private int resolveKind() {
+    if ("SurfaceView".equals(mSurfaceType)) {
+      return KIND_SURFACE_VIEW;
+    }
+    if ("TextureView".equals(mSurfaceType)) {
+      return KIND_TEXTURE_VIEW;
+    }
+    if ("HardwareBufferView".equals(mSurfaceType)) {
+      boolean hardwareBufferSupported = Build.VERSION.SDK_INT >= Build.VERSION_CODES.Q;
+      return hardwareBufferSupported ? KIND_HARDWARE_BUFFER_VIEW : KIND_TEXTURE_VIEW;
+    }
+    return mOpaque ? KIND_SURFACE_VIEW : KIND_TEXTURE_VIEW;
+  }
+
   // Apply the complete prop transaction once, after contextId and all rendering
   // options have arrived. Only a change of backing view (or of zOrderOnTop,
   // which a SurfaceView must know before it attaches) replaces the child;
   // opacity is applied in place.
   public void updateView() {
-    boolean textureView = "TextureView".equals(mSurfaceType)
-      || (!"SurfaceView".equals(mSurfaceType) && !mOpaque);
-    boolean zOrderOnTop = !textureView && mZOrderOnTop;
-    if (mView == null || textureView != mAppliedTextureView
-        || zOrderOnTop != mAppliedZOrderOnTop) {
+    int kind = resolveKind();
+    boolean zOrderOnTop = kind == KIND_SURFACE_VIEW && mZOrderOnTop;
+    if (mView == null || kind != mAppliedKind || zOrderOnTop != mAppliedZOrderOnTop) {
       if (mView != null) {
         removeView(mView);
       }
-      mAppliedTextureView = textureView;
+      mAppliedKind = kind;
       mAppliedZOrderOnTop = zOrderOnTop;
       Context ctx = getContext();
-      mView = textureView
-        ? new WebGPUTextureView(ctx, this, mOpaque)
-        : new WebGPUSurfaceView(ctx, this, zOrderOnTop, mOpaque);
+      switch (kind) {
+        case KIND_HARDWARE_BUFFER_VIEW:
+          mView = new WebGPUHardwareBufferView(ctx, this);
+          break;
+        case KIND_TEXTURE_VIEW:
+          mView = new WebGPUTextureView(ctx, this, mOpaque);
+          break;
+        default:
+          mView = new WebGPUSurfaceView(ctx, this, zOrderOnTop, mOpaque);
+          break;
+      }
       addView(mView);
       // ReactViewGroup.requestLayout() is a deliberate no-op, so addView outside
       // the layout pass never lays the child out; do it by hand or the new view
       // stays 0x0 and never gets a surface.
       layoutChild();
-    } else if (textureView) {
+    } else if (kind == KIND_TEXTURE_VIEW) {
       ((WebGPUTextureView) mView).setOpaque(mOpaque);
-    } else {
+    } else if (kind == KIND_SURFACE_VIEW) {
       ((WebGPUSurfaceView) mView).setOpaque(mOpaque);
     }
+    // WebGPUHardwareBufferView always composites with alpha; opacity is a no-op.
   }
 
   private void layoutChild() {
@@ -110,6 +142,11 @@ public class WebGPUView extends ReactViewGroup implements WebGPUAPI {
     switchToOffscreenSurface(mContextId);
   }
 
+  @Override
+  public int getContextId() {
+    return mContextId;
+  }
+
   /**
    * Called from WebGPUViewManager.onDropViewInstance when React removes this
    * view: the view dies with its Canvas, so it retires the registry entry.

package/cpp/jsi/JSICache.cpp

@@ -1,5 +1,6 @@
 #include "JSICache.h"
 
+#include <algorithm>
 #include <iterator>
 #include <stdexcept>
 #include <string>
@@ -94,4 +95,26 @@ jsi::Object &JSICache::setPrototype(PrototypeKey key, jsi::Object prototype) {
   return it->second;
 }
 
+jsi::Value JSICache::lockWrapper(jsi::Runtime &runtime,
+                                 const void *nativeObject) {
+  auto it = _wrappers.find(nativeObject);
+  if (it == _wrappers.end()) {
+    return jsi::Value::undefined();
+  }
+  return it->second.lock(runtime);
+}
+
+void JSICache::setWrapper(jsi::Runtime &runtime, const void *nativeObject,
+                          const jsi::Object &wrapper) {
+  _wrappers.insert_or_assign(nativeObject, jsi::WeakObject(runtime, wrapper));
+  if (_wrappers.size() < _pruneWrappersAt) {
+    return;
+  }
+  for (auto it = _wrappers.begin(); it != _wrappers.end();) {
+    it = it->second.lock(runtime).isUndefined() ? _wrappers.erase(it)
+                                                : std::next(it);
+  }
+  _pruneWrappersAt = std::max(kMinPruneWrappersAt, _wrappers.size() * 2);
+}
+
 } // namespace rnwgpu

package/cpp/jsi/JSICache.h

@@ -13,7 +13,8 @@ namespace jsi = facebook::jsi;
 
 /**
  * Per-runtime store for JSI values that native code wants to keep across
- * calls (currently: the prototype object of every NativeObject class).
+ * calls: the prototype object of every NativeObject class, and the JS
+ * wrapper of every native object that has one on the runtime.
  *
  * Ownership model (borrowed from react-native-nitro-modules' JSICache):
  *
@@ -72,6 +73,26 @@ public:
    */
   jsi::Object &setPrototype(PrototypeKey key, jsi::Object prototype);
 
+  /**
+   * Returns the JS wrapper stored for `nativeObject` on this runtime, or
+   * undefined if none was stored or the wrapper has been collected.
+   *
+   * The key is the native object's address. A live wrapper keeps its native
+   * object alive (the object is the wrapper's NativeState), so as long as
+   * the entry locks it still belongs to that object. Once the wrapper has
+   * been collected the entry is stale and is skipped or replaced, so an
+   * address reused by a new object can never resolve to the old wrapper.
+   */
+  jsi::Value lockWrapper(jsi::Runtime &runtime, const void *nativeObject);
+
+  /**
+   * Stores `wrapper` as the JS wrapper of `nativeObject` on this runtime.
+   * Only a weak reference is kept: the wrapper stays collectable, and the
+   * native object is released when it goes away.
+   */
+  void setWrapper(jsi::Runtime &runtime, const void *nativeObject,
+                  const jsi::Object &wrapper);
+
   JSICache() = default;
   ~JSICache() override = default;
   JSICache(const JSICache &) = delete;
@@ -81,6 +102,12 @@ private:
   static std::shared_ptr<JSICache> getOrCreateOnGlobal(jsi::Runtime &runtime);
 
   std::unordered_map<PrototypeKey, jsi::Object> _prototypes;
+  // Native object -> its wrapper on this runtime. Entries of collected
+  // wrappers are pruned in setWrapper once the map has doubled since the
+  // last pruning, so it stays proportional to the number of live wrappers.
+  std::unordered_map<const void *, jsi::WeakObject> _wrappers;
+  size_t _pruneWrappersAt = kMinPruneWrappersAt;
+  static constexpr size_t kMinPruneWrappersAt = 64;
 
   // Runtime -> cache. Weak on purpose: the runtime's global owns the cache.
   static std::mutex _registryMutex;

package/cpp/jsi/NativeObject.h

@@ -34,44 +34,51 @@ namespace jsi = facebook::jsi;
 template <typename Derived> class NativeObject;
 
 /**
- * Registry for NativeObject prototype installers.
- * This allows BoxedWebGPUObject::unbox() to install prototypes on any runtime
- * by looking up the brand name and calling the appropriate installer.
+ * Registry mapping a NativeObject class brand (CLASS_NAME) to a reconstructor
+ * that rebuilds a fully functional JS object (prototype + native state) on
+ * any runtime. This is what lets BoxedWebGPUObject::unbox() bring objects to
+ * worklet runtimes.
  */
 class NativeObjectRegistry {
 public:
-  using InstallerFunc = std::function<void(jsi::Runtime &)>;
+  using Reconstructor = std::function<jsi::Value(
+      jsi::Runtime &, std::shared_ptr<jsi::NativeState>)>;
 
   static NativeObjectRegistry &getInstance() {
     static NativeObjectRegistry instance;
     return instance;
   }
 
-  void registerInstaller(const std::string &brand, InstallerFunc installer) {
+  void registerClass(const std::string &brand, Reconstructor reconstructor) {
     std::lock_guard<std::mutex> lock(_mutex);
-    if (_installers.count(brand) != 0) {
+    if (_reconstructors.count(brand) != 0) {
       // The brand is the key unbox() uses to rebuild objects on worklet
       // runtimes - a duplicate would let one class hijack another's unboxing.
       throw std::runtime_error("Duplicate native object brand registered: " +
                                brand);
     }
-    _installers[brand] = std::move(installer);
+    _reconstructors[brand] = std::move(reconstructor);
   }
 
-  bool installPrototype(jsi::Runtime &runtime, const std::string &brand) {
-    std::lock_guard<std::mutex> lock(_mutex);
-    auto it = _installers.find(brand);
-    if (it != _installers.end()) {
-      it->second(runtime);
-      return true;
+  jsi::Value reconstruct(jsi::Runtime &runtime, const std::string &brand,
+                         std::shared_ptr<jsi::NativeState> state) {
+    Reconstructor reconstructor;
+    {
+      std::lock_guard<std::mutex> lock(_mutex);
+      auto it = _reconstructors.find(brand);
+      if (it == _reconstructors.end()) {
+        throw jsi::JSError(runtime,
+                           "No native class registered for brand: " + brand);
+      }
+      reconstructor = it->second;
     }
-    return false;
+    return reconstructor(runtime, std::move(state));
   }
 
 private:
   NativeObjectRegistry() = default;
   std::mutex _mutex;
-  std::unordered_map<std::string, InstallerFunc> _installers;
+  std::unordered_map<std::string, Reconstructor> _reconstructors;
 };
 
 /**
@@ -88,7 +95,8 @@ private:
  * Usage pattern with registerCustomSerializable:
  * - pack(): Call WebGPU.box(obj) to create a BoxedWebGPUObject (HostObject)
  * - The HostObject is serialized by Worklets and transferred to UI runtime
- * - unpack(): Call boxed.unbox() to get back the original object with prototype
+ * - unpack(): Call boxed.unbox() to get the object's wrapper on that runtime
+ *   (cached per runtime, see NativeObject::create), prototype included
  *
  * This is similar to NitroModules.box()/unbox() pattern.
  */
@@ -101,37 +109,15 @@ public:
   jsi::Value get(jsi::Runtime &runtime, const jsi::PropNameID &name) override {
     auto propName = name.utf8(runtime);
     if (propName == "unbox") {
+      auto state = _nativeState;
+      auto brand = _brand;
       return jsi::Function::createFromHostFunction(
           runtime, jsi::PropNameID::forUtf8(runtime, "unbox"), 0,
-          [this](jsi::Runtime &rt, const jsi::Value & /*thisVal*/,
-                 const jsi::Value * /*args*/,
-                 size_t /*count*/) -> jsi::Value {
-            // Try to get the prototype from the global constructor
-            auto ctor = rt.global().getProperty(rt, _brand.c_str());
-            if (!ctor.isObject()) {
-              // Constructor doesn't exist on this runtime - install it
-              NativeObjectRegistry::getInstance().installPrototype(rt, _brand);
-              ctor = rt.global().getProperty(rt, _brand.c_str());
-            }
-
-            // Create a new object and attach the native state
-            jsi::Object obj(rt);
-            obj.setNativeState(rt, _nativeState);
-
-            // Set the prototype if constructor exists
-            if (ctor.isObject()) {
-              auto ctorObj = ctor.getObject(rt);
-              auto proto = ctorObj.getProperty(rt, "prototype");
-              if (proto.isObject()) {
-                auto objectCtor =
-                    rt.global().getPropertyAsObject(rt, "Object");
-                auto setPrototypeOf =
-                    objectCtor.getPropertyAsFunction(rt, "setPrototypeOf");
-                setPrototypeOf.call(rt, obj, proto);
-              }
-            }
-
-            return std::move(obj);
+          [state, brand](jsi::Runtime &rt, const jsi::Value & /*thisVal*/,
+                         const jsi::Value * /*args*/,
+                         size_t /*count*/) -> jsi::Value {
+            return NativeObjectRegistry::getInstance().reconstruct(rt, brand,
+                                                                   state);
           });
     }
     if (propName == "__boxedWebGPU") {
@@ -305,15 +291,38 @@ public:
    * created internally by the native code).
    *
    * Also registers this class with NativeObjectRegistry so that
-   * BoxedWebGPUObject::unbox() can install prototypes on secondary runtimes.
+   * BoxedWebGPUObject::unbox() can rebuild its objects on secondary runtimes.
    */
   static void installConstructor(jsi::Runtime &runtime) {
-    // Register this class's installer in the registry (only needs to happen once)
+    // Register the reconstructor used by BoxedWebGPUObject::unbox() to
+    // rebuild this object (prototype + native state) on another runtime.
     static std::once_flag registryFlag;
     std::call_once(registryFlag, []() {
-      NativeObjectRegistry::getInstance().registerInstaller(
+      NativeObjectRegistry::getInstance().registerClass(
           Derived::CLASS_NAME,
-          [](jsi::Runtime &rt) { Derived::installConstructor(rt); });
+          [](jsi::Runtime &rt,
+             std::shared_ptr<jsi::NativeState> state) -> jsi::Value {
+            auto instance = std::dynamic_pointer_cast<Derived>(state);
+            if (instance == nullptr) {
+              throw jsi::JSError(rt, "Invalid boxed native object state");
+            }
+            // The constructor makes `instanceof` work on this runtime too.
+            if (!rt.global().hasProperty(rt, Derived::CLASS_NAME)) {
+              Derived::installConstructor(rt);
+            }
+            // Unboxing hands back the wrapper cached for the target runtime,
+            // creating it on first use. Keep the runtime the object was
+            // originally created on: async native code (GPUDevice error/lost
+            // events, GPUAdapter::requestDevice) delivers into the creation
+            // runtime, and rebinding it to a worklet runtime would invoke
+            // main-runtime jsi::Functions on the wrong runtime and thread.
+            auto *originalRuntime = instance->getCreationRuntime();
+            auto value = Derived::create(rt, instance);
+            if (originalRuntime != nullptr) {
+              instance->setCreationRuntime(originalRuntime);
+            }
+            return value;
+          });
     });
 
     installPrototype(runtime);
@@ -345,10 +354,36 @@ public:
   }
 
   /**
-   * Create a JS object with native state attached.
+   * Returns the JS wrapper of `instance` on `runtime`, creating it on first
+   * use.
+   *
+   * A native object has at most one live wrapper per runtime. The wrapper is
+   * cached (weakly, so it stays collectable) in the runtime's JSICache, and
+   * converting the same object again (returning it from another native
+   * call, unboxing it in a worklet on every frame, ...) hands back that
+   * wrapper instead of a new one. Each wrapper reports the native memory the
+   * object owns to the GC of its runtime through setExternalMemoryPressure,
+   * so the memory is charged exactly once per runtime that can reach the
+   * object, however often it is converted; charging every conversion would
+   * multiply the amount by the number of captures in a worklet until Hermes
+   * hit its max heap size. The hint is refreshed on every round trip for
+   * objects whose size changes after creation.
    */
   static jsi::Value create(jsi::Runtime &runtime,
                            std::shared_ptr<Derived> instance) {
+    auto &cache = JSICache::get(runtime);
+    const void *key = instance.get();
+
+    auto existing = cache.lockWrapper(runtime, key);
+    if (existing.isObject()) {
+      auto pressure = instance->getMemoryPressure();
+      if (pressure > 0) {
+        existing.getObject(runtime).setExternalMemoryPressure(runtime,
+                                                              pressure);
+      }
+      return existing;
+    }
+
     installPrototype(runtime);
 
     // Store creation runtime for logging etc.
@@ -374,6 +409,7 @@ public:
       obj.setExternalMemoryPressure(runtime, pressure);
     }
 
+    cache.setWrapper(runtime, key, obj);
     return std::move(obj);
   }