refactor(Worklets): move JSScheduler

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

@@ -53,22 +53,21 @@ namespace reanimated {
 ReanimatedModuleProxy::ReanimatedModuleProxy(
     const std::shared_ptr<WorkletsModuleProxy> &workletsModuleProxy,
     jsi::Runtime &rnRuntime,
-    const std::shared_ptr<JSScheduler> &jsScheduler,
+    const std::shared_ptr<CallInvoker> &jsCallInvoker,
     const std::shared_ptr<UIScheduler> &uiScheduler,
     const PlatformDepMethodsHolder &platformDepMethodsHolder,
     const bool isBridgeless,
     const bool isReducedMotion)
-    : ReanimatedModuleProxySpec(jsScheduler->getJSCallInvoker()),
+    : ReanimatedModuleProxySpec(jsCallInvoker),
       isBridgeless_(isBridgeless),
       isReducedMotion_(isReducedMotion),
       workletsModuleProxy_(workletsModuleProxy),
-      jsScheduler_(jsScheduler),
       uiScheduler_(uiScheduler),
       valueUnpackerCode_(workletsModuleProxy->getValueUnpackerCode()),
       uiWorkletRuntime_(std::make_shared<WorkletRuntime>(
           rnRuntime,
           workletsModuleProxy->getJSQueue(),
-          jsScheduler_,
+          workletsModuleProxy->getJSScheduler(),
           "Reanimated UI runtime",
           true /* supportsLocking */,
           valueUnpackerCode_)),
@@ -79,7 +78,8 @@ ReanimatedModuleProxy::ReanimatedModuleProxy(
         onRender(timestampMs);
       }),
       animatedSensorModule_(platformDepMethodsHolder),
-      jsLogger_(std::make_shared<JSLogger>(jsScheduler_)),
+      jsLogger_(
+          std::make_shared<JSLogger>(workletsModuleProxy->getJSScheduler())),
       layoutAnimationsManager_(
           std::make_shared<LayoutAnimationsManager>(jsLogger_)),
 #ifdef RCT_NEW_ARCH_ENABLED
