refactor(Worklets): move UI scheduling

packages/react-native-reanimated/Common/cpp/reanimated/NativeModules/ReanimatedModuleProxy.cpp

@@ -62,13 +62,6 @@ ReanimatedModuleProxy::ReanimatedModuleProxy(
       isReducedMotion_(isReducedMotion),
       workletsModuleProxy_(workletsModuleProxy),
       valueUnpackerCode_(workletsModuleProxy->getValueUnpackerCode()),
-      uiWorkletRuntime_(std::make_shared<WorkletRuntime>(
-          rnRuntime,
-          workletsModuleProxy->getJSQueue(),
-          workletsModuleProxy->getJSScheduler(),
-          "Reanimated UI runtime",
-          true /* supportsLocking */,
-          valueUnpackerCode_)),
       eventHandlerRegistry_(std::make_unique<EventHandlerRegistry>()),
       requestRender_(platformDepMethodsHolder.requestRender),
       onRenderCallback_([this](const double timestampMs) {
@@ -160,7 +153,8 @@ void ReanimatedModuleProxy::commonInit(
   };
 #endif
 
-  jsi::Runtime &uiRuntime = uiWorkletRuntime_->getJSIRuntime();
+  jsi::Runtime &uiRuntime =
+      workletsModuleProxy_->getUIWorkletRuntime()->getJSIRuntime();
   UIRuntimeDecorator::decorate(
       uiRuntime,
 #ifdef RCT_NEW_ARCH_ENABLED
@@ -197,56 +191,6 @@ ReanimatedModuleProxy::~ReanimatedModuleProxy() {
 #ifdef RCT_NEW_ARCH_ENABLED
   operationsInBatch_.clear();
 #endif // RCT_NEW_ARCH_ENABLED
-  uiWorkletRuntime_.reset();
-}
-
-void ReanimatedModuleProxy::scheduleOnUI(
-    jsi::Runtime &rt,
-    const jsi::Value &worklet) {
-  auto shareableWorklet = extractShareableOrThrow<ShareableWorklet>(
-      rt, worklet, "[Reanimated] Only worklets can be scheduled to run on UI.");
-  workletsModuleProxy_->getUIScheduler()->scheduleOnUI(COPY_CAPTURE_WITH_THIS {
-#if JS_RUNTIME_HERMES
-    // JSI's scope defined here allows for JSI-objects to be cleared up
-    // after each runtime loop. Within these loops we typically create some
-    // temporary JSI objects and hence it allows for such objects to be
-    // garbage collected much sooner. Apparently the scope API is only
-    // supported on Hermes at the moment.
-    const auto scope = jsi::Scope(uiWorkletRuntime_->getJSIRuntime());
-#endif
-    uiWorkletRuntime_->runGuarded(shareableWorklet);
-  });
-}
-
-jsi::Value ReanimatedModuleProxy::executeOnUIRuntimeSync(
-    jsi::Runtime &rt,
-    const jsi::Value &worklet) {
-  return uiWorkletRuntime_->executeSync(rt, worklet);
-}
-
-jsi::Value ReanimatedModuleProxy::createWorkletRuntime(
-    jsi::Runtime &rt,
-    const jsi::Value &name,
-    const jsi::Value &initializer) {
-  auto workletRuntime = std::make_shared<WorkletRuntime>(
-      rt,
-      workletsModuleProxy_->getJSQueue(),
-      workletsModuleProxy_->getJSScheduler(),
-      name.asString(rt).utf8(rt),
-      false /* supportsLocking */,
-      valueUnpackerCode_);
-  auto initializerShareable = extractShareableOrThrow<ShareableWorklet>(
-      rt, initializer, "[Reanimated] Initializer must be a worklet.");
-  workletRuntime->runGuarded(initializerShareable);
-  return jsi::Object::createFromHostObject(rt, workletRuntime);
-}
-
-jsi::Value ReanimatedModuleProxy::scheduleOnRuntime(
-    jsi::Runtime &rt,
-    const jsi::Value &workletRuntimeValue,
-    const jsi::Value &shareableWorkletValue) {
-  reanimated::scheduleOnRuntime(rt, workletRuntimeValue, shareableWorkletValue);
-  return jsi::Value::undefined();
 }
 
 jsi::Value ReanimatedModuleProxy::registerEventHandler(
@@ -348,7 +292,8 @@ jsi::Value ReanimatedModuleProxy::getViewProp(
       callback.getObject(rnRuntime).asFunction(rnRuntime));
   const auto shadowNode = shadowNodeFromValue(rnRuntime, shadowNodeWrapper);
   workletsModuleProxy_->getUIScheduler()->scheduleOnUI(COPY_CAPTURE_WITH_THIS {
-    jsi::Runtime &uiRuntime = uiWorkletRuntime_->getJSIRuntime();
+    jsi::Runtime &uiRuntime =
+        workletsModuleProxy_->getUIWorkletRuntime()->getJSIRuntime();
     const auto resultStr =
         obtainPropFromShadowNode(uiRuntime, propNameStr, shadowNode);
 
@@ -379,7 +324,8 @@ jsi::Value ReanimatedModuleProxy::getViewProp(
       COPY_CAPTURE_WITH_THIS
 
       () {
-        jsi::Runtime &uiRuntime = uiWorkletRuntime_->getJSIRuntime();
+        jsi::Runtime &uiRuntime =
+            workletsModuleProxy_->getUIWorkletRuntime()->getJSIRuntime();
         const jsi::Value propNameValue =
             jsi::String::createFromUtf8(uiRuntime, propNameStr);
         const auto resultValue =
@@ -488,15 +434,17 @@ void ReanimatedModuleProxy::requestAnimationFrame(
 void ReanimatedModuleProxy::maybeRequestRender() {
   if (!renderRequested_) {
     renderRequested_ = true;
-    jsi::Runtime &uiRuntime = uiWorkletRuntime_->getJSIRuntime();
+    jsi::Runtime &uiRuntime =
+        workletsModuleProxy_->getUIWorkletRuntime()->getJSIRuntime();
     requestRender_(onRenderCallback_, uiRuntime);
   }
 }
 
 void ReanimatedModuleProxy::onRender(double timestampMs) {
   auto callbacks = std::move(frameCallbacks_);
   frameCallbacks_.clear();
-  jsi::Runtime &uiRuntime = uiWorkletRuntime_->getJSIRuntime();
+  jsi::Runtime &uiRuntime =
+      workletsModuleProxy_->getUIWorkletRuntime()->getJSIRuntime();
   jsi::Value timestamp{timestampMs};
   for (const auto &callback : callbacks) {
     runOnRuntimeGuarded(uiRuntime, *callback, timestamp);
@@ -511,7 +459,7 @@ jsi::Value ReanimatedModuleProxy::registerSensor(
     const jsi::Value &sensorDataHandler) {
   return animatedSensorModule_.registerSensor(
       rt,
-      uiWorkletRuntime_,
+      workletsModuleProxy_->getUIWorkletRuntime(),
       sensorType,
       interval,
       iosReferenceFrame,
@@ -575,7 +523,11 @@ bool ReanimatedModuleProxy::handleEvent(
     const jsi::Value &payload,
     double currentTime) {
   eventHandlerRegistry_->processEvent(
-      uiWorkletRuntime_, currentTime, eventName, emitterReactTag, payload);
+      workletsModuleProxy_->getUIWorkletRuntime(),
+      currentTime,
+      eventName,
+      emitterReactTag,
+      payload);
 
   // TODO: return true if Reanimated successfully handled the event
   // to avoid sending it to JavaScript
@@ -601,7 +553,8 @@ bool ReanimatedModuleProxy::handleRawEvent(
   if (eventType.rfind("top", 0) == 0) {
     eventType = "on" + eventType.substr(3);
   }
-  jsi::Runtime &rt = uiWorkletRuntime_->getJSIRuntime();
+  jsi::Runtime &rt =
+      workletsModuleProxy_->getUIWorkletRuntime()->getJSIRuntime();
   const auto &eventPayload = rawEvent.eventPayload;
   jsi::Value payload = eventPayload->asJSIValue(rt);
 
@@ -638,7 +591,8 @@ void ReanimatedModuleProxy::performOperations() {
   auto copiedOperationsQueue = std::move(operationsInBatch_);
   operationsInBatch_.clear();
 
-  jsi::Runtime &rt = uiWorkletRuntime_->getJSIRuntime();
+  jsi::Runtime &rt =
+      workletsModuleProxy_->getUIWorkletRuntime()->getJSIRuntime();
 
   {
     auto lock = propsRegistry_->createLock();
@@ -778,7 +732,8 @@ jsi::String ReanimatedModuleProxy::obtainProp(
     jsi::Runtime &rt,
     const jsi::Value &shadowNodeWrapper,
     const jsi::Value &propName) {
-  jsi::Runtime &uiRuntime = uiWorkletRuntime_->getJSIRuntime();
+  jsi::Runtime &uiRuntime =
+      workletsModuleProxy_->getUIWorkletRuntime()->getJSIRuntime();
   const auto propNameStr = propName.asString(rt).utf8(rt);
   const auto shadowNode = shadowNodeFromValue(rt, shadowNodeWrapper);
   const auto resultStr =
@@ -856,7 +811,7 @@ void ReanimatedModuleProxy::initializeLayoutAnimationsProxy() {
         layoutAnimationsManager_,
         componentDescriptorRegistry,
         scheduler->getContextContainer(),
-        uiWorkletRuntime_->getJSIRuntime(),
+        workletsModuleProxy_->getUIWorkletRuntime()->getJSIRuntime(),
         workletsModuleProxy_->getUIScheduler());
   }
 }
@@ -876,7 +831,7 @@ jsi::Value ReanimatedModuleProxy::subscribeForKeyboardEvents(
       COPY_CAPTURE_WITH_THIS
 
       (int keyboardState, int height) {
-        uiWorkletRuntime_->runGuarded(
+        workletsModuleProxy_->getUIWorkletRuntime()->runGuarded(
             shareableHandler, jsi::Value(keyboardState), jsi::Value(height));
       },
       isStatusBarTranslucent.getBool(),

packages/react-native-reanimated/Common/cpp/reanimated/NativeModules/ReanimatedModuleProxy.h

@@ -42,19 +42,6 @@ class ReanimatedModuleProxy : public ReanimatedModuleProxySpec {
 
   ~ReanimatedModuleProxy();
 
-  void scheduleOnUI(jsi::Runtime &rt, const jsi::Value &worklet) override;
-  jsi::Value executeOnUIRuntimeSync(jsi::Runtime &rt, const jsi::Value &worklet)
-      override;
-
-  jsi::Value createWorkletRuntime(
-      jsi::Runtime &rt,
-      const jsi::Value &name,
-      const jsi::Value &initializer) override;
-  jsi::Value scheduleOnRuntime(
-      jsi::Runtime &rt,
-      const jsi::Value &workletRuntimeValue,
-      const jsi::Value &shareableWorkletValue) override;
-
   jsi::Value registerEventHandler(
       jsi::Runtime &rt,
       const jsi::Value &worklet,
@@ -161,10 +148,6 @@ class ReanimatedModuleProxy : public ReanimatedModuleProxySpec {
     return *layoutAnimationsManager_;
   }
 
-  [[nodiscard]] inline jsi::Runtime &getUIRuntime() const {
-    return uiWorkletRuntime_->getJSIRuntime();
-  }
-
   [[nodiscard]] inline bool isBridgeless() const {
     return isBridgeless_;
   }
@@ -194,7 +177,6 @@ class ReanimatedModuleProxy : public ReanimatedModuleProxySpec {
   const bool isReducedMotion_;
   const std::shared_ptr<WorkletsModuleProxy> workletsModuleProxy_;
   const std::string valueUnpackerCode_;
-  std::shared_ptr<WorkletRuntime> uiWorkletRuntime_;
 
   std::unique_ptr<EventHandlerRegistry> eventHandlerRegistry_;
   const RequestRenderFunction requestRender_;

packages/react-native-reanimated/Common/cpp/reanimated/NativeModules/ReanimatedModuleProxySpec.cpp

@@ -7,45 +7,6 @@
 
 namespace reanimated {
 
-// scheduler
-
-static jsi::Value REANIMATED_SPEC_PREFIX(scheduleOnUI)(
-    jsi::Runtime &rt,
-    TurboModule &turboModule,
-    const jsi::Value *args,
-    size_t) {
-  static_cast<ReanimatedModuleProxySpec *>(&turboModule)
-      ->scheduleOnUI(rt, std::move(args[0]));
-  return jsi::Value::undefined();
-}
-
-static jsi::Value REANIMATED_SPEC_PREFIX(executeOnUIRuntimeSync)(
-    jsi::Runtime &rt,
-    TurboModule &turboModule,
-    const jsi::Value *args,
-    size_t) {
-  return static_cast<ReanimatedModuleProxySpec *>(&turboModule)
-      ->executeOnUIRuntimeSync(rt, std::move(args[0]));
-}
-
-static jsi::Value REANIMATED_SPEC_PREFIX(createWorkletRuntime)(
-    jsi::Runtime &rt,
-    TurboModule &turboModule,
-    const jsi::Value *args,
-    size_t) {
-  return static_cast<ReanimatedModuleProxySpec *>(&turboModule)
-      ->createWorkletRuntime(rt, std::move(args[0]), std::move(args[1]));
-}
-
-static jsi::Value REANIMATED_SPEC_PREFIX(scheduleOnRuntime)(
-    jsi::Runtime &rt,
-    TurboModule &turboModule,
-    const jsi::Value *args,
-    size_t) {
-  return static_cast<ReanimatedModuleProxySpec *>(&turboModule)
-      ->scheduleOnRuntime(rt, std::move(args[0]), std::move(args[1]));
-}
-
 static jsi::Value REANIMATED_SPEC_PREFIX(registerEventHandler)(
     jsi::Runtime &rt,
     TurboModule &turboModule,
@@ -163,15 +124,6 @@ static jsi::Value REANIMATED_SPEC_PREFIX(setShouldAnimateExiting)(
 ReanimatedModuleProxySpec::ReanimatedModuleProxySpec(
     const std::shared_ptr<CallInvoker> &jsInvoker)
     : TurboModule("NativeReanimated", jsInvoker) {
-  methodMap_["scheduleOnUI"] =
-      MethodMetadata{1, REANIMATED_SPEC_PREFIX(scheduleOnUI)};
-  methodMap_["executeOnUIRuntimeSync"] =
-      MethodMetadata{1, REANIMATED_SPEC_PREFIX(executeOnUIRuntimeSync)};
-  methodMap_["createWorkletRuntime"] =
-      MethodMetadata{2, REANIMATED_SPEC_PREFIX(createWorkletRuntime)};
-  methodMap_["scheduleOnRuntime"] =
-      MethodMetadata{2, REANIMATED_SPEC_PREFIX(scheduleOnRuntime)};
-
   methodMap_["registerEventHandler"] =
       MethodMetadata{3, REANIMATED_SPEC_PREFIX(registerEventHandler)};
   methodMap_["unregisterEventHandler"] =

packages/react-native-reanimated/Common/cpp/reanimated/NativeModules/ReanimatedModuleProxySpec.h

@@ -18,22 +18,6 @@ class JSI_EXPORT ReanimatedModuleProxySpec : public TurboModule {
       const std::shared_ptr<CallInvoker> &jsInvoker);
 
  public:
-  // Scheduling
-  virtual void scheduleOnUI(jsi::Runtime &rt, const jsi::Value &worklet) = 0;
-  virtual jsi::Value executeOnUIRuntimeSync(
-      jsi::Runtime &rt,
-      const jsi::Value &worklet) = 0;
-
-  // Worklet runtime
-  virtual jsi::Value createWorkletRuntime(
-      jsi::Runtime &rt,
-      const jsi::Value &name,
-      const jsi::Value &initializer) = 0;
-  virtual jsi::Value scheduleOnRuntime(
-      jsi::Runtime &rt,
-      const jsi::Value &workletRuntimeValue,
-      const jsi::Value &shareableWorkletValue) = 0;
-
   // events
   virtual jsi::Value registerEventHandler(
       jsi::Runtime &rt,

packages/react-native-reanimated/Common/cpp/reanimated/RuntimeDecorators/RNRuntimeDecorator.cpp

@@ -5,8 +5,8 @@ namespace reanimated {
 
 void RNRuntimeDecorator::decorate(
     jsi::Runtime &rnRuntime,
+    jsi::Runtime &uiRuntime,
     const std::shared_ptr<ReanimatedModuleProxy> &reanimatedModuleProxy) {
-  jsi::Runtime &uiRuntime = reanimatedModuleProxy->getUIRuntime();
   auto workletRuntimeValue =
       rnRuntime.global()
           .getPropertyAsObject(rnRuntime, "ArrayBuffer")

packages/react-native-reanimated/Common/cpp/reanimated/RuntimeDecorators/RNRuntimeDecorator.h

@@ -14,6 +14,7 @@ class RNRuntimeDecorator {
  public:
   static void decorate(
       jsi::Runtime &rnRuntime,
+      jsi::Runtime &uiRuntime,
       const std::shared_ptr<ReanimatedModuleProxy> &reanimatedModuleProxy);
 };