[codegen] Don't drop directives in simple arrow fn

The compiler has an optimisation where it transforms a simple arrow
function with only a return statement to a implicit arrow function.

In the case, there's a directive in this simple arrow function, the
directive gets dropped.

Instead of dropping the directive, the compiler should perform this
optimisation only if there are no directives.

ghstack-source-id: 514cd2440025986a2d6d950694a7339d779b09f2
Pull Request resolved: https://github.com/facebook/react-forget/pull/2848

compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts

@@ -1646,7 +1646,7 @@ function codegenInstructionValue(
       ).unwrap();
       if (instrValue.expr.type === "ArrowFunctionExpression") {
         let body: t.BlockStatement | t.Expression = fn.body;
-        if (body.body.length === 1) {
+        if (body.body.length === 1 && loweredFunc.directives.length == 0) {
           const stmt = body.body[0]!;
           if (stmt.type === "ReturnStatement" && stmt.argument != null) {
             body = stmt.argument;

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

@@ -0,0 +1,52 @@
+
+## Input
+
+```javascript
+function Component() {
+  "use strict";
+  let [count, setCount] = React.useState(0);
+  const update = () => {
+    "worklet";
+    setCount((count) => count + 1);
+  };
+  return <button onClick={update}>{count}</button>;
+}
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Component() {
+  "use strict";
+  const $ = useMemoCache(3);
+
+  const [count, setCount] = React.useState(0);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = () => {
+      "worklet";
+
+      setCount((count_0) => count_0 + 1);
+    };
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const update = t0;
+  let t1;
+  if ($[1] !== count) {
+    t1 = <button onClick={update}>{count}</button>;
+    $[1] = count;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/arrow-expr-directive.js

@@ -0,0 +1,9 @@
+function Component() {
+  "use strict";
+  let [count, setCount] = React.useState(0);
+  const update = () => {
+    "worklet";
+    setCount((count) => count + 1);
+  };
+  return <button onClick={update}>{count}</button>;
+}

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/arrow-function-one-line-directive.expect.md

@@ -0,0 +1,50 @@
+
+## Input
+
+```javascript
+function useFoo() {
+  const update = () => {
+    "worklet";
+    return 1;
+  };
+  return update;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [],
+  isComponent: false,
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function useFoo() {
+  const $ = useMemoCache(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = () => {
+      "worklet";
+      return 1;
+    };
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const update = t0;
+  return update;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [],
+  isComponent: false,
+};
+
+```
+      
+### Eval output
+(kind: ok) "[[ function params=0 ]]"

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/arrow-function-one-line-directive.js

@@ -0,0 +1,13 @@
+function useFoo() {
+  const update = () => {
+    "worklet";
+    return 1;
+  };
+  return update;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [],
+  isComponent: false,
+};