@@ -233,7 +233,7 @@ jsi::Value ReanimatedModuleProxy::createWorkletRuntime(
   auto workletRuntime = std::make_shared<WorkletRuntime>(
       rt,
       workletsModuleProxy_->getJSQueue(),
-      jsScheduler_,
+      workletsModuleProxy_->getJSScheduler(),
       name.asString(rt).utf8(rt),
       false /* supportsLocking */,
       valueUnpackerCode_);
@@ -349,16 +349,17 @@ jsi::Value ReanimatedModuleProxy::getViewProp(
   const auto funPtr = std::make_shared<jsi::Function>(
       callback.getObject(rnRuntime).asFunction(rnRuntime));
   const auto shadowNode = shadowNodeFromValue(rnRuntime, shadowNodeWrapper);
-  uiScheduler_->scheduleOnUI([=]() {
+  uiScheduler_->scheduleOnUI(COPY_CAPTURE_WITH_THIS {
     jsi::Runtime &uiRuntime = uiWorkletRuntime_->getJSIRuntime();
     const auto resultStr =
         obtainPropFromShadowNode(uiRuntime, propNameStr, shadowNode);
 
-    jsScheduler_->scheduleOnJS([=](jsi::Runtime &rnRuntime) {
-      const auto resultValue =
-          jsi::String::createFromUtf8(rnRuntime, resultStr);
-      funPtr->call(rnRuntime, resultValue);
-    });
+    workletsModuleProxy_->getJSScheduler()->scheduleOnJS(
+        [=](jsi::Runtime &rnRuntime) {
+          const auto resultValue =
+              jsi::String::createFromUtf8(rnRuntime, resultStr);
+          funPtr->call(rnRuntime, resultValue);
+        });
   });
   return jsi::Value::undefined();
 }
@@ -386,8 +387,8 @@ jsi::Value ReanimatedModuleProxy::getViewProp(
         const auto resultValue =
             obtainPropFunction_(uiRuntime, viewTagInt, propNameValue);
         const auto resultStr = resultValue.asString(uiRuntime).utf8(uiRuntime);
-
-        jsScheduler_->scheduleOnJS([=](jsi::Runtime &rnRuntime) {
+        const auto jsScheduler = workletsModuleProxy_->getJSScheduler();
+        jsScheduler->scheduleOnJS([=](jsi::Runtime &rnRuntime) {
           const auto resultValue =
               jsi::String::createFromUtf8(rnRuntime, resultStr);
           funPtr->call(rnRuntime, resultValue);

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

@@ -35,7 +35,7 @@ class ReanimatedModuleProxy : public ReanimatedModuleProxySpec {
   ReanimatedModuleProxy(
       const std::shared_ptr<WorkletsModuleProxy> &workletsModuleProxy,
       jsi::Runtime &rnRuntime,
-      const std::shared_ptr<JSScheduler> &jsScheduler,
+      const std::shared_ptr<CallInvoker> &jsCallInvoker,
       const std::shared_ptr<UIScheduler> &uiScheduler,
       const PlatformDepMethodsHolder &platformDepMethodsHolder,
       const bool isBridgeless,
@@ -194,7 +194,6 @@ class ReanimatedModuleProxy : public ReanimatedModuleProxySpec {
   const bool isBridgeless_;
   const bool isReducedMotion_;
   const std::shared_ptr<WorkletsModuleProxy> workletsModuleProxy_;
-  const std::shared_ptr<JSScheduler> jsScheduler_;
   const std::shared_ptr<UIScheduler> uiScheduler_;
   const std::string valueUnpackerCode_;
   std::shared_ptr<WorkletRuntime> uiWorkletRuntime_;

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

@@ -20,10 +20,13 @@ namespace worklets {
 
 WorkletsModuleProxy::WorkletsModuleProxy(
     const std::string &valueUnpackerCode,
-    const std::shared_ptr<MessageQueueThread> &jsQueue)
-    : WorkletsModuleProxySpec(nullptr),
+    const std::shared_ptr<MessageQueueThread> &jsQueue,
+    const std::shared_ptr<CallInvoker> &jsCallInvoker,
+    const std::shared_ptr<JSScheduler> &jsScheduler)
+    : WorkletsModuleProxySpec(jsCallInvoker),
       valueUnpackerCode_(valueUnpackerCode),
-      jsQueue_(jsQueue) {}
+      jsQueue_(jsQueue),
+      jsScheduler_(jsScheduler) {}
 
 WorkletsModuleProxy::~WorkletsModuleProxy() {}
 

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

@@ -2,6 +2,7 @@
 
 #include <cxxreact/MessageQueueThread.h>
 #include <worklets/NativeModules/WorkletsModuleProxySpec.h>
+#include <worklets/Tools/JSScheduler.h>
 #include <worklets/Tools/SingleInstanceChecker.h>
 #include <worklets/WorkletRuntime/WorkletRuntime.h>
 #include <memory>
@@ -13,7 +14,9 @@ class WorkletsModuleProxy : public WorkletsModuleProxySpec {
  public:
   explicit WorkletsModuleProxy(
       const std::string &valueUnpackerCode,
-      const std::shared_ptr<MessageQueueThread> &jsQueue);
+      const std::shared_ptr<MessageQueueThread> &jsQueue,
+      const std::shared_ptr<CallInvoker> &jsCallInvoker,
+      const std::shared_ptr<JSScheduler> &jsScheduler);
 
   ~WorkletsModuleProxy();
 
@@ -31,9 +34,14 @@ class WorkletsModuleProxy : public WorkletsModuleProxySpec {
     return jsQueue_;
   }
 
+  [[nodiscard]] inline std::shared_ptr<JSScheduler> getJSScheduler() const {
+    return jsScheduler_;
+  }
+
  private:
   const std::string valueUnpackerCode_;
   const std::shared_ptr<MessageQueueThread> jsQueue_;
+  const std::shared_ptr<JSScheduler> jsScheduler_;
 #ifndef NDEBUG
   SingleInstanceChecker<WorkletsModuleProxy> singleInstanceChecker_;
 #endif // NDEBUG

packages/react-native-reanimated/Common/cpp/worklets/Tools/JSScheduler.cpp

@@ -17,24 +17,4 @@ JSScheduler::JSScheduler(
       rnRuntime_(rnRuntime),
       jsCallInvoker_(jsCallInvoker) {}
 
-#ifdef RCT_NEW_ARCH_ENABLED
-// With `runtimeExecutor`.
-JSScheduler::JSScheduler(
-    jsi::Runtime &rnRuntime,
-    RuntimeExecutor runtimeExecutor)
-    : scheduleOnJS([&](Job job) {
-        runtimeExecutor_(
-            [job = std::move(job)](jsi::Runtime &runtime) { job(runtime); });
-      }),
-      rnRuntime_(rnRuntime),
-      runtimeExecutor_(runtimeExecutor) {}
-#endif // RCT_NEW_ARCH_ENABLED
-
-const std::shared_ptr<CallInvoker> JSScheduler::getJSCallInvoker() const {
-  assert(
-      jsCallInvoker_ != nullptr &&
-      "[Reanimated] Expected jsCallInvoker, got nullptr instead.");
-  return jsCallInvoker_;
-}
-
 } // namespace worklets

packages/react-native-reanimated/Common/cpp/worklets/Tools/JSScheduler.h

@@ -15,26 +15,14 @@ class JSScheduler {
   using Job = std::function<void(jsi::Runtime &rt)>;
 
  public:
-  // With `jsCallInvoker`.
   explicit JSScheduler(
       jsi::Runtime &rnRuntime,
       const std::shared_ptr<CallInvoker> &jsCallInvoker);
 
-#ifdef RCT_NEW_ARCH_ENABLED
-  // With `runtimeExecutor`.
-  explicit JSScheduler(
-      jsi::Runtime &rnRuntime,
-      RuntimeExecutor runtimeExecutor);
-#endif // RCT_NEW_ARCH_ENABLED
-
   const std::function<void(Job)> scheduleOnJS = nullptr;
-  const std::shared_ptr<CallInvoker> getJSCallInvoker() const;
 
  protected:
   jsi::Runtime &rnRuntime_;
-#ifdef RCT_NEW_ARCH_ENABLED
-  RuntimeExecutor runtimeExecutor_ = nullptr;
-#endif // RCT_NEW_ARCH_ENABLED
   const std::shared_ptr<CallInvoker> jsCallInvoker_ = nullptr;
 };
 

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

@@ -45,7 +45,7 @@ NativeProxy::NativeProxy(
       reanimatedModuleProxy_(std::make_shared<ReanimatedModuleProxy>(
           workletsModuleProxy,
           *rnRuntime,
-          std::make_shared<JSScheduler>(*rnRuntime, jsCallInvoker),
+          jsCallInvoker,
           uiScheduler,
           getPlatformDependentMethods(),
           isBridgeless,

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

@@ -21,22 +21,37 @@ WorkletsModule::WorkletsModule(
     jni::alias_ref<WorkletsModule::javaobject> jThis,
     jsi::Runtime *rnRuntime,
     const std::string &valueUnpackerCode,
-    jni::alias_ref<JavaMessageQueueThread::javaobject> messageQueueThread)
+    jni::alias_ref<JavaMessageQueueThread::javaobject> messageQueueThread,
+    const std::shared_ptr<facebook::react::CallInvoker> &jsCallInvoker,
+    const std::shared_ptr<worklets::JSScheduler> &jsScheduler)
     : javaPart_(jni::make_global(jThis)),
       rnRuntime_(rnRuntime),
       workletsModuleProxy_(std::make_shared<WorkletsModuleProxy>(
           valueUnpackerCode,
-          std::make_shared<JMessageQueueThread>(messageQueueThread))) {
+          std::make_shared<JMessageQueueThread>(messageQueueThread),
+          jsCallInvoker,
+          jsScheduler)) {
   RNRuntimeWorkletDecorator::decorate(*rnRuntime_, workletsModuleProxy_);
 }
 
 jni::local_ref<WorkletsModule::jhybriddata> WorkletsModule::initHybrid(
     jni::alias_ref<jhybridobject> jThis,
     jlong jsContext,
     const std::string &valueUnpackerCode,
-    jni::alias_ref<JavaMessageQueueThread::javaobject> messageQueueThread) {
+    jni::alias_ref<JavaMessageQueueThread::javaobject> messageQueueThread,
+    jni::alias_ref<facebook::react::CallInvokerHolder::javaobject>
+        jsCallInvokerHolder) {
+  auto jsCallInvoker = jsCallInvokerHolder->cthis()->getCallInvoker();
+  auto rnRuntime = reinterpret_cast<jsi::Runtime *>(jsContext);
+  auto jsScheduler =
+      std::make_shared<worklets::JSScheduler>(*rnRuntime, jsCallInvoker);
   return makeCxxInstance(
-      jThis, (jsi::Runtime *)jsContext, valueUnpackerCode, messageQueueThread);
+      jThis,
+      rnRuntime,
+      valueUnpackerCode,
+      messageQueueThread,
+      jsCallInvoker,
+      jsScheduler);
 }
 
 void WorkletsModule::registerNatives() {

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

@@ -32,7 +32,9 @@ class WorkletsModule : public jni::HybridClass<WorkletsModule> {
       jni::alias_ref<jhybridobject> jThis,
       jlong jsContext,
       const std::string &valueUnpackerCode,
-      jni::alias_ref<JavaMessageQueueThread::javaobject> messageQueueThread);
+      jni::alias_ref<JavaMessageQueueThread::javaobject> messageQueueThread,
+      jni::alias_ref<facebook::react::CallInvokerHolder::javaobject>
+          jsCallInvokerHolder);
 
   static void registerNatives();
 
@@ -50,7 +52,9 @@ class WorkletsModule : public jni::HybridClass<WorkletsModule> {
       jni::alias_ref<WorkletsModule::jhybridobject> jThis,
       jsi::Runtime *rnRuntime,
       const std::string &valueUnpackerCode,
-      jni::alias_ref<JavaMessageQueueThread::javaobject> messageQueueThread);
+      jni::alias_ref<JavaMessageQueueThread::javaobject> messageQueueThread,
+      const std::shared_ptr<facebook::react::CallInvoker> &jsCallInvoker,
+      const std::shared_ptr<worklets::JSScheduler> &jsScheduler);
 };
 
 } // namespace worklets

packages/react-native-reanimated/android/src/main/java/com/swmansion/worklets/WorkletsModule.java

@@ -8,10 +8,12 @@
 import com.facebook.react.bridge.queue.MessageQueueThread;
 import com.facebook.react.common.annotations.FrameworkAPI;
 import com.facebook.react.module.annotations.ReactModule;
+import com.facebook.react.turbomodule.core.CallInvokerHolderImpl;
 import com.facebook.soloader.SoLoader;
 import com.swmansion.reanimated.NativeWorkletsModuleSpec;
 import java.util.Objects;
 
+/* * @noinspection JavaJniMissingFunction */
 @ReactModule(name = WorkletsModule.NAME)
 public class WorkletsModule extends NativeWorkletsModuleSpec {
   static {
@@ -22,21 +24,19 @@ public class WorkletsModule extends NativeWorkletsModuleSpec {
   @SuppressWarnings("unused")
   private HybridData mHybridData;
 
-  /**
-   * @noinspection unused
-   */
+  @SuppressWarnings("unused")
   protected HybridData getHybridData() {
     return mHybridData;
   }
 
   private final WorkletsMessageQueueThread mMessageQueueThread = new WorkletsMessageQueueThread();
 
-  /**
-   * @noinspection JavaJniMissingFunction
-   */
   @OptIn(markerClass = FrameworkAPI.class)
   private native HybridData initHybrid(
-      long jsContext, String valueUnpackerCode, MessageQueueThread messageQueueThread);
+      long jsContext,
+      String valueUnpackerCode,
+      MessageQueueThread messageQueueThread,
+      CallInvokerHolderImpl jsCallInvokerHolder);
 
   public WorkletsModule(ReactApplicationContext reactContext) {
     super(reactContext);
@@ -47,9 +47,10 @@ public WorkletsModule(ReactApplicationContext reactContext) {
   public boolean installTurboModule(String valueUnpackerCode) {
     var context = getReactApplicationContext();
     var jsContext = Objects.requireNonNull(context.getJavaScriptContextHolder()).get();
+    var jsCallInvokerHolder = JSCallInvokerResolver.getJSCallInvokerHolder(context);
 
-    mHybridData = initHybrid(jsContext, valueUnpackerCode, mMessageQueueThread);
-
+    mHybridData =
+        initHybrid(jsContext, valueUnpackerCode, mMessageQueueThread, jsCallInvokerHolder);
     return true;
   }
 }

packages/react-native-reanimated/apple/reanimated/apple/REAModule.mm

@@ -65,7 +65,6 @@ @implementation REAModule {
   SingleInstanceChecker<REAModule> singleInstanceChecker_;
 #endif // NDEBUG
   bool hasListeners;
-  bool _isBridgeless;
 }
 
 @synthesize moduleRegistry = _moduleRegistry;
@@ -174,7 +173,6 @@ - (void)attachReactEventListener
 - (void)setSurfacePresenter:(id<RCTSurfacePresenterStub>)surfacePresenter
 {
   _surfacePresenter = surfacePresenter;
-  _isBridgeless = true;
 }
 
 - (void)setBridge:(RCTBridge *)bridge
@@ -294,11 +292,12 @@ - (void)sendEventWithName:(NSString *)eventName body:(id)body
   auto jsCallInvoker = self.bridge.jsCallInvoker;
 #endif // defined(RCT_NEW_ARCH_ENABLED) && REACT_NATIVE_MINOR_VERSION >= 75
   auto jsiRuntime = reinterpret_cast<facebook::jsi::Runtime *>(self.bridge.runtime);
+  auto isBridgeless = ![self.bridge isKindOfClass:[RCTCxxBridge class]];
 
   assert(jsiRuntime != nullptr);
 
   auto reanimatedModuleProxy =
-      reanimated::createReanimatedModule(self, self.bridge, jsCallInvoker, workletsModule, _isBridgeless);
+      reanimated::createReanimatedModule(self, self.bridge, jsCallInvoker, workletsModule, isBridgeless);
 
   jsi::Runtime &rnRuntime = *jsiRuntime;
   [self commonInit:reanimatedModuleProxy withRnRuntime:rnRuntime];

packages/react-native-reanimated/apple/reanimated/apple/native/NativeProxy.mm

@@ -70,14 +70,13 @@ static inline bool getIsReducedMotion()
   PlatformDepMethodsHolder platformDepMethodsHolder = makePlatformDepMethodsHolder(bridge, nodesManager, reaModule);
 
   std::shared_ptr<UIScheduler> uiScheduler = std::make_shared<REAIOSUIScheduler>();
-  std::shared_ptr<JSScheduler> jsScheduler = std::make_shared<JSScheduler>(rnRuntime, jsInvoker);
 
   const auto workletsModuleProxy = [workletsModule getWorkletsModuleProxy];
 
   auto reanimatedModuleProxy = std::make_shared<ReanimatedModuleProxy>(
       workletsModuleProxy,
       rnRuntime,
-      jsScheduler,
+      jsInvoker,
       uiScheduler,
       platformDepMethodsHolder,
       isBridgeless,

packages/react-native-reanimated/apple/worklets/apple/WorkletsModule.mm

@@ -1,8 +1,8 @@
 #import <React/RCTBridge+Private.h>
+#import <worklets/Tools/SingleInstanceChecker.h>
 #import <worklets/WorkletRuntime/RNRuntimeWorkletDecorator.h>
 #import <worklets/apple/WorkletsMessageThread.h>
 #import <worklets/apple/WorkletsModule.h>
-#import <worklets/tools/SingleInstanceChecker.h>
 
 using worklets::RNRuntimeWorkletDecorator;
 using worklets::WorkletsModuleProxy;
@@ -39,7 +39,12 @@ @implementation WorkletsModule {
   auto jsQueue = std::make_shared<WorkletsMessageThread>([NSRunLoop currentRunLoop], ^(NSError *error) {
     throw error;
   });
-  workletsModuleProxy_ = std::make_shared<WorkletsModuleProxy>(std::string([valueUnpackerCode UTF8String]), jsQueue);
+
+  std::string valueUnpackerCodeStr = [valueUnpackerCode UTF8String];
+  auto jsCallInvoker = bridge.jsCallInvoker;
+  auto jsScheduler = std::make_shared<worklets::JSScheduler>(rnRuntime, jsCallInvoker);
+  workletsModuleProxy_ =
+      std::make_shared<WorkletsModuleProxy>(valueUnpackerCodeStr, jsQueue, jsCallInvoker, jsScheduler);
   RNRuntimeWorkletDecorator::decorate(rnRuntime, workletsModuleProxy_);
 
   return @YES;