feat: Add worklet() func that throws if value is not a worklet (#187)

* Add `worklet()` func that throws if value is not a worklet

* fix: Fix `isWorklet` check

* fix: Use `worklet` in `useWorklet`, and remove dependencylist

* fix: Better error

example/Tests/worklet-tests.ts

@@ -1,4 +1,4 @@
-import { Worklets } from "react-native-worklets-core";
+import { worklet, Worklets } from "react-native-worklets-core";
 import { ExpectException, ExpectValue, getWorkletInfo } from "./utils";
 
 export const worklet_tests = {
@@ -197,4 +197,20 @@ export const worklet_tests = {
     });
     return ExpectValue(result, 42);
   },
+  check_worklet_checker_works: () => {
+    const func = () => {
+      "worklet";
+      return 42;
+    };
+    const same = worklet(func);
+    return ExpectValue(same, func);
+  },
+  check_worklet_checker_throws_invalid_worklet: () => {
+    const func = () => {
+      return "not a worklet";
+    };
+    return ExpectException(() => {
+      worklet(func);
+    });
+  },
 };

example/src/App.tsx

@@ -1,18 +1,30 @@
-import React from "react";
+import React, { useEffect } from "react";
 import {
   ScrollView,
   StyleSheet,
   Text,
   TouchableOpacity,
   View,
 } from "react-native";
+import { useWorklet } from "react-native-worklets-core";
 
 import { useTestRunner } from "../Tests";
 import { TestWrapper } from "../Tests/TestWrapper";
 
 const App = () => {
   const { tests, categories, output, runTests, runSingleTest } =
     useTestRunner();
+
+  const dummyWorklet = useWorklet("default", (name: string): number => {
+    "worklet";
+    console.log(`useWorklet(${name}) called!`);
+    return name.length;
+  });
+
+  useEffect(() => {
+    dummyWorklet("marc");
+  }, [dummyWorklet]);
+
   return (
     <View style={styles.container}>
       <View style={styles.tests}>

src/hooks/useWorklet.ts

@@ -1,39 +1,33 @@
-import { DependencyList, useMemo } from "react";
-import type { IWorkletContext } from "src/types";
+import { useMemo } from "react";
+import type { IWorkletContext } from "../types";
+import { worklet } from "../worklet";
 
 /**
- * Create a Worklet function that persists between re-renders.
+ * Create a Worklet function that automatically memoizes itself using it's auto-captured closure.
  * The returned function can be called from both a Worklet context and the JS context, but will execute on a Worklet context.
  *
  * @worklet
  * @param context The context to run this Worklet in. Can be `default` to use the default background context, or a custom context.
  * @param callback The Worklet. Must be marked with the `'worklet'` directive.
- * @param dependencyList The React dependencies of this Worklet.
  * @returns A memoized Worklet
  * ```ts
  * const sayHello = useWorklet('default', (name: string) => {
  *   'worklet'
  *   console.log(`Hello ${name}, I am running on the Worklet Thread!`)
- * }, [])
+ * })
  * sayHello()
  * ```
  */
 export function useWorklet<T extends (...args: any[]) => any>(
   context: IWorkletContext | "default",
-  callback: T,
-  dependencyList: DependencyList
+  callback: T
 ): (...args: Parameters<T>) => Promise<ReturnType<T>> {
-  const worklet = useMemo(
-    () => {
-      if (context === "default") {
-        return Worklets.defaultContext.createRunAsync(callback);
-      } else {
-        return context.createRunAsync(callback);
-      }
-    },
+  const func = worklet(callback);
+  const ctx = context === "default" ? Worklets.defaultContext : context;
+
+  return useMemo(
+    () => ctx.createRunAsync(func),
     // eslint-disable-next-line react-hooks/exhaustive-deps
-    dependencyList
+    [...Object.values(func.__closure), ctx]
   );
-
-  return worklet;
 }

src/index.ts

@@ -1,5 +1,6 @@
 import "./NativeWorklets";
 export * from "./types";
+export * from "./worklet";
 export * from "./hooks/useRunOnJS";
 export * from "./hooks/useSharedValue";
 export * from "./hooks/useWorklet";

src/worklet.ts

@@ -0,0 +1,78 @@
+type AnyFunc = (...args: any[]) => any;
+
+type Workletize<TFunc extends () => any> = TFunc & {
+  __closure: Record<string, unknown>;
+  __initData: {
+    code: string;
+    location: string;
+    __sourceMap: string;
+  };
+  __workletHash: number;
+};
+const EXPECTED_KEYS: (keyof Workletize<AnyFunc>)[] = [
+  "__closure",
+  "__initData",
+  "__workletHash",
+];
+
+/**
+ * Checks whether the given function is a Worklet or not.
+ */
+export function isWorklet<TFunc extends AnyFunc>(
+  func: TFunc
+): func is Workletize<TFunc> {
+  const maybeWorklet = func as Partial<Workletize<TFunc>> & TFunc;
+  if (typeof maybeWorklet.__workletHash !== "number") return false;
+
+  if (
+    maybeWorklet.__closure == null ||
+    typeof maybeWorklet.__closure !== "object"
+  )
+    return false;
+
+  const initData = maybeWorklet.__initData;
+  if (initData == null || typeof initData !== "object") return false;
+
+  if (
+    typeof initData.__sourceMap !== "string" ||
+    typeof initData.code !== "string" ||
+    typeof initData.location !== "string"
+  )
+    return false;
+
+  return true;
+}
+
+class NotAWorkletError<TFunc extends AnyFunc> extends Error {
+  constructor(func: TFunc) {
+    let funcName = func.name;
+    if (funcName.length === 0) {
+      funcName = func.toString();
+    }
+
+    const expected = `[${EXPECTED_KEYS.join(", ")}]`;
+    const received = `[${Object.keys(func).join(", ")}]`;
+    super(
+      `The function "${funcName}" is not a Worklet! \n` +
+        `- Make sure the function "${funcName}" is decorated with the 'worklet' directive! \n` +
+        `- Make sure react-native-worklets-core is installed properly! \n` +
+        `- Make sure to add the react-native-worklets-core babel plugin to your babel.config.js! \n` +
+        `- Make sure that no other plugin overrides the react-native-worklets-core babel plugin! \n` +
+        `Expected "${funcName}" to contain ${expected}, but "${funcName}" only has these properties: ${received}`
+    );
+  }
+}
+
+/**
+ * Ensures the given function is a Worklet, and throws an error if not.
+ * @param func The function that should be a Worklet.
+ * @returns The same function that was passed in.
+ */
+export function worklet<TFunc extends () => any>(
+  func: TFunc
+): Workletize<TFunc> {
+  if (!isWorklet(func)) {
+    throw new NotAWorkletError(func);
+  }
+  return func;
+}