refactor(Worklets): move UIWorkletRuntime (#6792)

## Summary

Reanimated no longer creates UI Worklet Runtime. Now it's created by
Worklets and Reanimated only borrows it.

## Test plan

See that apps are working and zooming 👍

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,7 +191,6 @@ ReanimatedModuleProxy::~ReanimatedModuleProxy() {
 #ifdef RCT_NEW_ARCH_ENABLED
   operationsInBatch_.clear();
 #endif // RCT_NEW_ARCH_ENABLED
-  uiWorkletRuntime_.reset();
 }
 
 void ReanimatedModuleProxy::scheduleOnUI(
@@ -212,16 +205,17 @@ void ReanimatedModuleProxy::scheduleOnUI(
     // 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());
+    const auto scope = jsi::Scope(
+        workletsModuleProxy_->getUIWorkletRuntime()->getJSIRuntime());
 #endif
-    uiWorkletRuntime_->runGuarded(shareableWorklet);
+    workletsModuleProxy_->getUIWorkletRuntime()->runGuarded(shareableWorklet);
   });
 }
 
 jsi::Value ReanimatedModuleProxy::executeOnUIRuntimeSync(
     jsi::Runtime &rt,
     const jsi::Value &worklet) {
-  return uiWorkletRuntime_->executeSync(rt, worklet);
+  return workletsModuleProxy_->getUIWorkletRuntime()->executeSync(rt, worklet);
 }
 
 jsi::Value ReanimatedModuleProxy::createWorkletRuntime(
@@ -348,7 +342,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 +374,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 +484,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 +509,7 @@ jsi::Value ReanimatedModuleProxy::registerSensor(
     const jsi::Value &sensorDataHandler) {
   return animatedSensorModule_.registerSensor(
       rt,
-      uiWorkletRuntime_,
+      workletsModuleProxy_->getUIWorkletRuntime(),
       sensorType,
       interval,
       iosReferenceFrame,
@@ -575,7 +573,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 +603,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 +641,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 +782,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 +861,7 @@ void ReanimatedModuleProxy::initializeLayoutAnimationsProxy() {
         layoutAnimationsManager_,
         componentDescriptorRegistry,
         scheduler->getContextContainer(),
-        uiWorkletRuntime_->getJSIRuntime(),
+        workletsModuleProxy_->getUIWorkletRuntime()->getJSIRuntime(),
         workletsModuleProxy_->getUIScheduler());
   }
 }
@@ -876,7 +881,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

@@ -161,10 +161,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 +190,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/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);
 };
 

packages/react-native-reanimated/Common/cpp/worklets/NativeModules/WorkletsModuleProxy.cpp

