refactor(Worklets): move UIScheduler

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

@@ -54,15 +54,13 @@ ReanimatedModuleProxy::ReanimatedModuleProxy(
     const std::shared_ptr<WorkletsModuleProxy> &workletsModuleProxy,
     jsi::Runtime &rnRuntime,
     const std::shared_ptr<CallInvoker> &jsCallInvoker,
-    const std::shared_ptr<UIScheduler> &uiScheduler,
     const PlatformDepMethodsHolder &platformDepMethodsHolder,
     const bool isBridgeless,
     const bool isReducedMotion)
     : ReanimatedModuleProxySpec(jsCallInvoker),
       isBridgeless_(isBridgeless),
       isReducedMotion_(isReducedMotion),
       workletsModuleProxy_(workletsModuleProxy),
-      uiScheduler_(uiScheduler),
       valueUnpackerCode_(workletsModuleProxy->getValueUnpackerCode()),
       uiWorkletRuntime_(std::make_shared<WorkletRuntime>(
           rnRuntime,
@@ -207,7 +205,7 @@ void ReanimatedModuleProxy::scheduleOnUI(
     const jsi::Value &worklet) {
   auto shareableWorklet = extractShareableOrThrow<ShareableWorklet>(
       rt, worklet, "[Reanimated] Only worklets can be scheduled to run on UI.");
-  uiScheduler_->scheduleOnUI(COPY_CAPTURE_WITH_THIS {
+  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
@@ -264,7 +262,7 @@ jsi::Value ReanimatedModuleProxy::registerEventHandler(
       rt, worklet, "[Reanimated] Event handler must be a worklet.");
   int emitterReactTagInt = emitterReactTag.asNumber();
 
-  uiScheduler_->scheduleOnUI(COPY_CAPTURE_WITH_THIS {
+  workletsModuleProxy_->getUIScheduler()->scheduleOnUI(COPY_CAPTURE_WITH_THIS {
     auto handler = std::make_shared<WorkletEventHandler>(
         newRegistrationId, eventNameStr, emitterReactTagInt, handlerShareable);
     eventHandlerRegistry_->registerEventHandler(std::move(handler));
@@ -277,7 +275,7 @@ void ReanimatedModuleProxy::unregisterEventHandler(
     jsi::Runtime &,
     const jsi::Value &registrationId) {
   uint64_t id = registrationId.asNumber();
-  uiScheduler_->scheduleOnUI(
+  workletsModuleProxy_->getUIScheduler()->scheduleOnUI(
       COPY_CAPTURE_WITH_THIS
 
       { eventHandlerRegistry_->unregisterEventHandler(id); });
@@ -349,7 +347,7 @@ 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(COPY_CAPTURE_WITH_THIS {
+  workletsModuleProxy_->getUIScheduler()->scheduleOnUI(COPY_CAPTURE_WITH_THIS {
     jsi::Runtime &uiRuntime = uiWorkletRuntime_->getJSIRuntime();
     const auto resultStr =
         obtainPropFromShadowNode(uiRuntime, propNameStr, shadowNode);
@@ -377,7 +375,7 @@ jsi::Value ReanimatedModuleProxy::getViewProp(
 
   const int viewTagInt = viewTag.asNumber();
 
-  uiScheduler_->scheduleOnUI(
+  workletsModuleProxy_->getUIScheduler()->scheduleOnUI(
       COPY_CAPTURE_WITH_THIS
 
       () {
@@ -859,7 +857,7 @@ void ReanimatedModuleProxy::initializeLayoutAnimationsProxy() {
         componentDescriptorRegistry,
         scheduler->getContextContainer(),
         uiWorkletRuntime_->getJSIRuntime(),
-        uiScheduler_);
+        workletsModuleProxy_->getUIScheduler());
   }
 }
 

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

@@ -36,7 +36,6 @@ class ReanimatedModuleProxy : public ReanimatedModuleProxySpec {
       const std::shared_ptr<WorkletsModuleProxy> &workletsModuleProxy,
       jsi::Runtime &rnRuntime,
       const std::shared_ptr<CallInvoker> &jsCallInvoker,
-      const std::shared_ptr<UIScheduler> &uiScheduler,
       const PlatformDepMethodsHolder &platformDepMethodsHolder,
       const bool isBridgeless,
       const bool isReducedMotion);
@@ -194,7 +193,6 @@ class ReanimatedModuleProxy : public ReanimatedModuleProxySpec {
   const bool isBridgeless_;
   const bool isReducedMotion_;
   const std::shared_ptr<WorkletsModuleProxy> workletsModuleProxy_;
-  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

@@ -22,11 +22,13 @@ WorkletsModuleProxy::WorkletsModuleProxy(
     const std::string &valueUnpackerCode,
     const std::shared_ptr<MessageQueueThread> &jsQueue,
     const std::shared_ptr<CallInvoker> &jsCallInvoker,
-    const std::shared_ptr<JSScheduler> &jsScheduler)
+    const std::shared_ptr<JSScheduler> &jsScheduler,
+    const std::shared_ptr<UIScheduler> &uiScheduler)
     : WorkletsModuleProxySpec(jsCallInvoker),
       valueUnpackerCode_(valueUnpackerCode),
       jsQueue_(jsQueue),
-      jsScheduler_(jsScheduler) {}
+      jsScheduler_(jsScheduler),
+      uiScheduler_(uiScheduler) {}
 
 WorkletsModuleProxy::~WorkletsModuleProxy() {}
 

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

@@ -4,6 +4,7 @@
 #include <worklets/NativeModules/WorkletsModuleProxySpec.h>
 #include <worklets/Tools/JSScheduler.h>
 #include <worklets/Tools/SingleInstanceChecker.h>
+#include <worklets/Tools/UIScheduler.h>
 #include <worklets/WorkletRuntime/WorkletRuntime.h>
 #include <memory>
 #include <string>
@@ -16,7 +17,8 @@ class WorkletsModuleProxy : public WorkletsModuleProxySpec {
       const std::string &valueUnpackerCode,
       const std::shared_ptr<MessageQueueThread> &jsQueue,
       const std::shared_ptr<CallInvoker> &jsCallInvoker,
-      const std::shared_ptr<JSScheduler> &jsScheduler);
+      const std::shared_ptr<JSScheduler> &jsScheduler,
+      const std::shared_ptr<UIScheduler> &uiScheduler);
 
   ~WorkletsModuleProxy();
 
@@ -38,10 +40,15 @@ class WorkletsModuleProxy : public WorkletsModuleProxySpec {
     return jsScheduler_;
   }
 
+  [[nodiscard]] inline std::shared_ptr<UIScheduler> getUIScheduler() const {
+    return uiScheduler_;
+  }
+
  private:
   const std::string valueUnpackerCode_;
   const std::shared_ptr<MessageQueueThread> jsQueue_;
   const std::shared_ptr<JSScheduler> jsScheduler_;
+  const std::shared_ptr<UIScheduler> uiScheduler_;
 #ifndef NDEBUG
   SingleInstanceChecker<WorkletsModuleProxy> singleInstanceChecker_;
 #endif // NDEBUG

packages/react-native-reanimated/android/src/fabric/java/com/swmansion/reanimated/NativeProxy.java

@@ -41,7 +41,6 @@ public class NativeProxy extends NativeProxyCommon {
             workletsModule,
             Objects.requireNonNull(context.getJavaScriptContextHolder()).get(),
             callInvokerHolder,
-            mAndroidUIScheduler,
             LayoutAnimations,
             context.isBridgeless(),
             fabricUIManager);
@@ -58,7 +57,6 @@ private native HybridData initHybrid(
       WorkletsModule workletsModule,
       long jsContext,
       CallInvokerHolderImpl jsCallInvokerHolder,
-      AndroidUIScheduler androidUIScheduler,
       LayoutAnimations LayoutAnimations,
       boolean isBridgeless,
       FabricUIManager fabricUIManager);

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

@@ -1,14 +1,14 @@
 #include <reanimated/LayoutAnimations/LayoutAnimationsManager.h>
 #include <reanimated/RuntimeDecorators/RNRuntimeDecorator.h>
 #include <reanimated/Tools/PlatformDepMethodsHolder.h>
-#include <reanimated/android/AndroidUIScheduler.h>
 #include <reanimated/android/NativeProxy.h>
 
 #include <worklets/Tools/ReanimatedJSIUtils.h>
 #include <worklets/Tools/ReanimatedVersion.h>
 #include <worklets/WorkletRuntime/ReanimatedRuntime.h>
 #include <worklets/WorkletRuntime/WorkletRuntime.h>
 #include <worklets/WorkletRuntime/WorkletRuntimeCollector.h>
+#include <worklets/android/AndroidUIScheduler.h>
 
 #include <android/log.h>
 #include <fbjni/fbjni.h>
@@ -31,7 +31,6 @@ NativeProxy::NativeProxy(
     const std::shared_ptr<WorkletsModuleProxy> &workletsModuleProxy,
     jsi::Runtime *rnRuntime,
     const std::shared_ptr<facebook::react::CallInvoker> &jsCallInvoker,
-    const std::shared_ptr<UIScheduler> &uiScheduler,
     jni::global_ref<LayoutAnimations::javaobject> layoutAnimations,
     const bool isBridgeless
 #ifdef RCT_NEW_ARCH_ENABLED
@@ -46,7 +45,6 @@ NativeProxy::NativeProxy(
           workletsModuleProxy,
           *rnRuntime,
           jsCallInvoker,
-          uiScheduler,
           getPlatformDependentMethods(),
           isBridgeless,
           getIsReducedMotion())),
@@ -91,7 +89,6 @@ jni::local_ref<NativeProxy::jhybriddata> NativeProxy::initHybrid(
     jlong jsContext,
     jni::alias_ref<facebook::react::CallInvokerHolder::javaobject>
         jsCallInvokerHolder,
-    jni::alias_ref<AndroidUIScheduler::javaobject> androidUiScheduler,
     jni::alias_ref<LayoutAnimations::javaobject> layoutAnimations,
     bool isBridgeless
 #ifdef RCT_NEW_ARCH_ENABLED
@@ -101,14 +98,12 @@ jni::local_ref<NativeProxy::jhybriddata> NativeProxy::initHybrid(
 #endif
 ) {
   auto jsCallInvoker = jsCallInvokerHolder->cthis()->getCallInvoker();
-  auto uiScheduler = androidUiScheduler->cthis()->getUIScheduler();
   auto workletsModuleProxy = jWorkletsModule->cthis()->getWorkletsModuleProxy();
   return makeCxxInstance(
       jThis,
       workletsModuleProxy,
       (jsi::Runtime *)jsContext,
       jsCallInvoker,
-      uiScheduler,
       make_global(layoutAnimations),
       isBridgeless
 #ifdef RCT_NEW_ARCH_ENABLED

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

@@ -1,11 +1,9 @@
 #pragma once
 
 #include <reanimated/NativeModules/ReanimatedModuleProxy.h>
-#include <reanimated/android/AndroidUIScheduler.h>
 #include <reanimated/android/JNIHelper.h>
 #include <reanimated/android/LayoutAnimations.h>
 
-#include <worklets/Tools/UIScheduler.h>
 #include <worklets/android/WorkletsModule.h>
 
 #include <ReactCommon/CallInvokerHolder.h>
@@ -153,7 +151,6 @@ class NativeProxy : public jni::HybridClass<NativeProxy> {
       jlong jsContext,
       jni::alias_ref<facebook::react::CallInvokerHolder::javaobject>
           jsCallInvokerHolder,
-      jni::alias_ref<AndroidUIScheduler::javaobject> androidUiScheduler,
       jni::alias_ref<LayoutAnimations::javaobject> layoutAnimations,
       const bool isBridgeless
 #ifdef RCT_NEW_ARCH_ENABLED
@@ -268,7 +265,6 @@ class NativeProxy : public jni::HybridClass<NativeProxy> {
       const std::shared_ptr<WorkletsModuleProxy> &workletsModuleProxy,
       jsi::Runtime *rnRuntime,
       const std::shared_ptr<facebook::react::CallInvoker> &jsCallInvoker,
-      const std::shared_ptr<UIScheduler> &uiScheduler,
       jni::global_ref<LayoutAnimations::javaobject> layoutAnimations,
       const bool isBridgeless
 #ifdef RCT_NEW_ARCH_ENABLED

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

@@ -1,6 +1,5 @@
 #include <fbjni/fbjni.h>
 
-#include <reanimated/android/AndroidUIScheduler.h>
 #include <reanimated/android/LayoutAnimations.h>
 #include <reanimated/android/NativeProxy.h>
 
@@ -9,7 +8,6 @@ JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM *vm, void *) {
     reanimated::NativeProxy::registerNatives();
     reanimated::AnimationFrameCallback::registerNatives();
     reanimated::EventHandler::registerNatives();
-    reanimated::AndroidUIScheduler::registerNatives();
     reanimated::LayoutAnimations::registerNatives();
     reanimated::SensorSetter::registerNatives();
     reanimated::KeyboardWorkletWrapper::registerNatives();

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

@@ -1,10 +1,10 @@
-#include <reanimated/android/AndroidUIScheduler.h>
+#include <worklets/android/AndroidUIScheduler.h>
 
 #include <android/log.h>
 #include <fbjni/fbjni.h>
 #include <jsi/jsi.h>
 
-namespace reanimated {
+namespace worklets {
 
 using namespace facebook;
 using namespace react;
@@ -57,4 +57,4 @@ void AndroidUIScheduler::registerNatives() {
   });
 }
 
-} // namespace reanimated
+} // namespace worklets

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

@@ -10,15 +10,15 @@
 
 #include <memory>
 
-namespace reanimated {
+namespace worklets {
 
 using namespace facebook;
 using namespace worklets;
 
 class AndroidUIScheduler : public jni::HybridClass<AndroidUIScheduler> {
  public:
   static auto constexpr kJavaDescriptor =
-      "Lcom/swmansion/reanimated/AndroidUIScheduler;";
+      "Lcom/swmansion/worklets/AndroidUIScheduler;";
   static jni::local_ref<jhybriddata> initHybrid(
       jni::alias_ref<jhybridobject> jThis);
   static void registerNatives();
@@ -41,4 +41,4 @@ class AndroidUIScheduler : public jni::HybridClass<AndroidUIScheduler> {
       jni::alias_ref<AndroidUIScheduler::jhybridobject> jThis);
 };
 
-} // namespace reanimated
+} // namespace worklets

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

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

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

@@ -14,6 +14,7 @@
 #include <react/jni/WritableNativeMap.h>
 
 #include <worklets/NativeModules/WorkletsModuleProxy.h>
+#include <worklets/android/AndroidUIScheduler.h>
 
 #include <memory>
 #include <string>
@@ -34,7 +35,9 @@ class WorkletsModule : public jni::HybridClass<WorkletsModule> {
       const std::string &valueUnpackerCode,
       jni::alias_ref<JavaMessageQueueThread::javaobject> messageQueueThread,
       jni::alias_ref<facebook::react::CallInvokerHolder::javaobject>
-          jsCallInvokerHolder);
+          jsCallInvokerHolder,
+      jni::alias_ref<worklets::AndroidUIScheduler::javaobject>
+          androidUIScheduler);
 
   static void registerNatives();
 
@@ -54,7 +57,8 @@ class WorkletsModule : public jni::HybridClass<WorkletsModule> {
       const std::string &valueUnpackerCode,
       jni::alias_ref<JavaMessageQueueThread::javaobject> messageQueueThread,
       const std::shared_ptr<facebook::react::CallInvoker> &jsCallInvoker,
-      const std::shared_ptr<worklets::JSScheduler> &jsScheduler);
+      const std::shared_ptr<worklets::JSScheduler> &jsScheduler,
+      const std::shared_ptr<UIScheduler> &uiScheduler);
 };
 
 } // namespace worklets

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

@@ -1,8 +1,11 @@
 #include <fbjni/fbjni.h>
 
+#include <worklets/android/AndroidUIScheduler.h>
 #include <worklets/android/WorkletsModule.h>
 
 JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM *vm, void *) {
-  return facebook::jni::initialize(
-      vm, [] { worklets::WorkletsModule::registerNatives(); });
+  return facebook::jni::initialize(vm, [] {
+    worklets::WorkletsModule::registerNatives();
+    worklets::AndroidUIScheduler::registerNatives();
+  });
 }

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

@@ -128,14 +128,13 @@ public void invalidate() {
     }
 
     if (mNativeProxy != null) {
-      mNativeProxy.invalidate();
       mNativeProxy = null;
     }
   }
 
   public void initWithContext(ReactApplicationContext reactApplicationContext) {
     mNativeProxy = new NativeProxy(reactApplicationContext, mWorkletsModule);
-    mAnimationManager.setAndroidUIScheduler(getNativeProxy().getAndroidUIScheduler());
+    mAnimationManager.setAndroidUIScheduler(mWorkletsModule.getAndroidUIScheduler());
     compatibility = new ReaCompatibility(reactApplicationContext);
     compatibility.registerFabricEventListener(this);
   }