[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.
Sathya Gunasekaran committed
Apr 2, 2024 at 11:25 UTC
d5b6e584fbca6ae7dbdda8d2bbf16cf0af3831d4
17 files changed
+160
compiler/packages/babel-plugin-react-forget/src/HIR/BuildHIR.ts
+3
@@ -168,6 +168,7 @@ export function lower(
168
}
169
});
170
171
+ let directives: string[] = [];
172
const body = func.get("body");
173
if (body.isExpression()) {
174
const fallthrough = builder.reserve("block");
@@ -180,6 +181,7 @@ export function lower(
181
builder.terminateWithContinuation(terminal, fallthrough);
182
} else if (body.isBlockStatement()) {
183
lowerStatement(builder, body);
184
+ directives = body.get("directives").map((d) => d.node.value.value);
185
} else {
186
builder.errors.push({
187
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(
221
loc: func.node.loc ?? GeneratedSource,
222
env,
223
effects: null,
224
+ directives,
225
});
226
}
227
compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts
+2
@@ -55,6 +55,7 @@ export type ReactiveFunction = {
55
async: boolean;
56
body: ReactiveBlock;
57
env: Environment;
58
+ directives: string[];
59
};
60
61
export type ReactiveScopeBlock = {
@@ -280,6 +281,7 @@ export type HIRFunction = {
281
body: HIR;
282
generator: boolean;
283
async: boolean;
284
+ directives: string[];
285
};
286
287
export type FunctionEffect = {
compiler/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts
+1
@@ -64,6 +64,7 @@ export function printFunction(fn: HIRFunction): string {
64
output.push(definition);
65
}
66
output.push(printHIR(fn.body));
67
+ output.push(...fn.directives);
68
return output.join("\n");
69
}
70
compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/BuildReactiveFunction.ts
+1
@@ -49,6 +49,7 @@ export function buildReactiveFunction(fn: HIRFunction): ReactiveFunction {
49
async: fn.async,
50
body,
51
env: fn.env,
52
+ directives: fn.directives,
53
};
54
}
55
compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts
+3
@@ -167,6 +167,9 @@ function codegenReactiveFunction(
167
168
const params = fn.params.map((param) => convertParameter(param));
169
const body: t.BlockStatement = codegenBlock(cx, fn.body);
170
+ body.directives = fn.directives.map((d) =>
171
+ t.directive(t.directiveLiteral(d))
172
+ );
173
const statements = body.body;
174
if (statements.length !== 0) {
175
const last = statements[statements.length - 1];
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/codegen-instrument-forget-gating-test.expect.md
+2
@@ -28,6 +28,7 @@ import { useRenderCounter, shouldInstrument } from "react-forget-runtime";
28
import { unstable_useMemoCache as useMemoCache } from "react"; // @instrumentForget @compilationMode(annotation) @gating
29
const Bar = isForgetEnabled_Fixtures()
30
? function Bar(props) {
31
+ "use forget";
32
if (__DEV__ && shouldInstrument)
33
useRenderCounter("Bar", "/codegen-instrument-forget-gating-test.ts");
34
const $ = useMemoCache(2);
@@ -51,6 +52,7 @@ function NoForget(props) {
52
}
53
const Foo = isForgetEnabled_Fixtures()
54
? function Foo(props) {
55
+ "use forget";
56
if (__DEV__ && shouldInstrument)
57
useRenderCounter("Foo", "/codegen-instrument-forget-gating-test.ts");
58
const $ = useMemoCache(2);
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/codegen-instrument-forget-test.expect.md
+2
@@ -27,6 +27,7 @@ import { useRenderCounter, shouldInstrument } from "react-forget-runtime";
27
import { unstable_useMemoCache as useMemoCache } from "react"; // @instrumentForget @compilationMode(annotation)
28
29
function Bar(props) {
30
+ "use forget";
31
if (__DEV__ && shouldInstrument)
32
useRenderCounter("Bar", "/codegen-instrument-forget-test.ts");
33
const $ = useMemoCache(2);
@@ -46,6 +47,7 @@ function NoForget(props) {
47
}
48
49
function Foo(props) {
50
+ "use forget";
51
if (__DEV__ && shouldInstrument)
52
useRenderCounter("Foo", "/codegen-instrument-forget-test.ts");
53
const $ = useMemoCache(2);
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/function-expr-directive.expect.md
new
+64
@@ -0,0 +1,64 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+function Component() {
6
+ "use strict";
7
+ let [count, setCount] = React.useState(0);
8
+ function update() {
9
+ "worklet";
10
+ setCount((count) => count + 1);
11
+ }
12
+ return <button onClick={update}>{count}</button>;
13
+}
14
+
15
+export const FIXTURE_ENTRYPOINT = {
16
+ fn: Component,
17
+ params: [],
18
+ isComponent: true,
19
+};
20
+
21
+```
22
+
23
+## Code
24
+
25
+```javascript
26
+import { unstable_useMemoCache as useMemoCache } from "react";
27
+function Component() {
28
+ "use strict";
29
+ const $ = useMemoCache(3);
30
+
31
+ const [count, setCount] = React.useState(0);
32
+ let t0;
33
+ if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
34
+ t0 = function update() {
35
+ "worklet";
36
+
37
+ setCount((count_0) => count_0 + 1);
38
+ };
39
+ $[0] = t0;
40
+ } else {
41
+ t0 = $[0];
42
+ }
43
+ const update = t0;
44
+ let t1;
45
+ if ($[1] !== count) {
46
+ t1 = <button onClick={update}>{count}</button>;
47
+ $[1] = count;
48
+ $[2] = t1;
49
+ } else {
50
+ t1 = $[2];
51
+ }
52
+ return t1;
53
+}
54
+
55
+export const FIXTURE_ENTRYPOINT = {
56
+ fn: Component,
57
+ params: [],
58
+ isComponent: true,
59
+};
60
+
61
+```
62
+
63
+### Eval output
64
+(kind: ok) <button>0</button>
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/function-expr-directive.js
new
+15
@@ -0,0 +1,15 @@
1
+function Component() {
2
+ "use strict";
3
+ let [count, setCount] = React.useState(0);
4
+ function update() {
5
+ "worklet";
6
+ setCount((count) => count + 1);
7
+ }
8
+ return <button onClick={update}>{count}</button>;
9
+}
10
+
11
+export const FIXTURE_ENTRYPOINT = {
12
+ fn: Component,
13
+ params: [],
14
+ isComponent: true,
15
+};
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/gating-test-export-default-function.expect.md
+2
@@ -26,6 +26,7 @@ import { isForgetEnabled_Fixtures } from "ReactForgetFeatureFlag";
26
import { unstable_useMemoCache as useMemoCache } from "react"; // @gating @compilationMode(annotation)
27
export default isForgetEnabled_Fixtures()
28
? function Bar(props) {
29
+ "use forget";
30
const $ = useMemoCache(2);
31
let t0;
32
if ($[0] !== props.bar) {
@@ -47,6 +48,7 @@ function NoForget(props) {
48
}
49
const Foo = isForgetEnabled_Fixtures()
50
? function Foo(props) {
51
+ "use forget";
52
const $ = useMemoCache(2);
53
let t0;
54
if ($[0] !== props.bar) {
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/gating-test-export-function-and-default.expect.md
+2
@@ -26,6 +26,7 @@ import { isForgetEnabled_Fixtures } from "ReactForgetFeatureFlag";
26
import { unstable_useMemoCache as useMemoCache } from "react"; // @gating @compilationMode(annotation)
27
export default isForgetEnabled_Fixtures()
28
? function Bar(props) {
29
+ "use forget";
30
const $ = useMemoCache(2);
31
let t0;
32
if ($[0] !== props.bar) {
@@ -48,6 +49,7 @@ function NoForget(props) {
49
50
export const Foo = isForgetEnabled_Fixtures()
51
? function Foo(props) {
52
+ "use forget";
53
const $ = useMemoCache(2);
54
let t0;
55
if ($[0] !== props.bar) {
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/gating-test-export-function.expect.md
+2
@@ -26,6 +26,7 @@ import { isForgetEnabled_Fixtures } from "ReactForgetFeatureFlag";
26
import { unstable_useMemoCache as useMemoCache } from "react"; // @gating @compilationMode(annotation)
27
export const Bar = isForgetEnabled_Fixtures()
28
? function Bar(props) {
29
+ "use forget";
30
const $ = useMemoCache(2);
31
let t0;
32
if ($[0] !== props.bar) {
@@ -48,6 +49,7 @@ export function NoForget(props) {
49
50
export const Foo = isForgetEnabled_Fixtures()
51
? function Foo(props) {
52
+ "use forget";
53
const $ = useMemoCache(2);
54
let t0;
55
if ($[0] !== props.bar) {
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/gating-test.expect.md
+2
@@ -26,6 +26,7 @@ import { isForgetEnabled_Fixtures } from "ReactForgetFeatureFlag";
26
import { unstable_useMemoCache as useMemoCache } from "react"; // @gating @compilationMode(annotation)
27
const Bar = isForgetEnabled_Fixtures()
28
? function Bar(props) {
29
+ "use forget";
30
const $ = useMemoCache(2);
31
let t0;
32
if ($[0] !== props.bar) {
@@ -47,6 +48,7 @@ function NoForget(props) {
48
}
49
const Foo = isForgetEnabled_Fixtures()
50
? function Foo(props) {
51
+ "use forget";
52
const $ = useMemoCache(2);
53
let t0;
54
if ($[0] !== props.bar) {
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/ignore-use-no-forget.expect.md
+1
@@ -21,6 +21,7 @@ export const FIXTURE_ENTRYPOINT = {
21
```javascript
22
import { unstable_useMemoCache as useMemoCache } from "react"; // @ignoreUseNoForget
23
function Component(prop) {
24
+ "use no forget";
25
const $ = useMemoCache(4);
26
let t0;
27
if ($[0] !== prop.x) {
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/multi-directive.expect.md
new
+46
@@ -0,0 +1,46 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+function Component() {
6
+ "use foo";
7
+ "use bar";
8
+ return <div>"foo"</div>;
9
+}
10
+
11
+export const FIXTURE_ENTRYPOINT = {
12
+ fn: Component,
13
+ params: [],
14
+ isComponent: true,
15
+};
16
+
17
+```
18
+
19
+## Code
20
+
21
+```javascript
22
+import { unstable_useMemoCache as useMemoCache } from "react";
23
+function Component() {
24
+ "use foo";
25
+ "use bar";
26
+ const $ = useMemoCache(1);
27
+ let t0;
28
+ if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
29
+ t0 = <div>"foo"</div>;
30
+ $[0] = t0;
31
+ } else {
32
+ t0 = $[0];
33
+ }
34
+ return t0;
35
+}
36
+
37
+export const FIXTURE_ENTRYPOINT = {
38
+ fn: Component,
39
+ params: [],
40
+ isComponent: true,
41
+};
42
+
43
+```
44
+
45
+### Eval output
46
+(kind: ok) <div>"foo"</div>
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/multi-directive.js
new
+11
@@ -0,0 +1,11 @@
1
+function Component() {
2
+ "use foo";
3
+ "use bar";
4
+ return <div>"foo"</div>;
5
+}
6
+
7
+export const FIXTURE_ENTRYPOINT = {
8
+ fn: Component,
9
+ params: [],
10
+ isComponent: true,
11
+};
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-memo-simple.expect.md
+1
@@ -21,6 +21,7 @@ export const FIXTURE_ENTRYPOINT = {
21
```javascript
22
import { unstable_useMemoCache as useMemoCache } from "react";
23
function Component(props) {
24
+ "use memo";
25
const $ = useMemoCache(4);
26
let t0;
27
if ($[0] !== props.foo) {