refactor(Worklets): move UI scheduling (software-mansion#6793)

## Summary All methods related to scheduling are moved to Worklets. ## Test plan Apps are zooming.
mhoran · Dec 17, 2024 · c6d6815 · c6d6815
1 parent 4c885c2
commit c6d6815
Show file tree

Hide file tree

Showing 17 changed files with 238 additions and 215 deletions.
diff --git a/...ges/react-native-reanimated/Common/cpp/reanimated/NativeModules/ReanimatedModuleProxy.cpp b/...ges/react-native-reanimated/Common/cpp/reanimated/NativeModules/ReanimatedModuleProxy.cpp
@@ -193,56 +193,6 @@ ReanimatedModuleProxy::~ReanimatedModuleProxy() {
 #endif // RCT_NEW_ARCH_ENABLED
 }
 
-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(
-        workletsModuleProxy_->getUIWorkletRuntime()->getJSIRuntime());
-#endif
-    workletsModuleProxy_->getUIWorkletRuntime()->runGuarded(shareableWorklet);
-  });
-}
-
-jsi::Value ReanimatedModuleProxy::executeOnUIRuntimeSync(
-    jsi::Runtime &rt,
-    const jsi::Value &worklet) {
-  return workletsModuleProxy_->getUIWorkletRuntime()->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(
     jsi::Runtime &rt,
     const jsi::Value &worklet,

diff --git a/packages/react-native-reanimated/Common/cpp/reanimated/NativeModules/ReanimatedModuleProxy.h b/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,

diff --git a/...react-native-reanimated/Common/cpp/reanimated/NativeModules/ReanimatedModuleProxySpec.cpp b/...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"] =

diff --git a/...s/react-native-reanimated/Common/cpp/reanimated/NativeModules/ReanimatedModuleProxySpec.h b/...s/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,

diff --git a/packages/react-native-reanimated/Common/cpp/worklets/NativeModules/WorkletsModuleProxy.cpp b/packages/react-native-reanimated/Common/cpp/worklets/NativeModules/WorkletsModuleProxy.cpp
@@ -14,6 +14,18 @@
 
 #include <jsi/jsi.h>
 
+// Standard `__cplusplus` macro reference:
+// https://en.cppreference.com/w/cpp/preprocessor/replace#Predefined_macros
+#if REACT_NATIVE_MINOR_VERSION >= 75 || __cplusplus >= 202002L
+// Implicit copy capture of `this` is deprecated in NDK27, which uses C++20.
+#define COPY_CAPTURE_WITH_THIS [ =, this ] // NOLINT (whitespace/braces)
+#else
+// React Native 0.75 is the last one which allows NDK23. NDK23 uses C++17 and
+// explicitly disallows C++20 features, including the syntax above. Therefore we
+// fallback to the deprecated syntax here.
+#define COPY_CAPTURE_WITH_THIS [=] // NOLINT (whitespace/braces)
+#endif // REACT_NATIVE_MINOR_VERSION >= 75 || __cplusplus >= 202002L
+
 using namespace facebook;
 
 namespace worklets {
@@ -53,4 +65,53 @@ jsi::Value WorkletsModuleProxy::makeShareableClone(
       rt, value, shouldRetainRemote, nativeStateSource);
 }
 
+void WorkletsModuleProxy::scheduleOnUI(
+    jsi::Runtime &rt,
+    const jsi::Value &worklet) {
+  auto shareableWorklet = extractShareableOrThrow<ShareableWorklet>(
+      rt, worklet, "[Worklets] Only worklets can be scheduled to run on UI.");
+  uiScheduler_->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 WorkletsModuleProxy::executeOnUIRuntimeSync(
+    jsi::Runtime &rt,
+    const jsi::Value &worklet) {
+  return uiWorkletRuntime_->executeSync(rt, worklet);
+}
+
+jsi::Value WorkletsModuleProxy::createWorkletRuntime(
+    jsi::Runtime &rt,
+    const jsi::Value &name,
+    const jsi::Value &initializer) {
+  auto workletRuntime = std::make_shared<WorkletRuntime>(
+      rt,
+      jsQueue_,
+      jsScheduler_,
+      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 WorkletsModuleProxy::scheduleOnRuntime(
+    jsi::Runtime &rt,
+    const jsi::Value &workletRuntimeValue,
+    const jsi::Value &shareableWorkletValue) {
+  worklets::scheduleOnRuntime(rt, workletRuntimeValue, shareableWorkletValue);
+  return jsi::Value::undefined();
+}
+
 } // namespace worklets
diff --git a/packages/react-native-reanimated/Common/cpp/worklets/NativeModules/WorkletsModuleProxy.h b/packages/react-native-reanimated/Common/cpp/worklets/NativeModules/WorkletsModuleProxy.h
@@ -29,6 +29,21 @@ class WorkletsModuleProxy : public WorkletsModuleProxySpec {
       const jsi::Value &shouldRetainRemote,
       const jsi::Value &nativeStateSource) override;
 
+  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;
+
   [[nodiscard]] inline std::string getValueUnpackerCode() const {
     return valueUnpackerCode_;
   }

diff --git a/...ges/react-native-reanimated/Common/cpp/worklets/NativeModules/WorkletsModuleProxySpec.cpp b/...ges/react-native-reanimated/Common/cpp/worklets/NativeModules/WorkletsModuleProxySpec.cpp
@@ -17,11 +17,56 @@ static jsi::Value WORKLETS_SPEC_PREFIX(makeShareableClone)(
           rt, std::move(args[0]), std::move(args[1]), std::move(args[2]));
 }
 
+static jsi::Value WORKLETS_SPEC_PREFIX(scheduleOnUI)(
+    jsi::Runtime &rt,
+    TurboModule &turboModule,
+    const jsi::Value *args,
+    size_t) {
+  static_cast<WorkletsModuleProxySpec *>(&turboModule)
+      ->scheduleOnUI(rt, std::move(args[0]));
+  return jsi::Value::undefined();
+}
+
+static jsi::Value WORKLETS_SPEC_PREFIX(executeOnUIRuntimeSync)(
+    jsi::Runtime &rt,
+    TurboModule &turboModule,
+    const jsi::Value *args,
+    size_t) {
+  return static_cast<WorkletsModuleProxySpec *>(&turboModule)
+      ->executeOnUIRuntimeSync(rt, std::move(args[0]));
+}
+
+static jsi::Value WORKLETS_SPEC_PREFIX(createWorkletRuntime)(
+    jsi::Runtime &rt,
+    TurboModule &turboModule,
+    const jsi::Value *args,
+    size_t) {
+  return static_cast<WorkletsModuleProxySpec *>(&turboModule)
+      ->createWorkletRuntime(rt, std::move(args[0]), std::move(args[1]));
+}
+
+static jsi::Value WORKLETS_SPEC_PREFIX(scheduleOnRuntime)(
+    jsi::Runtime &rt,
+    TurboModule &turboModule,
+    const jsi::Value *args,
+    size_t) {
+  return static_cast<WorkletsModuleProxySpec *>(&turboModule)
+      ->scheduleOnRuntime(rt, std::move(args[0]), std::move(args[1]));
+}
+
 WorkletsModuleProxySpec::WorkletsModuleProxySpec(
     const std::shared_ptr<CallInvoker> jsInvoker)
     : TurboModule("NativeWorklets", jsInvoker) {
   methodMap_["makeShareableClone"] =
       MethodMetadata{2, WORKLETS_SPEC_PREFIX(makeShareableClone)};
+  methodMap_["scheduleOnUI"] =
+      MethodMetadata{1, WORKLETS_SPEC_PREFIX(scheduleOnUI)};
+  methodMap_["executeOnUIRuntimeSync"] =
+      MethodMetadata{1, WORKLETS_SPEC_PREFIX(executeOnUIRuntimeSync)};
+  methodMap_["createWorkletRuntime"] =
+      MethodMetadata{2, WORKLETS_SPEC_PREFIX(createWorkletRuntime)};
+  methodMap_["scheduleOnRuntime"] =
+      MethodMetadata{2, WORKLETS_SPEC_PREFIX(scheduleOnRuntime)};
 }
 
 } // namespace worklets
diff --git a/packages/react-native-reanimated/Common/cpp/worklets/NativeModules/WorkletsModuleProxySpec.h b/packages/react-native-reanimated/Common/cpp/worklets/NativeModules/WorkletsModuleProxySpec.h
@@ -15,11 +15,30 @@ class JSI_EXPORT WorkletsModuleProxySpec : public TurboModule {
       const std::shared_ptr<CallInvoker> jsInvoker);
 
  public:
+  // Shareables
   virtual jsi::Value makeShareableClone(
       jsi::Runtime &rt,
       const jsi::Value &value,
       const jsi::Value &shouldRetainRemote,
       const jsi::Value &nativeStateSource) = 0;
+
+  // 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;
 };
 
 } // namespace worklets
