Transform '_WORKLET' in worklets to 'true'.

plugin/jestUtils.ts

@@ -1,10 +1,11 @@
-export {};
+import { UI_RUNTIME_CHECK_REGEX } from './src/makeWorklet';
 
 declare global {
   namespace jest {
     interface Matchers<R> {
       toHaveWorkletData(times?: number): R;
       toHaveInlineStyleWarning(times?: number): R;
+      toHaveUIRuntimeCheck(times?: number): R;
     }
   }
 }
@@ -52,3 +53,22 @@ expect.extend({
     };
   },
 });
+
+expect.extend({
+  toHaveUIRuntimeCheck(received: string, expectedMatchCount: number = 1) {
+    const receivedMatchCount = received.match(UI_RUNTIME_CHECK_REGEX)?.length;
+
+    if (receivedMatchCount === expectedMatchCount) {
+      return {
+        message: () =>
+          `Reanimated: expected to have UI runtime check ${expectedMatchCount} times`,
+        pass: true,
+      };
+    }
+    return {
+      message: () =>
+        `Reanimated: expected to have UI runtime check ${expectedMatchCount} times, but found ${receivedMatchCount}`,
+      pass: false,
+    };
+  },
+});

plugin/src/buildWorkletString.ts

@@ -6,11 +6,8 @@ import {
 } from '@babel/core';
 import generate from '@babel/generator';
 import {
-  ObjectMethod,
   isObjectMethod,
   FunctionDeclaration,
-  FunctionExpression,
-  ArrowFunctionExpression,
   identifier,
   Identifier,
   objectProperty,
@@ -34,6 +31,7 @@ import * as fs from 'fs';
 import * as convertSourceMap from 'convert-source-map';
 import { strict as assert } from 'assert';
 import { isRelease } from './utils';
+import { WorkletizableFunction } from './types';
 
 export function buildWorkletString(
   fun: BabelFile,
@@ -127,12 +125,7 @@ function shouldGenerateSourceMap() {
 }
 
 function prependClosure(
-  path: NodePath<
-    | FunctionDeclaration
-    | FunctionExpression
-    | ArrowFunctionExpression
-    | ObjectMethod
-  >,
+  path: NodePath<WorkletizableFunction>,
   closureVariables: Array<Identifier>,
   closureDeclaration: VariableDeclaration
 ) {
@@ -145,14 +138,7 @@ function prependClosure(
   }
 }
 
-function prependRecursiveDeclaration(
-  path: NodePath<
-    | FunctionDeclaration
-    | FunctionExpression
-    | ArrowFunctionExpression
-    | ObjectMethod
-  >
-) {
+function prependRecursiveDeclaration(path: NodePath<WorkletizableFunction>) {
   if (
     isProgram(path.parent) &&
     !isArrowFunctionExpression(path.node) &&
@@ -197,14 +183,7 @@ function prependClosureVariablesIfNecessary(
   return {
     visitor: {
       'FunctionDeclaration|FunctionExpression|ArrowFunctionExpression|ObjectMethod':
-        (
-          path: NodePath<
-            | FunctionDeclaration
-            | FunctionExpression
-            | ArrowFunctionExpression
-            | ObjectMethod
-          >
-        ) => {
+        (path: NodePath<WorkletizableFunction>) => {
           prependClosure(path, closureVariables, closureDeclaration);
           prependRecursiveDeclaration(path);
         },

plugin/src/makeWorklet.ts

@@ -1,11 +1,8 @@
 import { NodePath, transformSync, traverse } from '@babel/core';
 import generate from '@babel/generator';
 import {
-  ObjectMethod,
   isObjectMethod,
-  FunctionDeclaration,
   FunctionExpression,
-  ArrowFunctionExpression,
   identifier,
   Identifier,
   objectProperty,
@@ -36,7 +33,7 @@ import {
   isIdentifier,
   File as BabelFile,
 } from '@babel/types';
-import { ReanimatedPluginPass } from './types';
+import { ReanimatedPluginPass, WorkletizableFunction } from './types';
 import { isRelease } from './utils';
 import { strict as assert } from 'assert';
 import { globals } from './commonObjects';
@@ -45,28 +42,18 @@ import { buildWorkletString } from './buildWorkletString';
 
 const version = require('../../package.json').version;
 
+export const UI_RUNTIME_CHECK_REGEX = /(global.)?_WORKLET/g;
+
 export function makeWorklet(
-  fun: NodePath<
-    | FunctionDeclaration
-    | FunctionExpression
-    | ObjectMethod
-    | ArrowFunctionExpression
-  >,
+  fun: NodePath<WorkletizableFunction>,
   state: ReanimatedPluginPass
 ): FunctionExpression {
   // Returns a new FunctionExpression which is a workletized version of provided
   // FunctionDeclaration, FunctionExpression, ArrowFunctionExpression or ObjectMethod.
 
   const functionName = makeWorkletName(fun);
 
-  // remove 'worklet'; directive before generating string
-  fun.traverse({
-    DirectiveLiteral(path) {
-      if (path.node.value === 'worklet' && path.getFunctionParent() === fun) {
-        path.parentPath.remove();
-      }
-    },
-  });
+  removeWorkletDirective(fun);
 
   // We use copy because some of the plugins don't update bindings and
   // some even break them
@@ -77,6 +64,10 @@ export function makeWorklet(
     sourceFileName: state.file.opts.filename,
   });
 
+  // We change _WORKLET to true to simplify conditionals and avoid
+  // situations when _WORKLET is not yet defined but referenced.
+  codeObject.code = codeObject.code.replace(UI_RUNTIME_CHECK_REGEX, 'true');
+
   // We need to add a newline at the end, because there could potentially be a
   // comment after the function that gets included here, and then the closing
   // bracket would become part of the comment thus resulting in an error, since
@@ -260,6 +251,17 @@ export function makeWorklet(
   return newFun;
 }
 
+function removeWorkletDirective(fun: NodePath<WorkletizableFunction>) {
+  fun.traverse({
+    // Remove 'worklet'; directive before generating string.
+    DirectiveLiteral(path) {
+      if (path.node.value === 'worklet' && path.getFunctionParent() === fun) {
+        path.parentPath.remove();
+      }
+    },
+  });
+}
+
 function shouldInjectVersion() {
   // We don't inject version in release since cache is reset there anyway
   if (isRelease()) {
@@ -292,14 +294,7 @@ function hash(str: string) {
   return (hash1 >>> 0) * 4096 + (hash2 >>> 0);
 }
 
-function makeWorkletName(
-  fun: NodePath<
-    | FunctionDeclaration
-    | FunctionExpression
-    | ObjectMethod
-    | ArrowFunctionExpression
-  >
-) {
+function makeWorkletName(fun: NodePath<WorkletizableFunction>) {
   if (isObjectMethod(fun.node) && 'name' in fun.node.key) {
     return fun.node.key.name;
   }
@@ -314,12 +309,7 @@ function makeWorkletName(
 
 function makeArrayFromCapturedBindings(
   ast: BabelFile,
-  fun: NodePath<
-    | FunctionDeclaration
-    | FunctionExpression
-    | ObjectMethod
-    | ArrowFunctionExpression
-  >
+  fun: NodePath<WorkletizableFunction>
 ) {
   const closure = new Map<string, Identifier>();
 

plugin/src/plugin.ts

@@ -1,13 +1,8 @@
 import { PluginItem, NodePath } from '@babel/core';
 import { globals } from './commonObjects';
-import {
-  CallExpression,
-  FunctionDeclaration,
-  FunctionExpression,
-  ArrowFunctionExpression,
-} from '@babel/types';
+import { CallExpression } from '@babel/types';
 import { processForCalleesWorklets } from './processForCalleesWorklets';
-import { ReanimatedPluginPass } from './types';
+import { ExplicitWorklet, ReanimatedPluginPass } from './types';
 import { processIfWorkletNode } from './processIfWorkletNode';
 import { processIfGestureHandlerEventCallbackFunctionNode } from './processIfGestureHandlerEventCallbackFunctionNode';
 import { processInlineStylesWarning } from './processInlineStylesWarning';
@@ -29,12 +24,7 @@ module.exports = function (): PluginItem {
         },
       },
       'FunctionDeclaration|FunctionExpression|ArrowFunctionExpression': {
-        enter(
-          path: NodePath<
-            FunctionDeclaration | FunctionExpression | ArrowFunctionExpression
-          >,
-          state: ReanimatedPluginPass
-        ) {
+        enter(path: NodePath<ExplicitWorklet>, state: ReanimatedPluginPass) {
           processIfWorkletNode(path, state);
           processIfGestureHandlerEventCallbackFunctionNode(path, state);
         },

plugin/src/processIfGestureHandlerEventCallbackFunctionNode.ts

@@ -1,16 +1,13 @@
 import { NodePath } from '@babel/core';
 import {
-  FunctionDeclaration,
-  FunctionExpression,
-  ArrowFunctionExpression,
   isIdentifier,
   isCallExpression,
   Expression,
   isMemberExpression,
   isExpression,
 } from '@babel/types';
 import { processIfWorkletFunction } from './processIfWorkletFunction';
-import { ReanimatedPluginPass } from './types';
+import { ExplicitWorklet, ReanimatedPluginPass } from './types';
 
 const gestureHandlerGestureObjects = new Set([
   // from https://github.com/software-mansion/react-native-gesture-handler/blob/new-api/src/handlers/gestures/gestureObjects.ts
@@ -89,9 +86,7 @@ const gestureHandlerBuilderMethods = new Set([
   )
   */
 export function processIfGestureHandlerEventCallbackFunctionNode(
-  path: NodePath<
-    FunctionDeclaration | FunctionExpression | ArrowFunctionExpression
-  >,
+  path: NodePath<ExplicitWorklet>,
   state: ReanimatedPluginPass
 ) {
   if (

plugin/src/processIfWorkletFunction.ts

@@ -1,15 +1,12 @@
 import { NodePath, Node } from '@babel/core';
 import {
-  FunctionDeclaration,
-  FunctionExpression,
-  ArrowFunctionExpression,
   callExpression,
   isScopable,
   isExportNamedDeclaration,
   variableDeclaration,
   variableDeclarator,
 } from '@babel/types';
-import { ReanimatedPluginPass } from './types';
+import { ExplicitWorklet, ReanimatedPluginPass } from './types';
 import { makeWorklet } from './makeWorklet';
 
 // Replaces FunctionDeclaration, FunctionExpression or ArrowFunctionExpression
@@ -29,9 +26,7 @@ export function processIfWorkletFunction(
 }
 
 function processWorkletFunction(
-  path: NodePath<
-    FunctionDeclaration | FunctionExpression | ArrowFunctionExpression
-  >,
+  path: NodePath<ExplicitWorklet>,
   state: ReanimatedPluginPass
 ) {
   const newFun = makeWorklet(path, state);

plugin/src/processIfWorkletNode.ts

@@ -1,18 +1,10 @@
 import { NodePath } from '@babel/core';
-import {
-  FunctionDeclaration,
-  FunctionExpression,
-  ArrowFunctionExpression,
-  isBlockStatement,
-  isDirectiveLiteral,
-} from '@babel/types';
+import { isBlockStatement, isDirectiveLiteral } from '@babel/types';
 import { processIfWorkletFunction } from './processIfWorkletFunction';
-import { ReanimatedPluginPass } from './types';
+import { ExplicitWorklet, ReanimatedPluginPass } from './types';
 
 export function processIfWorkletNode(
-  fun: NodePath<
-    FunctionDeclaration | FunctionExpression | ArrowFunctionExpression
-  >,
+  fun: NodePath<ExplicitWorklet>,
   state: ReanimatedPluginPass
 ) {
   fun.traverse({

plugin/src/types.ts

@@ -1,4 +1,10 @@
 import { BabelFile } from '@babel/core';
+import {
+  FunctionDeclaration,
+  FunctionExpression,
+  ObjectMethod,
+  ArrowFunctionExpression,
+} from '@babel/types';
 
 export interface ReanimatedPluginPass {
   file: BabelFile;
@@ -13,3 +19,10 @@ export interface ReanimatedPluginPass {
   set(key: unknown, value: unknown): void;
   [key: string]: unknown;
 }
+
+export type ExplicitWorklet =
+  | FunctionDeclaration
+  | FunctionExpression
+  | ArrowFunctionExpression;
+
+export type WorkletizableFunction = ExplicitWorklet | ObjectMethod;