Add HostFunction handling to _scheduleOnJS

packages/react-native-reanimated/Common/cpp/worklets/WorkletRuntime/WorkletRuntimeDecorator.cpp

@@ -19,6 +19,22 @@ static inline double performanceNow() {
   return duration / NANOSECONDS_IN_MILLISECOND;
 }
 
+static inline std::vector<jsi::Value> parseArgs(
+    jsi::Runtime &rt,
+    std::shared_ptr<ShareableArray> shareableArgs) {
+  if (shareableArgs == nullptr) {
+    return {};
+  }
+
+  auto argsArray = shareableArgs->toJSValue(rt).asObject(rt).asArray(rt);
+  auto argsSize = argsArray.size(rt);
+  std::vector<jsi::Value> result(argsSize);
+  for (size_t i = 0; i < argsSize; i++) {
+    result[i] = argsArray.getValueAtIndex(rt, i);
+  }
+  return result;
+}
+
 void WorkletRuntimeDecorator::decorate(
     jsi::Runtime &rt,
     const std::string &name,
@@ -78,39 +94,61 @@ void WorkletRuntimeDecorator::decorate(
 
   jsi_utils::installJsiFunction(
       rt,
-      "_scheduleOnJS",
+      "_scheduleRemoteFunctionOnJS",
       [jsScheduler](
           jsi::Runtime &rt,
-          const jsi::Value &remoteFun,
+          const jsi::Value &funValue,
           const jsi::Value &argsValue) {
         auto shareableRemoteFun = extractShareableOrThrow<
             ShareableRemoteFunction>(
             rt,
-            remoteFun,
+            funValue,
             "[Reanimated] Incompatible object passed to scheduleOnJS. It is only allowed to schedule worklets or functions defined on the React Native JS runtime this way.");
+
         auto shareableArgs = argsValue.isUndefined()
             ? nullptr
             : extractShareableOrThrow<ShareableArray>(
                   rt, argsValue, "[Reanimated] Args must be an array.");
+
         jsScheduler->scheduleOnJS([=](jsi::Runtime &rt) {
-          auto remoteFun = shareableRemoteFun->toJSValue(rt);
+          auto fun =
+              shareableRemoteFun->toJSValue(rt).asObject(rt).asFunction(rt);
           if (shareableArgs == nullptr) {
             // fast path for remote function w/o arguments
-            remoteFun.asObject(rt).asFunction(rt).call(rt);
+            fun.call(rt);
           } else {
-            auto argsArray =
-                shareableArgs->toJSValue(rt).asObject(rt).asArray(rt);
-            auto argsSize = argsArray.size(rt);
-            std::vector<jsi::Value> args(argsSize);
-            for (size_t i = 0; i < argsSize; i++) {
-              args[i] = argsArray.getValueAtIndex(rt, i);
-            }
-            remoteFun.asObject(rt).asFunction(rt).call(
+            auto args = parseArgs(rt, shareableArgs);
+            fun.call(
                 rt, const_cast<const jsi::Value *>(args.data()), args.size());
           }
         });
       });
 
+  jsi_utils::installJsiFunction(
+      rt,
+      "_scheduleHostFunctionOnJS",
+      [jsScheduler](
+          jsi::Runtime &rt,
+          const jsi::Value &hostFunValue,
+          const jsi::Value &argsValue) {
+        auto hostFun =
+            hostFunValue.asObject(rt).asFunction(rt).getHostFunction(rt);
+
+        auto shareableArgs = argsValue.isUndefined()
+            ? nullptr
+            : extractShareableOrThrow<ShareableArray>(
+                  rt, argsValue, "[Reanimated] Args must be an array.");
+
+        jsScheduler->scheduleOnJS([=](jsi::Runtime &rt) {
+          auto args = parseArgs(rt, shareableArgs);
+          hostFun(
+              rt,
+              jsi::Value::undefined(),
+              const_cast<const jsi::Value *>(args.data()),
+              args.size());
+        });
+      });
+
   jsi_utils::installJsiFunction(
       rt,
       "_scheduleOnRuntime",

packages/react-native-reanimated/plugin/src/globals.ts

@@ -133,7 +133,8 @@ const notCapturedIdentifiers_DEPRECATED = [
   '_IS_FABRIC',
   '_log',
   '_toString',
-  '_scheduleOnJS',
+  '_scheduleHostFunctionOnJS',
+  '_scheduleRemoteFunctionOnJS',
   '_scheduleOnRuntime',
   '_makeShareableClone',
   '_updatePropsPaper',

packages/react-native-reanimated/src/js-reanimated/index.ts

@@ -18,9 +18,15 @@ global._makeShareableClone = () => {
   );
 };
 
-global._scheduleOnJS = () => {
+global._scheduleHostFunctionOnJS = () => {
   throw new ReanimatedError(
-    '_scheduleOnJS should never be called in JSReanimated.'
+    '_scheduleHostFunctionOnJS should never be called in JSReanimated.'
+  );
+};
+
+global._scheduleRemoteFunctionOnJS = () => {
+  throw new ReanimatedError(
+    '_scheduleHostFunctionOnJS should never be called in JSReanimated.'
   );
 };
 

packages/react-native-reanimated/src/privateGlobals.d.ts

@@ -49,7 +49,8 @@ declare global {
     value: T,
     nativeStateSource?: object
   ) => FlatShareableRef<T>;
-  var _scheduleOnJS: (fun: (...args: A) => R, args?: A) => void;
+  var _scheduleHostFunctionOnJS: (fun: (...args: A) => R, args?: A) => void;
+  var _scheduleRemoteFunctionOnJS: (fun: (...args: A) => R, args?: A) => void;
   var _scheduleOnRuntime: (
     runtime: WorkletRuntime,
     worklet: ShareableRef<() => void>

packages/react-native-reanimated/src/threads.ts

@@ -254,8 +254,14 @@ export function runOnJS<Args extends unknown[], ReturnValue>(
     // reference to the original remote function in the `__remoteFunction` property.
     fun = (fun as FunDevRemote).__remoteFunction;
   }
+
+  const scheduleOnJS =
+    typeof fun === 'function'
+      ? global._scheduleHostFunctionOnJS
+      : global._scheduleRemoteFunctionOnJS;
+
   return (...args) => {
-    global._scheduleOnJS(
+    scheduleOnJS(
       fun as
         | ((...args: Args) => ReturnValue)
         | WorkletFunction<Args, ReturnValue>,