feat: Add fallback detection methods for Worklet Classes (#6706)

## Summary

Sometimes the class property

```javascript
	class Clazz {
		__workletClass = true;
	};
```

Can get stripped away by some Babel plugins. We add some additional
fallback methods so the users doesn't have to change his Babel pipeline
to use this feature.

## Test plan

🚀

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

@@ -43,12 +43,7 @@ export function processIfWorkletClass(
   classPath: NodePath<ClassDeclaration>,
   state: ReanimatedPluginPass
 ): boolean {
-  if (!classPath.node.id) {
-    // We don't support unnamed classes yet.
-    return false;
-  }
-
-  if (!hasWorkletClassMarker(classPath.node.body)) {
+  if (!isWorkletizableClass(classPath, state)) {
     return false;
   }
 
@@ -342,3 +337,42 @@ type Polyfill = {
   index: number;
   dependencies: Set<string>;
 };
+
+function isWorkletizableClass(
+  classPath: NodePath<ClassDeclaration>,
+  state: ReanimatedPluginPass
+): boolean {
+  const className = classPath.node.id?.name;
+  const classNode = classPath.node;
+
+  // We don't support unnamed classes yet.
+  if (!className) {
+    return false;
+  }
+
+  // Primary method of determining if a class is workletizable. However, some
+  // Babel plugins might remove Class Properties.
+  const isMarked = hasWorkletClassMarker(classNode.body);
+
+  // Secondary method of determining if a class is workletizable. We look for the
+  // reference we memoized earlier. However, some plugin could've changed the reference.
+  const isMemoizedNode = state.classesToWorkletize.some(
+    (record) => record.node === classNode
+  );
+
+  // Fallback for the name of the class.
+  // We bail on non-top-level declarations.
+  const isTopLevelMemoizedName =
+    classPath.parentPath.isProgram() &&
+    state.classesToWorkletize.some((record) => record.name === className);
+
+  // Remove the class from the list of classes to workletize. There are some edge
+  // cases when leaving it as is would lead to multiple workletizations.
+  state.classesToWorkletize = state.classesToWorkletize.filter(
+    (record) => record.node !== classNode && record.name !== className
+  );
+
+  const result = isMarked || isMemoizedNode || isTopLevelMemoizedName;
+
+  return result;
+}

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

@@ -36,7 +36,7 @@ import { contextObjectMarker } from './contextObject';
 
 export function processIfWorkletFile(
   path: NodePath<Program>,
-  _state: ReanimatedPluginPass
+  state: ReanimatedPluginPass
 ): boolean {
   if (
     !path.node.directives.some(
@@ -50,18 +50,21 @@ export function processIfWorkletFile(
   path.node.directives = path.node.directives.filter(
     (functionDirective) => functionDirective.value.value !== 'worklet'
   );
-  processWorkletFile(path);
+  processWorkletFile(path, state);
 
   return true;
 }
 
 /** Adds a worklet directive to each viable top-level entity in the file. */
-function processWorkletFile(programPath: NodePath<Program>) {
+function processWorkletFile(
+  programPath: NodePath<Program>,
+  state: ReanimatedPluginPass
+) {
   const statements = programPath.get('body');
   dehoistCommonJSExports(programPath.node);
   statements.forEach((statement) => {
     const candidatePath = getCandidate(statement);
-    processWorkletizableEntity(candidatePath);
+    processWorkletizableEntity(candidatePath, state);
   });
 }
 
@@ -76,7 +79,10 @@ function getCandidate(statementPath: NodePath<Statement>) {
   }
 }
 
-function processWorkletizableEntity(nodePath: NodePath<unknown>) {
+function processWorkletizableEntity(
+  nodePath: NodePath<unknown>,
+  state: ReanimatedPluginPass
+) {
   if (isWorkletizableFunctionPath(nodePath)) {
     if (nodePath.isArrowFunctionExpression()) {
       replaceImplicitReturnWithBlock(nodePath.node);
@@ -86,35 +92,46 @@ function processWorkletizableEntity(nodePath: NodePath<unknown>) {
     if (isImplicitContextObject(nodePath)) {
       appendWorkletContextObjectMarker(nodePath.node);
     } else {
-      processWorkletAggregator(nodePath);
+      processWorkletAggregator(nodePath, state);
     }
   } else if (nodePath.isVariableDeclaration()) {
-    processVariableDeclaration(nodePath);
+    processVariableDeclaration(nodePath, state);
   } else if (nodePath.isClassDeclaration()) {
     appendWorkletClassMarker(nodePath.node.body);
+    if (nodePath.node.id?.name) {
+      // We don't support unnamed classes yet.
+      state.classesToWorkletize.push({
+        node: nodePath.node,
+        name: nodePath.node.id.name,
+      });
+    }
   }
 }
 
 function processVariableDeclaration(
-  variableDeclarationPath: NodePath<VariableDeclaration>
+  variableDeclarationPath: NodePath<VariableDeclaration>,
+  state: ReanimatedPluginPass
 ) {
   const declarations = variableDeclarationPath.get('declarations');
   declarations.forEach((declaration) => {
     const initPath = declaration.get('init');
     if (initPath.isExpression()) {
-      processWorkletizableEntity(initPath);
+      processWorkletizableEntity(initPath, state);
     }
   });
 }
 
-function processWorkletAggregator(objectPath: NodePath<ObjectExpression>) {
+function processWorkletAggregator(
+  objectPath: NodePath<ObjectExpression>,
+  state: ReanimatedPluginPass
+) {
   const properties = objectPath.get('properties');
   properties.forEach((property) => {
     if (property.isObjectMethod()) {
       appendWorkletDirective(property.node.body);
     } else if (property.isObjectProperty()) {
       const valuePath = property.get('value');
-      processWorkletizableEntity(valuePath);
+      processWorkletizableEntity(valuePath, state);
     }
   });
 }

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

@@ -156,6 +156,7 @@ const notCapturedIdentifiers_DEPRECATED = [
 
 export function initializeState(state: ReanimatedPluginPass) {
   state.workletNumber = 1;
+  state.classesToWorkletize = [];
   initializeGlobals();
   addCustomGlobals(state);
 }

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

@@ -32,6 +32,7 @@ export interface ReanimatedPluginPass {
   cwd: string;
   filename: string | undefined;
   workletNumber: number;
+  classesToWorkletize: { node: BabelNode; name: string }[];
 }
 
 export type WorkletizableFunction =