[hir] Add support for directives

Previously, we would drop directives inside a component or hook but this is problematic with reanimated which uses `'worklet'` to mark components from compilation. This PR adds a directive to HIRFunction and ReactiveFunction and codegens the directive add the end. No processing is done on the directives themselves. Babel seems to store the directives on a BlockStatement, rather than on the Function but I've stored it on the Function types because we only support compiling functions and the spec defines directives as occuring in the initial statement list of a function: > A Directive Prologue is the longest sequence of ExpressionStatements > occurring as the initial StatementListItems or ModuleItems of a > FunctionBody, a ScriptBody, or a ModuleBody and where each > ExpressionStatement in the sequence consists entirely of a > StringLiteral token followed by a semicolon.
2025-11-01 09:12:30 +00:00 · 2024-04-02 11:25:10 +01:00
parent 1c3313707f
commit d5b6e584fb
17 changed files with 160 additions and 0 deletions
@@ -168,6 +168,7 @@ export function lower(
    }
  });

+  let directives: string[] = [];
  const body = func.get("body");
  if (body.isExpression()) {
    const fallthrough = builder.reserve("block");
@@ -180,6 +181,7 @@ export function lower(
    builder.terminateWithContinuation(terminal, fallthrough);
  } else if (body.isBlockStatement()) {
    lowerStatement(builder, body);
+    directives = body.get("directives").map((d) => d.node.value.value);
  } else {
    builder.errors.push({
      reason: `Unexpected function body kind: ${body.type}}. This error is likely caused by a bug in React Compiler. Please file an issue`,
@@ -219,6 +221,7 @@ export function lower(
    loc: func.node.loc ?? GeneratedSource,
    env,
    effects: null,
+    directives,
  });
 }

@@ -55,6 +55,7 @@ export type ReactiveFunction = {
  async: boolean;
  body: ReactiveBlock;
  env: Environment;
+  directives: string[];
 };

 export type ReactiveScopeBlock = {
@@ -280,6 +281,7 @@ export type HIRFunction = {
  body: HIR;
  generator: boolean;
  async: boolean;
+  directives: string[];
 };

 export type FunctionEffect = {
@@ -64,6 +64,7 @@ export function printFunction(fn: HIRFunction): string {
    output.push(definition);
  }
  output.push(printHIR(fn.body));
+  output.push(...fn.directives);
  return output.join("\n");
 }

@@ -49,6 +49,7 @@ export function buildReactiveFunction(fn: HIRFunction): ReactiveFunction {
    async: fn.async,
    body,
    env: fn.env,
+    directives: fn.directives,
  };
 }

@@ -167,6 +167,9 @@ function codegenReactiveFunction(

  const params = fn.params.map((param) => convertParameter(param));
  const body: t.BlockStatement = codegenBlock(cx, fn.body);
+  body.directives = fn.directives.map((d) =>
+    t.directive(t.directiveLiteral(d))
+  );
  const statements = body.body;
  if (statements.length !== 0) {
    const last = statements[statements.length - 1];
@@ -28,6 +28,7 @@ import { useRenderCounter, shouldInstrument } from "react-forget-runtime";
 import { unstable_useMemoCache as useMemoCache } from "react"; // @instrumentForget @compilationMode(annotation) @gating
 const Bar = isForgetEnabled_Fixtures()
  ? function Bar(props) {
+      "use forget";
      if (__DEV__ && shouldInstrument)
        useRenderCounter("Bar", "/codegen-instrument-forget-gating-test.ts");
      const $ = useMemoCache(2);
@@ -51,6 +52,7 @@ function NoForget(props) {
 }
 const Foo = isForgetEnabled_Fixtures()
  ? function Foo(props) {
+      "use forget";
      if (__DEV__ && shouldInstrument)
        useRenderCounter("Foo", "/codegen-instrument-forget-gating-test.ts");
      const $ = useMemoCache(2);
@@ -27,6 +27,7 @@ import { useRenderCounter, shouldInstrument } from "react-forget-runtime";
 import { unstable_useMemoCache as useMemoCache } from "react"; // @instrumentForget @compilationMode(annotation)

 function Bar(props) {
+  "use forget";
  if (__DEV__ && shouldInstrument)
    useRenderCounter("Bar", "/codegen-instrument-forget-test.ts");
  const $ = useMemoCache(2);
@@ -46,6 +47,7 @@ function NoForget(props) {
 }

 function Foo(props) {
+  "use forget";
  if (__DEV__ && shouldInstrument)
    useRenderCounter("Foo", "/codegen-instrument-forget-test.ts");
  const $ = useMemoCache(2);
@@ -0,0 +1,64 @@
+
+## Input
+
+```javascript
+function Component() {
+  "use strict";
+  let [count, setCount] = React.useState(0);
+  function update() {
+    "worklet";
+    setCount((count) => count + 1);
+  }
+  return <button onClick={update}>{count}</button>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+  isComponent: true,
+};
+
+```
+
+## 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 = function update() {
+      "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;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+  isComponent: true,
+};
+
+```
+      
+### Eval output
+(kind: ok) <button>0</button>
@@ -0,0 +1,15 @@
+function Component() {
+  "use strict";
+  let [count, setCount] = React.useState(0);
+  function update() {
+    "worklet";
+    setCount((count) => count + 1);
+  }
+  return <button onClick={update}>{count}</button>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+  isComponent: true,
+};
@@ -26,6 +26,7 @@ import { isForgetEnabled_Fixtures } from "ReactForgetFeatureFlag";
 import { unstable_useMemoCache as useMemoCache } from "react"; // @gating @compilationMode(annotation)
 export default isForgetEnabled_Fixtures()
  ? function Bar(props) {
+      "use forget";
      const $ = useMemoCache(2);
      let t0;
      if ($[0] !== props.bar) {
@@ -47,6 +48,7 @@ function NoForget(props) {
 }
 const Foo = isForgetEnabled_Fixtures()
  ? function Foo(props) {
+      "use forget";
      const $ = useMemoCache(2);
      let t0;
      if ($[0] !== props.bar) {
@@ -26,6 +26,7 @@ import { isForgetEnabled_Fixtures } from "ReactForgetFeatureFlag";
 import { unstable_useMemoCache as useMemoCache } from "react"; // @gating @compilationMode(annotation)
 export default isForgetEnabled_Fixtures()
  ? function Bar(props) {
+      "use forget";
      const $ = useMemoCache(2);
      let t0;
      if ($[0] !== props.bar) {
@@ -48,6 +49,7 @@ function NoForget(props) {

 export const Foo = isForgetEnabled_Fixtures()
  ? function Foo(props) {
+      "use forget";
      const $ = useMemoCache(2);
      let t0;
      if ($[0] !== props.bar) {
@@ -26,6 +26,7 @@ import { isForgetEnabled_Fixtures } from "ReactForgetFeatureFlag";
 import { unstable_useMemoCache as useMemoCache } from "react"; // @gating @compilationMode(annotation)
 export const Bar = isForgetEnabled_Fixtures()
  ? function Bar(props) {
+      "use forget";
      const $ = useMemoCache(2);
      let t0;
      if ($[0] !== props.bar) {
@@ -48,6 +49,7 @@ export function NoForget(props) {

 export const Foo = isForgetEnabled_Fixtures()
  ? function Foo(props) {
+      "use forget";
      const $ = useMemoCache(2);
      let t0;
      if ($[0] !== props.bar) {
@@ -26,6 +26,7 @@ import { isForgetEnabled_Fixtures } from "ReactForgetFeatureFlag";
 import { unstable_useMemoCache as useMemoCache } from "react"; // @gating @compilationMode(annotation)
 const Bar = isForgetEnabled_Fixtures()
  ? function Bar(props) {
+      "use forget";
      const $ = useMemoCache(2);
      let t0;
      if ($[0] !== props.bar) {
@@ -47,6 +48,7 @@ function NoForget(props) {
 }
 const Foo = isForgetEnabled_Fixtures()
  ? function Foo(props) {
+      "use forget";
      const $ = useMemoCache(2);
      let t0;
      if ($[0] !== props.bar) {
@@ -21,6 +21,7 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { unstable_useMemoCache as useMemoCache } from "react"; // @ignoreUseNoForget
 function Component(prop) {
+  "use no forget";
  const $ = useMemoCache(4);
  let t0;
  if ($[0] !== prop.x) {
@@ -0,0 +1,46 @@
+
+## Input
+
+```javascript
+function Component() {
+  "use foo";
+  "use bar";
+  return <div>"foo"</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+  isComponent: true,
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Component() {
+  "use foo";
+  "use bar";
+  const $ = useMemoCache(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = <div>"foo"</div>;
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+  isComponent: true,
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>"foo"</div>
@@ -0,0 +1,11 @@
+function Component() {
+  "use foo";
+  "use bar";
+  return <div>"foo"</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+  isComponent: true,
+};
@@ -21,6 +21,7 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { unstable_useMemoCache as useMemoCache } from "react";
 function Component(props) {
+  "use memo";
  const $ = useMemoCache(4);
  let t0;
  if ($[0] !== props.foo) {