react-native-wgpu 0.3.0 → 0.3.1

package/android/CMakeLists.txt

@@ -46,8 +46,7 @@ add_library(${PACKAGE_NAME} SHARED
     ../cpp/rnwgpu/RNWebGPUManager.cpp
     ../cpp/jsi/RNFPromise.cpp
     ../cpp/jsi/RNFHybridObject.cpp
-    ../cpp/jsi/RNFRuntimeCache.cpp
-    ../cpp/jsi/RNFWorkletRuntimeRegistry.cpp
+    ../cpp/jsi/RNFRuntimeState.cpp
     ../cpp/rnwgpu/async/AsyncRunner.cpp
     ../cpp/rnwgpu/async/AsyncTaskHandle.cpp
     ../cpp/rnwgpu/async/JSIMicrotaskDispatcher.cpp

package/cpp/jsi/RNFHybridObject.cpp

@@ -129,6 +129,7 @@ void HybridObject::ensureInitialized(facebook::jsi::Runtime& runtime) {
   if (!_didLoadMethods) {
     [[unlikely]];
     _creationRuntime = &runtime;
+    _runtimeState = rnwgpu::RNFRuntimeState::get(runtime);
     // lazy-load all exposed methods
     loadHybridMethods();
     _didLoadMethods = true;
@@ -136,11 +137,14 @@ void HybridObject::ensureInitialized(facebook::jsi::Runtime& runtime) {
 }
 
 bool HybridObject::isRuntimeAlive() {
-  if (_creationRuntime == nullptr) {
-    [[unlikely]];
-    return false;
+  return !_runtimeState.expired();
+}
+
+facebook::jsi::Runtime* HybridObject::getCreationRuntime() const {
+  if (_runtimeState.expired()) {
+    return nullptr;
   }
-  return RNFWorkletRuntimeRegistry::isRuntimeAlive(_creationRuntime);
+  return _creationRuntime;
 }
 
 } // namespace margelo

package/cpp/jsi/RNFHybridObject.h

@@ -5,7 +5,7 @@
 #pragma once
 
 #include "WGPULogger.h"
-#include "RNFWorkletRuntimeRegistry.h"
+#include "RNFRuntimeState.h"
 #include <functional>
 #include <jsi/jsi.h>
 #include <memory>
@@ -91,8 +91,9 @@ private:
 
 protected:
   const char* _name = TAG;
-  // Store a pointer to the runtime. Needed for checking if the runtime is still active, see WorkletRuntimeRegistry.
+  // Store a pointer to the runtime for convenience; ensure it is only used while the runtime state is alive.
   jsi::Runtime* _creationRuntime = nullptr;
+  std::weak_ptr<rnwgpu::RNFRuntimeState> _runtimeState;
 
 private:
   inline void ensureInitialized(facebook::jsi::Runtime& runtime);
@@ -128,6 +129,10 @@ private:
     };
   }
 
+protected:
+  facebook::jsi::Runtime* getCreationRuntime() const;
+  std::weak_ptr<rnwgpu::RNFRuntimeState> getRuntimeStateWeak() const { return _runtimeState; }
+
 protected:
   template <typename Derived, typename ReturnType, typename... Args>
   void registerHybridMethod(std::string name, ReturnType (Derived::*method)(Args...), Derived* derivedInstance, bool override = false) {

package/cpp/jsi/RNFJSIConverter.h

@@ -22,7 +22,6 @@
 #include "RNFEnumMapper.h"
 #include "RNFJSIHelper.h"
 #include "RNFPromise.h"
-#include "RNFWorkletRuntimeRegistry.h"
 
 #include "Unions.h"
 
@@ -225,52 +224,6 @@ template <> struct JSIConverter<rnwgpu::async::AsyncTaskHandle> {
   }
 };
 
-
-// [](Args...) -> T {} <> (Args...) => T
-template <typename ReturnType, typename... Args> struct JSIConverter<std::function<ReturnType(Args...)>> {
-  static std::function<ReturnType(Args...)> fromJSI(jsi::Runtime& runtime, const jsi::Value& arg, bool outOfBound) {
-    jsi::Function function = arg.asObject(runtime).asFunction(runtime);
-
-    std::shared_ptr<jsi::Function> sharedFunction = JSIHelper::createSharedJsiFunction(runtime, std::move(function));
-    return [&runtime, sharedFunction, outOfBound](Args... args) -> ReturnType {
-      jsi::Value result = sharedFunction->call(runtime, JSIConverter<std::decay_t<Args>>::toJSI(runtime, args)...);
-      if constexpr (std::is_same_v<ReturnType, void>) {
-        // it is a void function (returns undefined)
-        return;
-      } else {
-        // it returns a custom type, parse it from the JSI value.
-        return JSIConverter<ReturnType>::fromJSI(runtime, std::move(result), outOfBound);
-      }
-    };
-  }
-
-  template <size_t... Is>
-  static jsi::Value callHybridFunction(const std::function<ReturnType(Args...)>& function, jsi::Runtime& runtime, const jsi::Value* args,
-                                       std::index_sequence<Is...>, size_t count) {
-    if constexpr (std::is_same_v<ReturnType, void>) {
-      // it is a void function (will return undefined in JS)
-      function(JSIConverter<std::decay_t<Args>>::fromJSI(runtime, args[Is], Is >= count)...);
-      return jsi::Value::undefined();
-    } else {
-      // it is a custom type, parse it to a JS value
-      ReturnType result = function(JSIConverter<std::decay_t<Args>>::fromJSI(runtime, args[Is], Is >= count)...);
-      return JSIConverter<ReturnType>::toJSI(runtime, result);
-    }
-  }
-  static jsi::Value toJSI(jsi::Runtime& runtime, const std::function<ReturnType(Args...)>& function) {
-    jsi::HostFunctionType jsFunction = [function = std::move(function)](jsi::Runtime& runtime, const jsi::Value& thisValue,
-                                                                        const jsi::Value* args, size_t count) -> jsi::Value {
-      if (count != sizeof...(Args)) {
-        [[unlikely]];
-        throw jsi::JSError(runtime, "Function expected " + std::to_string(sizeof...(Args)) + " arguments, but received " +
-                                        std::to_string(count) + "!");
-      }
-      return callHybridFunction(function, runtime, args, std::index_sequence_for<Args...>{}, count);
-    };
-    return jsi::Function::createFromHostFunction(runtime, jsi::PropNameID::forUtf8(runtime, "hostFunction"), sizeof...(Args), jsFunction);
-  }
-};
-
 // std::map<std::string, T> <> Record<string, T>
 template <typename ValueType> struct JSIConverter<std::map<std::string, ValueType>> {
   static std::map<std::string, ValueType> fromJSI(jsi::Runtime& runtime, const jsi::Value& arg, bool outOfBound) {

package/cpp/jsi/RNFJSIHelper.h

@@ -4,7 +4,7 @@
 
 #pragma once
 
-#include "RNFWorkletRuntimeRegistry.h"
+#include "RNFRuntimeState.h"
 #include <jsi/jsi.h>
 #include <memory>
 #include <utility>
@@ -20,8 +20,10 @@ public:
    * Every jsi::Function you intend to share or hold should be wrapped using this function.
    */
   static std::shared_ptr<jsi::Function> createSharedJsiFunction(jsi::Runtime& runtime, jsi::Function&& function) {
-    std::shared_ptr<jsi::Function> sharedFunction(new jsi::Function(std::move(function)), [&runtime](jsi::Function* ptr) {
-      if (RNFWorkletRuntimeRegistry::isRuntimeAlive(&runtime)) {
+    auto runtimeState = rnwgpu::RNFRuntimeState::get(runtime);
+    std::weak_ptr<rnwgpu::RNFRuntimeState> runtimeStateWeak = runtimeState;
+    std::shared_ptr<jsi::Function> sharedFunction(new jsi::Function(std::move(function)), [runtimeStateWeak](jsi::Function* ptr) {
+      if (!runtimeStateWeak.expired()) {
         // Only delete the jsi::Function when the runtime it created is still alive.
         // Otherwise leak memory. We do this on purpose, as sometimes we would keep
         // references to JSI objects past the lifetime of its runtime (e.g.,

package/cpp/jsi/RNFRuntimeState.cpp

@@ -0,0 +1,18 @@
+#include "RNFRuntimeState.h"
+
+namespace rnwgpu {
+
+const facebook::jsi::UUID RNFRuntimeState::kRuntimeStateKey = facebook::jsi::UUID();
+
+std::shared_ptr<RNFRuntimeState> RNFRuntimeState::get(facebook::jsi::Runtime& runtime) {
+  auto existing = runtime.getRuntimeData(kRuntimeStateKey);
+  if (existing) {
+    return std::static_pointer_cast<RNFRuntimeState>(existing);
+  }
+
+  auto state = std::shared_ptr<RNFRuntimeState>(new RNFRuntimeState());
+  runtime.setRuntimeData(kRuntimeStateKey, state);
+  return state;
+}
+
+} // namespace rnwgpu

package/cpp/jsi/RNFRuntimeState.h

@@ -0,0 +1,106 @@
+#pragma once
+
+#include <jsi/jsi.h>
+
+#include <memory>
+#include <mutex>
+#include <unordered_map>
+#include <unordered_set>
+
+namespace rnwgpu {
+
+namespace jsi = facebook::jsi;
+
+/**
+ * Runtime state management using Hermes' runtime.setRuntimeData API.
+ * This replaces the old RuntimeLifecycleMonitor/RuntimeAwareCache system.
+ */
+class RNFRuntimeState {
+public:
+  // UUID key for storing our runtime state
+  static const jsi::UUID kRuntimeStateKey;
+
+  /**
+   * Get or create the runtime state for the given runtime
+   */
+  static std::shared_ptr<RNFRuntimeState> get(jsi::Runtime& runtime);
+
+  /**
+   * Template cache that can store any type T per object pointer
+   */
+  template <typename T>
+  class ObjectCache {
+  public:
+    std::shared_ptr<T> getOrCreate(void* object) {
+      std::lock_guard<std::mutex> lock(mutex_);
+      auto it = cache_.find(object);
+      if (it != cache_.end()) {
+        return it->second;
+      }
+
+      auto value = std::make_shared<T>();
+      cache_[object] = value;
+      return value;
+    }
+
+    void remove(void* object) {
+      std::lock_guard<std::mutex> lock(mutex_);
+      cache_.erase(object);
+    }
+
+    void clear() {
+      std::lock_guard<std::mutex> lock(mutex_);
+      cache_.clear();
+    }
+
+  private:
+    std::mutex mutex_;
+    std::unordered_map<void*, std::shared_ptr<T>> cache_;
+  };
+
+  /**
+   * Get or create a cache for a specific type T
+   */
+  template <typename T>
+  std::shared_ptr<ObjectCache<T>> getCache() {
+    std::lock_guard<std::mutex> lock(mutex_);
+
+    // Use type_info as key for the cache type
+    const std::type_info& typeId = typeid(T);
+    auto it = typeCaches_.find(&typeId);
+
+    if (it != typeCaches_.end()) {
+      return std::static_pointer_cast<ObjectCache<T>>(it->second);
+    }
+
+    auto cache = std::make_shared<ObjectCache<T>>();
+    typeCaches_[&typeId] = cache;
+    return cache;
+  }
+
+private:
+  RNFRuntimeState() = default;
+
+  std::mutex mutex_;
+  // Map from type_info to cache instance
+  std::unordered_map<const std::type_info*, std::shared_ptr<void>> typeCaches_;
+};
+
+/**
+ * Template helper for runtime-aware caching compatible with the old API
+ * This provides a migration path from RuntimeAwareCache
+ */
+template <typename T>
+class RuntimeAwareCache {
+public:
+  T& get(jsi::Runtime& rt) {
+    auto state = RNFRuntimeState::get(rt);
+    auto cache = state->getCache<T>();
+
+    // For compatibility, we use the runtime pointer as the object key
+    auto ptr = cache->getOrCreate(&rt);
+    return *ptr;
+  }
+};
+
+} // namespace rnwgpu

package/cpp/rnwgpu/api/GPUAdapter.cpp

@@ -79,18 +79,19 @@ async::AsyncTaskHandle GPUAdapter::requestDevice(
   std::string label =
       descriptor.has_value() ? descriptor.value()->label.value_or("") : "";
 
+  auto creationRuntime = getCreationRuntime();
   return _async->postTask(
       [this, aDescriptor, descriptor, label = std::move(label),
-       deviceLostBinding](
-          const async::AsyncTaskHandle::ResolveFunction &resolve,
-          const async::AsyncTaskHandle::RejectFunction &reject) {
+       deviceLostBinding,
+       creationRuntime](const async::AsyncTaskHandle::ResolveFunction &resolve,
+                        const async::AsyncTaskHandle::RejectFunction &reject) {
         (void)descriptor;
         _instance.RequestDevice(
             &aDescriptor, wgpu::CallbackMode::AllowProcessEvents,
-            [asyncRunner = _async, resolve, reject, label,
-             creationRuntime = _creationRuntime, deviceLostBinding](
-                wgpu::RequestDeviceStatus status, wgpu::Device device,
-                wgpu::StringView message) mutable {
+            [asyncRunner = _async, resolve, reject, label, creationRuntime,
+             deviceLostBinding](wgpu::RequestDeviceStatus status,
+                                wgpu::Device device,
+                                wgpu::StringView message) mutable {
               if (message.length) {
                 fprintf(stderr, "%s", message.data);
               }
@@ -106,6 +107,9 @@ async::AsyncTaskHandle GPUAdapter::requestDevice(
               device.SetLoggingCallback([creationRuntime](
                                             wgpu::LoggingType type,
                                             wgpu::StringView msg) {
+                if (creationRuntime == nullptr) {
+                  return;
+                }
                 const char *logLevel = "";
                 switch (type) {
                 case wgpu::LoggingType::Warning:

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "react-native-wgpu",
-  "version": "0.3.0",
+  "version": "0.3.1",
   "description": "React Native WebGPU",
   "main": "lib/commonjs/index",
   "module": "lib/module/index",

package/cpp/jsi/RNFRuntimeCache.cpp

@@ -1,57 +0,0 @@
-//
-// Created by Marc Rousavy on 22.02.24.
-//
-#include <jsi/jsi.h>
-
-#include <memory>
-#include <unordered_map>
-#include <unordered_set>
-#include <utility>
-
-#include "RNFRuntimeCache.h"
-
-namespace margelo {
-
-static std::unordered_map<jsi::Runtime*, std::unordered_set<RuntimeLifecycleListener*>> listeners;
-
-struct RuntimeLifecycleMonitorObject : public jsi::HostObject {
-  jsi::Runtime* _rt;
-  explicit RuntimeLifecycleMonitorObject(jsi::Runtime* rt) : _rt(rt) {}
-  ~RuntimeLifecycleMonitorObject() {
-    auto listenersSet = listeners.find(_rt);
-    if (listenersSet != listeners.end()) {
-      for (auto listener : listenersSet->second) {
-        listener->onRuntimeDestroyed(_rt);
-      }
-      listeners.erase(listenersSet);
-    }
-  }
-};
-
-void RuntimeLifecycleMonitor::addListener(jsi::Runtime& rt, RuntimeLifecycleListener* listener) {
-  auto listenersSet = listeners.find(&rt);
-  if (listenersSet == listeners.end()) {
-    // We install a global host object in the provided runtime, this way we can
-    // use that host object destructor to get notified when the runtime is being
-    // terminated. We use a unique name for the object as it gets saved with the
-    // runtime's global object.
-    rt.global().setProperty(rt, "__rnfl_rt_lifecycle_monitor",
-                            jsi::Object::createFromHostObject(rt, std::make_shared<RuntimeLifecycleMonitorObject>(&rt)));
-    std::unordered_set<RuntimeLifecycleListener*> newSet;
-    newSet.insert(listener);
-    listeners.emplace(&rt, std::move(newSet));
-  } else {
-    listenersSet->second.insert(listener);
-  }
-}
-
-void RuntimeLifecycleMonitor::removeListener(jsi::Runtime& rt, RuntimeLifecycleListener* listener) {
-  auto listenersSet = listeners.find(&rt);
-  if (listenersSet == listeners.end()) {
-    // nothing to do here
-  } else {
-    listenersSet->second.erase(listener);
-  }
-}
-
-} // namespace margelo

package/cpp/jsi/RNFRuntimeCache.h

@@ -1,79 +0,0 @@
-//
-// Created by Marc Rousavy on 22.02.24.
-//
-#pragma once
-
-#include <jsi/jsi.h>
-
-#include <memory>
-#include <unordered_map>
-#include <unordered_set>
-#include <utility>
-
-namespace margelo {
-
-namespace jsi = facebook::jsi;
-
-/**
- * Listener interface that allows for getting notified when a jsi::Runtime
- * instance is destroyed.
- */
-struct RuntimeLifecycleListener {
-  virtual ~RuntimeLifecycleListener() {}
-  virtual void onRuntimeDestroyed(jsi::Runtime*) = 0;
-};
-
-/**
- * This class provides an API via static methods for registering and
- * unregistering runtime lifecycle listeners. The listeners can be used to
- * cleanup any data that references a given jsi::Runtime instance before it gets
- * destroyed.
- */
-struct RuntimeLifecycleMonitor {
-  static void addListener(jsi::Runtime& rt, RuntimeLifecycleListener* listener);
-  static void removeListener(jsi::Runtime& rt, RuntimeLifecycleListener* listener);
-};
-
-/**
- * Provides a way to keep data specific to a jsi::Runtime instance that gets
- * cleaned up when that runtime is destroyed. This is necessary because JSI does
- * not allow for its associated objects to be retained past the runtime
- * lifetime. If an object (e.g. jsi::Values or jsi::Function instances) is kept
- * after the runtime is torn down, its destructor (once it is destroyed
- * eventually) will result in a crash (JSI objects keep a pointer to memory
- * managed by the runtime, accessing that portion of the memory after runtime is
- * deleted is the root cause of that crash).
- */
-template <typename T> class RuntimeAwareCache : public RuntimeLifecycleListener {
-
-public:
-  void onRuntimeDestroyed(jsi::Runtime* rt) override {
-    // A runtime has been destroyed, so destroy the related cache.
-    _runtimeCaches.erase(rt);
-  }
-
-  ~RuntimeAwareCache() {
-    for (auto& cache : _runtimeCaches) {
-      // remove all `onRuntimeDestroyed` listeners.
-      RuntimeLifecycleMonitor::removeListener(*cache.first, this);
-    }
-  }
-
-  T& get(jsi::Runtime& rt) {
-    if (_runtimeCaches.count(&rt) == 0) {
-      // This is the first time this Runtime has been accessed.
-      // We set up a `onRuntimeDestroyed` listener for it and
-      // initialize the cache map.
-      RuntimeLifecycleMonitor::addListener(rt, this);
-
-      T cache;
-      _runtimeCaches.emplace(&rt, std::move(cache));
-    }
-    return _runtimeCaches.at(&rt);
-  }
-
-private:
-  std::unordered_map<jsi::Runtime*, T> _runtimeCaches;
-};
-
-} // namespace margelo

package/cpp/jsi/RNFWorkletRuntimeCollector.h

@@ -1,43 +0,0 @@
-//
-// Created by Marc Rousavy on 22.02.24.
-//
-#pragma once
-
-#include "RNFWorkletRuntimeRegistry.h"
-
-#include <jsi/jsi.h>
-
-#include <memory>
-
-namespace margelo {
-
-// From:
-// https://github.com/software-mansion/react-native-reanimated/blob/6cb1a66f1a68cac8079de2b6b305d22359847e51/Common/cpp/ReanimatedRuntime/WorkletRuntimeCollector.h
-class WorkletRuntimeCollector : public jsi::HostObject {
-  // When worklet runtime is created, we inject an instance of this class as a
-  // `jsi::HostObject` into the global object. When worklet runtime is
-  // terminated, the object is garbage-collected, which runs the C++ destructor.
-  // In the destructor, we unregister the worklet runtime from the registry.
-
-public:
-  explicit WorkletRuntimeCollector(jsi::Runtime& runtime) : _runtime(runtime) {
-    Logger::log("WorkletRuntimeCollector", "Registering WorkletRuntime %p", &runtime);
-    RNFWorkletRuntimeRegistry::registerRuntime(_runtime);
-  }
-
-  ~WorkletRuntimeCollector() {
-    Logger::log("WorkletRuntimeCollector", "Unregistering WorkletRuntime %p", &_runtime);
-    RNFWorkletRuntimeRegistry::unregisterRuntime(_runtime);
-  }
-
-  static void install(jsi::Runtime& rt) {
-    auto collector = std::make_shared<WorkletRuntimeCollector>(rt);
-    auto object = jsi::Object::createFromHostObject(rt, collector);
-    rt.global().setProperty(rt, "__workletRuntimeCollector", object);
-  }
-
-private:
-  jsi::Runtime& _runtime;
-};
-
-} // namespace margelo

package/cpp/jsi/RNFWorkletRuntimeRegistry.cpp

@@ -1,12 +0,0 @@
-//
-// Created by Marc Rousavy on 22.02.24.
-//
-#include "RNFWorkletRuntimeRegistry.h"
-#include <set>
-
-namespace margelo {
-
-std::set<jsi::Runtime*> RNFWorkletRuntimeRegistry::registry_{};
-std::mutex RNFWorkletRuntimeRegistry::mutex_{};
-
-} // namespace margelo

package/cpp/jsi/RNFWorkletRuntimeRegistry.h

@@ -1,44 +0,0 @@
-//
-// Created by Marc Rousavy on 22.02.24.
-//
-#pragma once
-
-#include <jsi/jsi.h>
-
-#include <mutex>
-#include <set>
-
-namespace jsi = facebook::jsi;
-
-namespace margelo {
-
-// From:
-// https://github.com/software-mansion/react-native-reanimated/blob/6cb1a66f1a68cac8079de2b6b305d22359847e51/Common/cpp/ReanimatedRuntime/WorkletRuntimeRegistry.h
-class RNFWorkletRuntimeRegistry {
-private:
-  static std::set<jsi::Runtime*> registry_;
-  static std::mutex mutex_; // Protects `registry_`.
-
-  RNFWorkletRuntimeRegistry() {} // private ctor
-
-  static void registerRuntime(jsi::Runtime& runtime) {
-    std::lock_guard<std::mutex> lock(mutex_);
-    registry_.insert(&runtime);
-  }
-
-  static void unregisterRuntime(jsi::Runtime& runtime) {
-    std::lock_guard<std::mutex> lock(mutex_);
-    registry_.erase(&runtime);
-  }
-
-  friend class WorkletRuntimeCollector;
-
-public:
-  static bool isRuntimeAlive(jsi::Runtime* runtime) {
-    assert(runtime != nullptr);
-    std::lock_guard<std::mutex> lock(mutex_);
-    return registry_.find(runtime) != registry_.end();
-  }
-};
-
-} // namespace margelo