@@ -19,6 +19,7 @@ using namespace facebook;
 namespace worklets {
 
 WorkletsModuleProxy::WorkletsModuleProxy(
+    jsi::Runtime &rnRuntime,
     const std::string &valueUnpackerCode,
     const std::shared_ptr<MessageQueueThread> &jsQueue,
     const std::shared_ptr<CallInvoker> &jsCallInvoker,
@@ -28,9 +29,18 @@ WorkletsModuleProxy::WorkletsModuleProxy(
       valueUnpackerCode_(valueUnpackerCode),
       jsQueue_(jsQueue),
       jsScheduler_(jsScheduler),
-      uiScheduler_(uiScheduler) {}
-
-WorkletsModuleProxy::~WorkletsModuleProxy() {}
+      uiScheduler_(uiScheduler),
+      uiWorkletRuntime_(std::make_shared<WorkletRuntime>(
+          rnRuntime,
+          jsQueue,
+          jsScheduler,
+          "Reanimated UI runtime",
+          true /* supportsLocking */,
+          valueUnpackerCode_)) {}
+
+WorkletsModuleProxy::~WorkletsModuleProxy() {
+  uiWorkletRuntime_.reset();
+}
 
 jsi::Value WorkletsModuleProxy::makeShareableClone(
     jsi::Runtime &rt,

packages/react-native-reanimated/Common/cpp/worklets/NativeModules/WorkletsModuleProxy.h

@@ -14,6 +14,7 @@ namespace worklets {
 class WorkletsModuleProxy : public WorkletsModuleProxySpec {
  public:
   explicit WorkletsModuleProxy(
+      jsi::Runtime &rnRuntime,
       const std::string &valueUnpackerCode,
       const std::shared_ptr<MessageQueueThread> &jsQueue,
       const std::shared_ptr<CallInvoker> &jsCallInvoker,
@@ -44,11 +45,17 @@ class WorkletsModuleProxy : public WorkletsModuleProxySpec {
     return uiScheduler_;
   }
 
+  [[nodiscard]] inline std::shared_ptr<WorkletRuntime> getUIWorkletRuntime()
+      const {
+    return uiWorkletRuntime_;
+  }
+
  private:
   const std::string valueUnpackerCode_;
   const std::shared_ptr<MessageQueueThread> jsQueue_;
   const std::shared_ptr<JSScheduler> jsScheduler_;
   const std::shared_ptr<UIScheduler> uiScheduler_;
+  std::shared_ptr<WorkletRuntime> uiWorkletRuntime_;
 #ifndef NDEBUG
   SingleInstanceChecker<WorkletsModuleProxy> singleInstanceChecker_;
 #endif // NDEBUG

packages/react-native-reanimated/android/src/main/cpp/reanimated/android/NativeProxy.cpp

@@ -41,6 +41,7 @@ NativeProxy::NativeProxy(
     )
     : javaPart_(jni::make_global(jThis)),
       rnRuntime_(rnRuntime),
+      workletsModuleProxy_(workletsModuleProxy),
       reanimatedModuleProxy_(std::make_shared<ReanimatedModuleProxy>(
           workletsModuleProxy,
           *rnRuntime,
@@ -155,7 +156,10 @@ void NativeProxy::injectCppVersion() {
 void NativeProxy::installJSIBindings() {
   jsi::Runtime &rnRuntime = *rnRuntime_;
   WorkletRuntimeCollector::install(rnRuntime);
-  RNRuntimeDecorator::decorate(rnRuntime, reanimatedModuleProxy_);
+  RNRuntimeDecorator::decorate(
+      rnRuntime,
+      workletsModuleProxy_->getUIWorkletRuntime()->getJSIRuntime(),
+      reanimatedModuleProxy_);
 #ifndef NDEBUG
   checkJavaVersion(rnRuntime);
   injectCppVersion();
@@ -407,7 +411,8 @@ void NativeProxy::handleEvent(
     return;
   }
 
-  jsi::Runtime &rt = reanimatedModuleProxy_->getUIRuntime();
+  jsi::Runtime &rt =
+      workletsModuleProxy_->getUIWorkletRuntime()->getJSIRuntime();
   jsi::Value payload;
   try {
     payload = jsi::Value::createFromJsonUtf8(
@@ -505,34 +510,43 @@ PlatformDepMethodsHolder NativeProxy::getPlatformDependentMethods() {
 void NativeProxy::setupLayoutAnimations() {
   auto weakReanimatedModuleProxy =
       std::weak_ptr<ReanimatedModuleProxy>(reanimatedModuleProxy_);
+  auto weakWorkletsModuleProxy =
+      std::weak_ptr<WorkletsModuleProxy>(workletsModuleProxy_);
 
   layoutAnimations_->cthis()->setAnimationStartingBlock(
-      [weakReanimatedModuleProxy](
+      [weakReanimatedModuleProxy, weakWorkletsModuleProxy](
           int tag, int type, alias_ref<JMap<jstring, jstring>> values) {
         if (auto reanimatedModuleProxy = weakReanimatedModuleProxy.lock()) {
-          jsi::Runtime &rt = reanimatedModuleProxy->getUIRuntime();
-          jsi::Object yogaValues(rt);
-          for (const auto &entry : *values) {
-            try {
-              std::string keyString = entry.first->toStdString();
-              std::string valueString = entry.second->toStdString();
-              auto key = jsi::String::createFromAscii(rt, keyString);
-              if (keyString == "currentTransformMatrix" ||
-                  keyString == "targetTransformMatrix") {
-                jsi::Array matrix =
-                    jsi_utils::convertStringToArray(rt, valueString, 9);
-                yogaValues.setProperty(rt, key, matrix);
-              } else {
-                auto value = stod(valueString);
-                yogaValues.setProperty(rt, key, value);
+          if (auto workletsModuleProxy = weakWorkletsModuleProxy.lock()) {
+            jsi::Runtime &rt =
+                workletsModuleProxy->getUIWorkletRuntime()->getJSIRuntime();
+            jsi::Object yogaValues(rt);
+            for (const auto &entry : *values) {
+              try {
+                std::string keyString = entry.first->toStdString();
+                std::string valueString = entry.second->toStdString();
+                auto key = jsi::String::createFromAscii(rt, keyString);
+                if (keyString == "currentTransformMatrix" ||
+                    keyString == "targetTransformMatrix") {
+                  jsi::Array matrix =
+                      jsi_utils::convertStringToArray(rt, valueString, 9);
+                  yogaValues.setProperty(rt, key, matrix);
+                } else {
+                  auto value = stod(valueString);
+                  yogaValues.setProperty(rt, key, value);
+                }
+              } catch (std::invalid_argument e) {
+                throw std::runtime_error(
+                    "[Reanimated] Failed to convert value to number.");
               }
-            } catch (std::invalid_argument e) {
-              throw std::runtime_error(
-                  "[Reanimated] Failed to convert value to number.");
             }
+            reanimatedModuleProxy->layoutAnimationsManager()
+                .startLayoutAnimation(
+                    rt,
+                    tag,
+                    static_cast<LayoutAnimationType>(type),
+                    yogaValues);
           }
-          reanimatedModuleProxy->layoutAnimationsManager().startLayoutAnimation(
-              rt, tag, static_cast<LayoutAnimationType>(type), yogaValues);
         }
       });
 
@@ -573,11 +587,14 @@ void NativeProxy::setupLayoutAnimations() {
       });
 
   layoutAnimations_->cthis()->setCancelAnimationForTag(
-      [weakReanimatedModuleProxy](int tag) {
+      [weakReanimatedModuleProxy, weakWorkletsModuleProxy](int tag) {
         if (auto reanimatedModuleProxy = weakReanimatedModuleProxy.lock()) {
-          jsi::Runtime &rt = reanimatedModuleProxy->getUIRuntime();
-          reanimatedModuleProxy->layoutAnimationsManager()
-              .cancelLayoutAnimation(rt, tag);
+          if (auto workletsModuleProxy = weakWorkletsModuleProxy.lock()) {
+            jsi::Runtime &rt =
+                workletsModuleProxy->getUIWorkletRuntime()->getJSIRuntime();
+            reanimatedModuleProxy->layoutAnimationsManager()
+                .cancelLayoutAnimation(rt, tag);
+          }
         }
       });
 

packages/react-native-reanimated/android/src/main/cpp/reanimated/android/NativeProxy.h

@@ -168,6 +168,7 @@ class NativeProxy : public jni::HybridClass<NativeProxy> {
   friend HybridBase;
   jni::global_ref<NativeProxy::javaobject> javaPart_;
   jsi::Runtime *rnRuntime_;
+  std::shared_ptr<WorkletsModuleProxy> workletsModuleProxy_;
   std::shared_ptr<ReanimatedModuleProxy> reanimatedModuleProxy_;
   jni::global_ref<LayoutAnimations::javaobject> layoutAnimations_;
 #ifndef NDEBUG

packages/react-native-reanimated/android/src/main/cpp/worklets/android/WorkletsModule.cpp

@@ -28,6 +28,7 @@ WorkletsModule::WorkletsModule(
     : javaPart_(jni::make_global(jThis)),
       rnRuntime_(rnRuntime),
       workletsModuleProxy_(std::make_shared<WorkletsModuleProxy>(
+          *rnRuntime,
           valueUnpackerCode,
           std::make_shared<JMessageQueueThread>(messageQueueThread),
           jsCallInvoker,