diff --git a/packages/react-native-reanimated/src/ReanimatedModule/NativeReanimated.ts b/packages/react-native-reanimated/src/ReanimatedModule/NativeReanimated.ts
@@ -10,7 +10,6 @@ import type {
 } from '../commonTypes';
 import { checkCppVersion } from '../platform-specific/checkCppVersion';
 import { jsVersion } from '../platform-specific/jsVersion';
-import type { WorkletRuntime } from '../runtimes';
 import { isFabric } from '../PlatformChecker';
 import type React from 'react';
 import { getShadowNodeWrapperFromRef } from '../fabricUtils';
@@ -65,28 +64,6 @@ See https://docs.swmansion.com/react-native-reanimated/docs/guides/troubleshooti
     this.#reanimatedModuleProxy = global.__reanimatedModuleProxy;
   }
 
-  scheduleOnUI<T>(shareable: ShareableRef<T>) {
-    return this.#reanimatedModuleProxy.scheduleOnUI(shareable);
-  }
-
-  executeOnUIRuntimeSync<T, R>(shareable: ShareableRef<T>): R {
-    return this.#reanimatedModuleProxy.executeOnUIRuntimeSync(shareable);
-  }
-
-  createWorkletRuntime(name: string, initializer: ShareableRef<() => void>) {
-    return this.#reanimatedModuleProxy.createWorkletRuntime(name, initializer);
-  }
-
-  scheduleOnRuntime<T>(
-    workletRuntime: WorkletRuntime,
-    shareableWorklet: ShareableRef<T>
-  ) {
-    return this.#reanimatedModuleProxy.scheduleOnRuntime(
-      workletRuntime,
-      shareableWorklet
-    );
-  }
-
   registerSensor(
     sensorType: number,
     interval: number,