fix: outlining worklet functions

compiler/packages/babel-plugin-react-compiler/src/Optimization/OutlineFunctions.ts

@@ -19,8 +19,15 @@ export function outlineFunctions(
         value.kind === 'FunctionExpression' ||
         value.kind === 'ObjectMethod'
       ) {
-        // Recurse in case there are inner functions which can be outlined
-        outlineFunctions(value.loweredFunc.func, fbtOperands);
+        // Don't outline inner functions from parents with 'worklet' directive as it might break multi-threading assumptions.
+        const canOutlineInnerFunctions = value.loweredFunc.func.directives
+          ? !value.loweredFunc.func.directives.includes('worklet')
+          : true;
+        if (canOutlineInnerFunctions) {
+         {
+          // Recurse in case there are inner functions which can be outlined
+          outlineFunctions(value.loweredFunc.func, fbtOperands);
+        }
       }
       if (
         value.kind === 'FunctionExpression' &&

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/arrow-expr-directive.expect.md

@@ -28,7 +28,7 @@ function Component() {
     t0 = () => {
       "worklet";
 
-      setCount(_temp);
+      setCount((count_0) => count_0 + 1);
     };
     $[0] = t0;
   } else {
@@ -45,9 +45,6 @@ function Component() {
   }
   return t1;
 }
-function _temp(count_0) {
-  return count_0 + 1;
-}
 
 ```
       

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/function-expr-directive.expect.md

@@ -34,7 +34,7 @@ function Component() {
     t0 = function update() {
       "worklet";
 
-      setCount(_temp);
+      setCount((count_0) => count_0 + 1);
     };
     $[0] = t0;
   } else {
@@ -51,9 +51,6 @@ function Component() {
   }
   return t1;
 }
-function _temp(count_0) {
-  return count_0 + 1;
-}
 
 export const FIXTURE_ENTRYPOINT = {
   fn: Component,

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/worklet-directive-skips-outlining.expect.md

@@ -0,0 +1,48 @@
+
+## Input
+
+```javascript
+const useSomeHook = () => {};
+
+const Component = () => {
+  useSomeHook(() => {
+    'worklet';
+    return [1, 2, 3].map(() => null);
+  });
+
+  return null;
+};
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+  isComponent: true,
+};
+
+```
+
+## Code
+
+```javascript
+const useSomeHook = () => {};
+
+const Component = () => {
+  useSomeHook(_temp);
+
+  return null;
+};
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+  isComponent: true,
+};
+function _temp() {
+  "worklet";
+  return [1, 2, 3].map(() => null);
+}
+
+```
+
+### Eval output
+(kind: ok) null

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/worklet-directive-skips-outlining.js

@@ -0,0 +1,16 @@
+const useSomeHook = () => {};
+
+const Component = () => {
+  useSomeHook(() => {
+    'worklet';
+    return [1, 2, 3].map(() => null);
+  });
+
+  return null;
+};
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+  isComponent: true,
+};