[babel] Add gating module support
Instead of replacing original function with compiled code, this adds an option to append the code and switch between the two based on an `isForgetEnabled` test condition that's imported from the specified gatingModule.
Sathya Gunasekaran committed
Mar 6, 2023 at 19:54 UTC
39898c1858959f31e2fd89756c28eb8409892032
4 files changed
+151
-5
compiler/forget/src/Babel/BabelPlugin.ts
+65
-3
@@ -11,6 +11,7 @@ import type * as BabelCore from "@babel/core";
11
import generate from "@babel/generator";
12
import jsx from "@babel/plugin-syntax-jsx";
13
import * as t from "@babel/types";
14
+import invariant from "invariant";
15
import prettier from "prettier";
16
import { compile } from "../CompilerPipeline";
17
import { parsePluginOptions, PluginOptions } from "./PluginOptions";
@@ -19,6 +20,8 @@ type BabelPluginPass = {
20
opts: PluginOptions;
21
};
22
23
+const testId = t.identifier("isForgetEnabled");
24
+
25
function hasUseForgetDirective(directives: t.Directive[]): boolean {
26
for (const directive of directives) {
27
if (directive.value.value === "use forget") {
@@ -36,6 +39,8 @@ function hasUseForgetDirective(directives: t.Directive[]): boolean {
39
export default function ReactForgetBabelPlugin(
40
_babel: typeof BabelCore
41
): BabelCore.PluginObj {
42
+ let hasForgetCompiledCode: boolean = false;
43
+
44
const visitor = {
45
FunctionDeclaration(
46
fn: BabelCore.NodePath<t.FunctionDeclaration>,
@@ -50,12 +55,37 @@ export default function ReactForgetBabelPlugin(
55
if (fn.scope.getProgramParent() !== fn.scope.parent) {
56
return;
57
}
58
+
59
+ hasForgetCompiledCode = true;
60
const ast = compile(fn, pass.opts.environment);
61
55
- // We are generating a new FunctionDeclaration node, so we must skip over it or this
56
- // traversal will loop infinitely.
62
try {
58
- fn.replaceWith(ast);
63
+ if (pass.opts.gatingModule) {
64
+ // Rename existing function
65
+ invariant(fn.node.id, "FunctionDeclaration must have a name");
66
+ const original = fn.node.id;
67
+ fn.node.id = addSuffix(fn.node.id, "_uncompiled");
68
+
69
+ // Rename and append compiled function
70
+ invariant(ast.id, "FunctionDeclaration must produce a name");
71
+ ast.id = addSuffix(ast.id, "_forget");
72
+ const compiledFn = fn.insertAfter(ast)[0];
73
+ compiledFn.skip();
74
+
75
+ // Build and append gating test
76
+ compiledFn.insertAfter(
77
+ buildTest({
78
+ compiled: ast.id,
79
+ uncompiled: fn.node.id,
80
+ original,
81
+ })
82
+ );
83
+ } else {
84
+ fn.replaceWith(ast);
85
+ }
86
+
87
+ // We are generating a new FunctionDeclaration node, so we must skip over it or this
88
+ // traversal will loop infinitely.
89
fn.skip();
90
} catch (err) {
91
const result = generate(ast);
@@ -82,6 +112,13 @@ export default function ReactForgetBabelPlugin(
112
...pass,
113
opts: { ...pass.opts, ...options },
114
});
115
+
116
+ if (options.gatingModule && hasForgetCompiledCode) {
117
+ path.unshiftContainer(
118
+ "body",
119
+ buildImportForGatingModule(options.gatingModule)
120
+ );
121
+ }
122
} catch (err) {
123
if (options.logger && err) {
124
options.logger.logEvent("err", err);
@@ -92,3 +129,28 @@ export default function ReactForgetBabelPlugin(
129
},
130
};
131
}
132
+
133
+function addSuffix(id: t.Identifier, suffix: string): t.Identifier {
134
+ return t.identifier(`${id.name}${suffix}`);
135
+}
136
+
137
+function buildTest(ids: {
138
+ uncompiled: t.Identifier;
139
+ compiled: t.Identifier;
140
+ original: t.Identifier;
141
+}): t.Declaration {
142
+ return t.variableDeclaration("var", [
143
+ t.variableDeclarator(
144
+ ids.original,
145
+ t.conditionalExpression(testId, ids.compiled, ids.uncompiled)
146
+ ),
147
+ ]);
148
+}
149
+
150
+function buildImportForGatingModule(gatingModule: string) {
151
+ const importDefaultSpecifier = t.importDefaultSpecifier(testId);
152
+ return t.importDeclaration(
153
+ [importDefaultSpecifier],
154
+ t.stringLiteral(gatingModule)
155
+ );
156
+}
compiler/forget/src/__tests__/fixtures/hir/gating-test.expect.md
new
+69
@@ -0,0 +1,69 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+// @gatingModule @forgetDirective
6
+function Bar(props) {
7
+ "use forget";
8
+ return <div>{props.bar}</div>;
9
+}
10
+
11
+function NoForget(props) {
12
+ return <Bar>{props.noForget}</Bar>;
13
+}
14
+
15
+function Foo(props) {
16
+ "use forget";
17
+ return <Foo>{props.bar}</Foo>;
18
+}
19
+
20
+```
21
+
22
+## Code
23
+
24
+```javascript
25
+import isForgetEnabled from "ReactForgetFeatureFlag"; // @gatingModule @forgetDirective
26
+function Bar_uncompiled(props) {
27
+ "use forget";
28
+ return <div>{props.bar}</div>;
29
+}
30
+function Bar_forget(props) {
31
+ const $ = React.unstable_useMemoCache(2);
32
+ const c_0 = $[0] !== props.bar;
33
+ let t0;
34
+ if (c_0) {
35
+ t0 = <div>{props.bar}</div>;
36
+ $[0] = props.bar;
37
+ $[1] = t0;
38
+ } else {
39
+ t0 = $[1];
40
+ }
41
+ return t0;
42
+}
43
+var Bar = isForgetEnabled ? Bar_forget : Bar_uncompiled;
44
+
45
+function NoForget(props) {
46
+ return <Bar>{props.noForget}</Bar>;
47
+}
48
+
49
+function Foo_uncompiled(props) {
50
+ "use forget";
51
+ return <Foo>{props.bar}</Foo>;
52
+}
53
+function Foo_forget(props) {
54
+ const $ = React.unstable_useMemoCache(2);
55
+ const c_0 = $[0] !== props.bar;
56
+ let t0;
57
+ if (c_0) {
58
+ t0 = <Foo>{props.bar}</Foo>;
59
+ $[0] = props.bar;
60
+ $[1] = t0;
61
+ } else {
62
+ t0 = $[1];
63
+ }
64
+ return t0;
65
+}
66
+var Foo = isForgetEnabled ? Foo_forget : Foo_uncompiled;
67
+
68
+```
69
+
\ No newline at end of file
compiler/forget/src/__tests__/fixtures/hir/gating-test.js
new
+14
@@ -0,0 +1,14 @@
1
+// @gatingModule @forgetDirective
2
+function Bar(props) {
3
+ "use forget";
4
+ return <div>{props.bar}</div>;
5
+}
6
+
7
+function NoForget(props) {
8
+ return <Bar>{props.noForget}</Bar>;
9
+}
10
+
11
+function Foo(props) {
12
+ "use forget";
13
+ return <Foo>{props.bar}</Foo>;
14
+}
compiler/forget/src/__tests__/hir-test.ts
+3
-2
@@ -62,7 +62,8 @@ describe("React Forget (HIR version)", () => {
62
try {
63
items.push({
64
js: runReactForgetBabelPlugin(input, file, language, {
65
- enableOnlyOnUseForgetDirective: false,
65
+ enableOnlyOnUseForgetDirective:
66
+ options.enableOnlyOnUseForgetDirective,
67
environment: {
68
customHooks: new Map([
69
[
@@ -77,7 +78,7 @@ describe("React Forget (HIR version)", () => {
78
]),
79
},
80
logger: null,
80
- gatingModule: null,
81
+ gatingModule: options.gatingModule,
82
}).code,
83
});
84
} catch (e) {