fix: Allow TypeScript in worklet classes (#6667)

## Summary

Turns out `@babel/preset-typescript` always has to be included when
calling `transformSync` from `babel` (it cannot be included in a later
call). Due to that I made all `transformSync` calls into a wrapper call
that uses a required set of plugins and presets.

Fixes 
- #6642

## Test plan

- [x] Plugin unit tests pass
- [x] Runtime tests pass

Author: tjzel

apps/common-app/src/examples/RuntimeTests/tests/plugin/fileWorkletization.ts

@@ -23,12 +23,17 @@ export const implicitContextObject = {
   },
 };
 
-export class ImplicitWorkletClass {
-  getSix() {
+interface IWorkletClass {
+  getSix(): number;
+  getSeven(): number;
+}
+
+export class ImplicitWorkletClass implements IWorkletClass {
+  getSix(): number {
     return 6;
   }
 
-  getSeven() {
+  getSeven(): number {
     return this.getSix() + 1;
   }
 }

apps/common-app/src/examples/RuntimeTests/tests/plugin/recursion.test.tsx

@@ -1,6 +1,6 @@
 import React, { useEffect } from 'react';
 import { View } from 'react-native';
-import { useSharedValue, runOnUI, runOnJS } from 'react-native-reanimated';
+import { useSharedValue, runOnUI } from 'react-native-reanimated';
 import { render, wait, describe, getRegisteredValue, registerValue, test, expect } from '../../ReJest/RuntimeTestsApi';
 
 const SHARED_VALUE_REF = 'SHARED_VALUE_REF';
@@ -85,15 +85,13 @@ describe('Test recursion in worklets', () => {
       const output = useSharedValue<number | null>(null);
       registerValue(SHARED_VALUE_REF, output);
       function recursiveWorklet(a: number) {
-        if (a === 2) {
+        if (a === 1) {
           output.value = a;
-        } else if (a === 1) {
-          try {
-            // TODO: Such case isn't supported at the moment -
-            // a function can't be a Worklet and a Remote function at the same time.
-            // Consider supporting it in the future.
-            runOnJS(recursiveWorklet)(a + 1);
-          } catch {}
+        } else if (a === 2) {
+          // TODO: Such case isn't supported at the moment -
+          // a function can't be a Worklet and a Remote function at the same time.
+          // Consider supporting it in the future.
+          // runOnJS(recursiveWorklet)(a + 1);
         } else {
           recursiveWorklet(a + 1);
         }
@@ -108,6 +106,6 @@ describe('Test recursion in worklets', () => {
     await render(<ExampleComponent />);
     await wait(100);
     const sharedValue = await getRegisteredValue(SHARED_VALUE_REF);
-    expect(sharedValue.onJS).toBe(null);
+    expect(sharedValue.onJS).toBe(1);
   });
 });

apps/common-app/src/examples/RuntimeTests/tests/plugin/workletClasses.test.tsx

@@ -21,6 +21,28 @@ class WorkletClass {
   }
 }
 
+interface ITypeScriptClass {
+  getOne(): number;
+  getTwo(): number;
+  getIncremented(): number;
+}
+
+class TypeScriptClass implements ITypeScriptClass {
+  __workletClass: boolean = true;
+  value: number = 0;
+  getOne(): number {
+    return 1;
+  }
+
+  getTwo(): number {
+    return this.getOne() + 1;
+  }
+
+  getIncremented(): number {
+    return ++this.value;
+  }
+}
+
 describe('Test worklet classes', () => {
   test('class works on React runtime', async () => {
     const ExampleComponent = () => {
@@ -134,5 +156,26 @@ describe('Test worklet classes', () => {
     expect(sharedValue.onUI).toBe(true);
   });
 
+  test('TypeScript classes work on Worklet runtime', async () => {
+    const ExampleComponent = () => {
+      const output = useSharedValue<number | null>(null);
+      registerValue(SHARED_VALUE_REF, output);
+
+      useEffect(() => {
+        runOnUI(() => {
+          const clazz = new TypeScriptClass();
+          output.value = clazz.getOne();
+        })();
+      });
+
+      return <View />;
+    };
+    await render(<ExampleComponent />);
+    await wait(100);
+    const sharedValue = await getRegisteredValue(SHARED_VALUE_REF);
+    expect(sharedValue.onUI).toBe(1);
+  });
+
   // TODO: Add a test that throws when class is sent from React to Worklet runtime.
+  // TODO: Add a test that throws when trying to use Worklet Class with inheritance.
 });

packages/docs-reanimated/docs/reanimated-babel-plugin/about.md

@@ -146,6 +146,7 @@ runOnUI(() => new Clazz().foo())(); // Logs 'Hello from WorkletClass'
 
 **Pitfalls:**
 
+- Worklet Classes don't support inheritance.
 - Worklet Classes don't support static methods and properties.
 - Class instances cannot be shared between JS and UI threads.
 

packages/react-native-reanimated/plugin/index.js

@@ -1,5 +1,4 @@
 import type { NodePath } from '@babel/core';
-import { transformSync } from '@babel/core';
 import generate from '@babel/generator';
 import traverse from '@babel/traverse';
 import type {
@@ -36,6 +35,7 @@ import { strict as assert } from 'assert';
 import type { ReanimatedPluginPass } from './types';
 import { workletClassFactorySuffix } from './types';
 import { replaceWithFactoryCall } from './utils';
+import { workletTransformSync } from './transform';
 
 const classWorkletMarker = '__workletClass';
 
@@ -96,8 +96,8 @@ function getPolyfilledAst(
 ) {
   const classCode = generate(classNode).code;
 
-  const classWithPolyfills = transformSync(classCode, {
-    plugins: [
+  const classWithPolyfills = workletTransformSync(classCode, {
+    extraPlugins: [
       '@babel/plugin-transform-class-properties',
       '@babel/plugin-transform-classes',
       '@babel/plugin-transform-unicode-regex',

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

@@ -0,0 +1,28 @@
+import { transformSync } from '@babel/core';
+import type { PluginItem, TransformOptions } from '@babel/core';
+
+export function workletTransformSync(
+  code: string,
+  opts: WorkletTransformOptions
+) {
+  const { extraPlugins = [], extraPresets = [], ...rest } = opts;
+
+  return transformSync(code, {
+    ...rest,
+    plugins: [...defaultPlugins, ...extraPlugins],
+    presets: [...defaultPresets, ...extraPresets],
+  });
+}
+
+const defaultPresets: PluginItem[] = [
+  require.resolve('@babel/preset-typescript'),
+];
+
+const defaultPlugins: PluginItem[] = [];
+
+interface WorkletTransformOptions
+  extends Omit<TransformOptions, 'plugins' | 'presets'> {
+  extraPlugins?: PluginItem[];
+  extraPresets?: PluginItem[];
+  filename: TransformOptions['filename'];
+}

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

@@ -1,6 +1,6 @@
 /* eslint-disable @typescript-eslint/no-var-requires */
 import type { NodePath } from '@babel/core';
-import { transformSync, traverse } from '@babel/core';
+import { traverse } from '@babel/core';
 import generate from '@babel/generator';
 import type {
   File as BabelFile,
@@ -44,25 +44,11 @@ import type { ReanimatedPluginPass, WorkletizableFunction } from './types';
 import { workletClassFactorySuffix } from './types';
 import { isRelease } from './utils';
 import { buildWorkletString } from './workletStringCode';
+import { workletTransformSync } from './transform';
 
 const REAL_VERSION = require('../../package.json').version;
 const MOCK_VERSION = 'x.y.z';
 
-const workletStringTransformPresets = [
-  require.resolve('@babel/preset-typescript'),
-];
-
-const workletStringTransformPlugins = [
-  require.resolve('@babel/plugin-transform-shorthand-properties'),
-  require.resolve('@babel/plugin-transform-arrow-functions'),
-  require.resolve('@babel/plugin-transform-optional-chaining'),
-  require.resolve('@babel/plugin-transform-nullish-coalescing-operator'),
-  [
-    require.resolve('@babel/plugin-transform-template-literals'),
-    { loose: true },
-  ],
-];
-
 export function makeWorkletFactory(
   fun: NodePath<WorkletizableFunction>,
   state: ReanimatedPluginPass
@@ -91,10 +77,9 @@ export function makeWorkletFactory(
   codeObject.code =
     '(' + (fun.isObjectMethod() ? 'function ' : '') + codeObject.code + '\n)';
 
-  const transformed = transformSync(codeObject.code, {
+  const transformed = workletTransformSync(codeObject.code, {
+    extraPlugins,
     filename: state.file.opts.filename,
-    presets: workletStringTransformPresets,
-    plugins: workletStringTransformPlugins,
     ast: true,
     babelrc: false,
     configFile: false,
@@ -469,3 +454,14 @@ function makeArrayFromCapturedBindings(
 
   return Array.from(closure.values());
 }
+
+const extraPlugins = [
+  require.resolve('@babel/plugin-transform-shorthand-properties'),
+  require.resolve('@babel/plugin-transform-arrow-functions'),
+  require.resolve('@babel/plugin-transform-optional-chaining'),
+  require.resolve('@babel/plugin-transform-nullish-coalescing-operator'),
+  [
+    require.resolve('@babel/plugin-transform-template-literals'),
+    { loose: true },
+  ],
+];

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

@@ -1,5 +1,5 @@
 import type { BabelFileResult, NodePath, PluginItem } from '@babel/core';
-import { transformSync, traverse } from '@babel/core';
+import { traverse } from '@babel/core';
 import generate from '@babel/generator';
 import type {
   File as BabelFile,
@@ -34,6 +34,7 @@ import * as fs from 'fs';
 import type { ReanimatedPluginPass, WorkletizableFunction } from './types';
 import { workletClassFactorySuffix } from './types';
 import { isRelease } from './utils';
+import { workletTransformSync } from './transform';
 
 const MOCK_SOURCE_MAP = 'mock source map';
 
@@ -130,8 +131,9 @@ export function buildWorkletString(
     }
   }
 
-  const transformed = transformSync(code, {
-    plugins: [prependClosureVariablesIfNecessary(closureVariables)],
+  const transformed = workletTransformSync(code, {
+    filename: state.file.opts.filename,
+    extraPlugins: [getClosurePlugin(closureVariables)],
     compact: true,
     sourceMaps: includeSourceMap,
     inputSourceMap: inputMap,
@@ -222,9 +224,8 @@ function prependRecursiveDeclaration(path: NodePath<WorkletizableFunction>) {
   }
 }
 
-function prependClosureVariablesIfNecessary(
-  closureVariables: Array<Identifier>
-): PluginItem {
+/** Prepends necessary closure variables to the worklet function. */
+function getClosurePlugin(closureVariables: Array<Identifier>): PluginItem {
   const closureDeclaration = variableDeclaration('const', [
     variableDeclarator(
       objectPattern(