From d5b6e584fbca6ae7dbdda8d2bbf16cf0af3831d4 Mon Sep 17 00:00:00 2001 From: Sathya Gunasekaran Date: Tue, 2 Apr 2024 11:25:10 +0100 Subject: [PATCH] [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. --- .../src/HIR/BuildHIR.ts | 3 + .../babel-plugin-react-forget/src/HIR/HIR.ts | 2 + .../src/HIR/PrintHIR.ts | 1 + .../ReactiveScopes/BuildReactiveFunction.ts | 1 + .../ReactiveScopes/CodegenReactiveFunction.ts | 3 + ...en-instrument-forget-gating-test.expect.md | 2 + .../codegen-instrument-forget-test.expect.md | 2 + .../function-expr-directive.expect.md | 64 +++++++++++++++++++ .../compiler/function-expr-directive.js | 15 +++++ ...ing-test-export-default-function.expect.md | 2 + ...test-export-function-and-default.expect.md | 2 + .../gating-test-export-function.expect.md | 2 + .../fixtures/compiler/gating-test.expect.md | 2 + .../compiler/ignore-use-no-forget.expect.md | 1 + .../compiler/multi-directive.expect.md | 46 +++++++++++++ .../fixtures/compiler/multi-directive.js | 11 ++++ .../compiler/use-memo-simple.expect.md | 1 + 17 files changed, 160 insertions(+) create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/function-expr-directive.expect.md create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/function-expr-directive.js create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/multi-directive.expect.md create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/multi-directive.js diff --git a/compiler/packages/babel-plugin-react-forget/src/HIR/BuildHIR.ts b/compiler/packages/babel-plugin-react-forget/src/HIR/BuildHIR.ts index 14b6e0edb7..642736e32e 100644 --- a/compiler/packages/babel-plugin-react-forget/src/HIR/BuildHIR.ts +++ b/compiler/packages/babel-plugin-react-forget/src/HIR/BuildHIR.ts @@ -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, }); } diff --git a/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts b/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts index 3785dbaa38..69ab35fc12 100644 --- a/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts +++ b/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts @@ -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 = { diff --git a/compiler/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts b/compiler/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts index 3b0e74cb8c..11853968eb 100644 --- a/compiler/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts +++ b/compiler/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts @@ -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"); } diff --git a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/BuildReactiveFunction.ts b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/BuildReactiveFunction.ts index 1d96f689b5..689c9a5d4c 100644 --- a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/BuildReactiveFunction.ts +++ b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/BuildReactiveFunction.ts @@ -49,6 +49,7 @@ export function buildReactiveFunction(fn: HIRFunction): ReactiveFunction { async: fn.async, body, env: fn.env, + directives: fn.directives, }; } diff --git a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts index 7a0af6656c..31d13943bd 100644 --- a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts +++ b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts @@ -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]; diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/codegen-instrument-forget-gating-test.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/codegen-instrument-forget-gating-test.expect.md index 05fc4caddb..d9e4ddd27a 100644 --- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/codegen-instrument-forget-gating-test.expect.md +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/codegen-instrument-forget-gating-test.expect.md @@ -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); diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/codegen-instrument-forget-test.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/codegen-instrument-forget-test.expect.md index ddc79408f7..6cc46e68c6 100644 --- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/codegen-instrument-forget-test.expect.md +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/codegen-instrument-forget-test.expect.md @@ -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); diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/function-expr-directive.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/function-expr-directive.expect.md new file mode 100644 index 0000000000..ffe21f9b31 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/function-expr-directive.expect.md @@ -0,0 +1,64 @@ + +## Input + +```javascript +function Component() { + "use strict"; + let [count, setCount] = React.useState(0); + function update() { + "worklet"; + setCount((count) => count + 1); + } + return ; +} + +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 = ; + $[1] = count; + $[2] = t1; + } else { + t1 = $[2]; + } + return t1; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [], + isComponent: true, +}; + +``` + +### Eval output +(kind: ok) \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/function-expr-directive.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/function-expr-directive.js new file mode 100644 index 0000000000..f04ffddff4 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/function-expr-directive.js @@ -0,0 +1,15 @@ +function Component() { + "use strict"; + let [count, setCount] = React.useState(0); + function update() { + "worklet"; + setCount((count) => count + 1); + } + return ; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [], + isComponent: true, +}; diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/gating-test-export-default-function.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/gating-test-export-default-function.expect.md index d6a9af708d..5b3950fd5f 100644 --- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/gating-test-export-default-function.expect.md +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/gating-test-export-default-function.expect.md @@ -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) { diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/gating-test-export-function-and-default.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/gating-test-export-function-and-default.expect.md index afe7fbdd1d..aa0b4c4121 100644 --- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/gating-test-export-function-and-default.expect.md +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/gating-test-export-function-and-default.expect.md @@ -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) { diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/gating-test-export-function.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/gating-test-export-function.expect.md index c800dcd897..84ff4dbab4 100644 --- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/gating-test-export-function.expect.md +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/gating-test-export-function.expect.md @@ -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) { diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/gating-test.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/gating-test.expect.md index 11c1d920fd..fd10885e85 100644 --- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/gating-test.expect.md +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/gating-test.expect.md @@ -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) { diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/ignore-use-no-forget.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/ignore-use-no-forget.expect.md index 56de5f68fc..187eee864a 100644 --- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/ignore-use-no-forget.expect.md +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/ignore-use-no-forget.expect.md @@ -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) { diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/multi-directive.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/multi-directive.expect.md new file mode 100644 index 0000000000..894d310cc8 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/multi-directive.expect.md @@ -0,0 +1,46 @@ + +## Input + +```javascript +function Component() { + "use foo"; + "use bar"; + return
"foo"
; +} + +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 =
"foo"
; + $[0] = t0; + } else { + t0 = $[0]; + } + return t0; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [], + isComponent: true, +}; + +``` + +### Eval output +(kind: ok)
"foo"
\ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/multi-directive.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/multi-directive.js new file mode 100644 index 0000000000..2c027648be --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/multi-directive.js @@ -0,0 +1,11 @@ +function Component() { + "use foo"; + "use bar"; + return
"foo"
; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [], + isComponent: true, +}; diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-memo-simple.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-memo-simple.expect.md index d21466648d..23e87987ff 100644 --- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-memo-simple.expect.md +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-memo-simple.expect.md @@ -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) {