@samitouri / QOS-React-2 / commits / 0b697d1903

[babel plugin] Add codegen for useRenderCounter

Mofei Zhang committed May 17, 2023 at 14:39 UTC 0b697d1903318b75c1e61b749365a44d39d67ebb
10 files changed +338 -18
compiler/forget/packages/snap/src/compiler-worker.ts
+16
@@ -91,6 +91,7 @@ export async function compile(
91 }
92
93 let gating = null;
94 + let instrumentForget = null;
95 if (firstLine.indexOf("@gating") !== -1) {
96 gating = {
97 source: "ReactForgetFeatureFlag",
@@ -98,6 +99,19 @@ export async function compile(
99 };
100 }
101
102 + if (firstLine.indexOf("@instrumentForget") !== -1) {
103 + instrumentForget = {
104 + gating: {
105 + source: "ReactInstrumentForgetFeatureFlag",
106 + importSpecifierName: "isInstrumentForgetEnabled_Fixtures",
107 + },
108 + instrumentFn: {
109 + source: "react-forget-runtime",
110 + importSpecifierName: "useRenderCounter",
111 + },
112 + };
113 + }
114 +
115 let panicOnBailout = true;
116 if (firstLine.indexOf("@panicOnBailout false") !== -1) {
117 panicOnBailout = false;
@@ -120,9 +134,11 @@ export async function compile(
134 ],
135 ]),
136 validateHooksUsage: true,
137 + inlineUseMemo: true,
138 },
139 logger: null,
140 gating,
141 + instrumentForget,
142 panicOnBailout,
143 isDev: true,
144 }).code;
compiler/forget/src/Babel/BabelPlugin.ts
+58 -14
@@ -17,8 +17,9 @@ import {
17 } from "../CompilerError";
18 import { compile } from "../CompilerPipeline";
19 import { GeneratedSource } from "../HIR";
20 +import { addInstrumentForget } from "./InstrumentForgetBabelPlugin";
21 import {
21 - GatingOptions,
22 + ExternalFunction,
23 PluginOptions,
24 parsePluginOptions,
25 } from "./PluginOptions";
@@ -57,16 +58,16 @@ export default function ReactForgetBabelPlugin(
58 ): void {
59 try {
60 const compiled = compile(fn, pass.opts.environment);
61 + if (fn.node.id == null) {
62 + CompilerError.invariant(
63 + "FunctionDeclaration must have a name",
64 + fn.node.loc ?? GeneratedSource
65 + );
66 + }
67 + const originalIdent = fn.node.id;
68
69 if (pass.opts.gating != null) {
70 // Rename existing function
63 - if (fn.node.id == null) {
64 - CompilerError.invariant(
65 - "FunctionDeclaration must have a name",
66 - fn.node.loc ?? GeneratedSource
67 - );
68 - }
69 - const original = fn.node.id;
71 fn.node.id = addSuffix(fn.node.id, "_uncompiled");
72
73 // Rename and append compiled function
@@ -85,12 +86,42 @@ export default function ReactForgetBabelPlugin(
86 buildGatingTest({
87 originalFnDecl: fn,
88 compiledIdent: compiled.id,
88 - originalIdent: original,
89 + originalIdent,
90 gating: pass.opts.gating,
91 })
92 );
93 + if (pass.opts.instrumentForget != null) {
94 + const gatingIdentifierName =
95 + pass.opts.instrumentForget.gating.importSpecifierName;
96 + const instrumentFnName =
97 + pass.opts.instrumentForget.instrumentFn.importSpecifierName;
98 + addInstrumentForget(
99 + fn,
100 + originalIdent.name,
101 + gatingIdentifierName,
102 + instrumentFnName
103 + );
104 + addInstrumentForget(
105 + compiledFn,
106 + originalIdent.name,
107 + gatingIdentifierName,
108 + instrumentFnName
109 + );
110 + }
111 } else {
112 fn.replaceWith(compiled);
113 + if (pass.opts.instrumentForget != null) {
114 + const gatingIdentifierName =
115 + pass.opts.instrumentForget.gating.importSpecifierName;
116 + const instrumentFnName =
117 + pass.opts.instrumentForget.instrumentFn.importSpecifierName;
118 + addInstrumentForget(
119 + fn,
120 + originalIdent.name,
121 + gatingIdentifierName,
122 + instrumentFnName
123 + );
124 + }
125 }
126
127 hasForgetCompiledCode = true;
@@ -302,10 +333,23 @@ export default function ReactForgetBabelPlugin(
333 );
334 }
335 }
336 + // TODO: check for duplicate import specifiers
337 if (options.gating != null) {
338 path.unshiftContainer(
339 "body",
308 - buildImportForGatingModule(options.gating)
340 + buildImportForExternalFunction(options.gating)
341 + );
342 + }
343 + if (options.instrumentForget != null) {
344 + path.unshiftContainer(
345 + "body",
346 + buildImportForExternalFunction(options.instrumentForget.gating)
347 + );
348 + path.unshiftContainer(
349 + "body",
350 + buildImportForExternalFunction(
351 + options.instrumentForget.instrumentFn
352 + )
353 );
354 }
355 }
@@ -430,7 +474,7 @@ type GatingTestOptions = {
474 originalFnDecl: BabelCore.NodePath<t.FunctionDeclaration>;
475 compiledIdent: t.Identifier;
476 originalIdent: t.Identifier;
433 - gating: GatingOptions;
477 + gating: ExternalFunction;
478 };
479 function buildGatingTest({
480 originalFnDecl,
@@ -473,8 +517,8 @@ function addSuffix(id: t.Identifier, suffix: string): t.Identifier {
517 return t.identifier(`${id.name}${suffix}`);
518 }
519
476 -function buildImportForGatingModule(
477 - gating: GatingOptions
520 +function buildImportForExternalFunction(
521 + gating: ExternalFunction
522 ): t.ImportDeclaration {
523 const specifierIdent = buildSpecifierIdent(gating);
524 return t.importDeclaration(
@@ -483,7 +527,7 @@ function buildImportForGatingModule(
527 );
528 }
529
486 -function buildSpecifierIdent(gating: GatingOptions): t.Identifier {
530 +function buildSpecifierIdent(gating: ExternalFunction): t.Identifier {
531 return t.identifier(gating.importSpecifierName);
532 }
533
compiler/forget/src/Babel/InstrumentForgetBabelPlugin.ts new
+34
@@ -0,0 +1,34 @@
1 +/**
2 + * Copyright (c) Meta Platforms, Inc. and affiliates.
3 + *
4 + * This source code is licensed under the MIT license found in the
5 + * LICENSE file in the root directory of this source tree.
6 + */
7 +
8 +import { NodePath } from "@babel/core";
9 +import * as t from "@babel/types";
10 +import "./BabelPlugin";
11 +
12 +export function addInstrumentForget(
13 + fn: NodePath<t.FunctionDeclaration>,
14 + fnName: string,
15 + gatingIdentifierName: string,
16 + instrumentFnName: string
17 +): void {
18 + const fnBody = fn.get("body");
19 + // Technically, this is a conditional hook call. However, we expect
20 + // __DEV__ and gatingIdentifier to be runtime constants
21 + const testExpr: t.Node = t.ifStatement(
22 + t.logicalExpression(
23 + "&&",
24 + t.identifier("__DEV__"),
25 + t.identifier(gatingIdentifierName)
26 + ),
27 + t.expressionStatement(
28 + t.callExpression(t.identifier(instrumentFnName), [
29 + t.stringLiteral(fnName),
30 + ])
31 + )
32 + );
33 + fnBody.unshiftContainer("body", testExpr);
34 +}
compiler/forget/src/Babel/PluginOptions.ts
+36 -2
@@ -7,7 +7,7 @@
7
8 import { EnvironmentConfig } from "../HIR/Environment";
9
10 -export type GatingOptions = {
10 +export type ExternalFunction = {
11 /**
12 * Source for the imported module that exports the `importSpecifierName` functions
13 */
@@ -18,6 +18,11 @@ export type GatingOptions = {
18 importSpecifierName: string;
19 };
20
21 +export type InstrumentForgetOptions = {
22 + gating: ExternalFunction;
23 + instrumentFn: ExternalFunction;
24 +};
25 +
26 export type PluginOptions = {
27 /**
28 * Enable to make Forget only compile functions containing the 'use forget' directive.
@@ -48,7 +53,35 @@ export type PluginOptions = {
53 *
54 * var Foo = isForgetEnabled_Pokes() ? Foo_forget : Foo_uncompiled;
55 */
51 - gating: GatingOptions | null;
56 + gating: ExternalFunction | null;
57 + /**
58 + * Enables instrumentation codegen. This emits a dev-mode only + gated call to
59 + * an instrumentation function, for components and hooks that Forget compiles.
60 + * For example:
61 + * instrumentForget: {
62 + * gating: {
63 + * source: 'ReactInstrumentForgetFeatureFlag',
64 + * importSpecifierName: 'isInstrumentForgetEnabled_Pokes',
65 + * },
66 + * instrumentFn: {
67 + * source: 'react-forget-runtime',
68 + * importSpecifierName: 'useRenderCounter',
69 + * }
70 + * }
71 + *
72 + * produces:
73 + * import {isInstrumentForgetEnabled_Pokes} from 'ReactInstrumentForgetFeatureFlag';
74 + * import {useRenderCounter} from 'react-forget-runtime';
75 + *
76 + * function Component(props) {
77 + * if (__DEV__ && isInstrumentForgetEnabled_Pokes) {
78 + * useRenderCounter();
79 + * }
80 + * // ...
81 + * }
82 + *
83 + */
84 + instrumentForget: InstrumentForgetOptions | null;
85
86 panicOnBailout: boolean;
87
@@ -66,6 +99,7 @@ export const defaultOptions: PluginOptions = {
99 logger: null,
100 gating: null,
101 isDev: false,
102 + instrumentForget: null,
103 } as const;
104
105 export function parsePluginOptions(obj: unknown): PluginOptions {
compiler/forget/src/__tests__/compiler-test.ts
+1
@@ -62,6 +62,7 @@ describe("React Forget", () => {
62 },
63 logger: null,
64 gating: options.gating,
65 + instrumentForget: options.instrumentForget,
66 panicOnBailout: options.panicOnBailout,
67 isDev: true,
68 }).code,
compiler/forget/src/__tests__/fixtures/compiler/instrument-forget-gating-test.expect.md new
+78
@@ -0,0 +1,78 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @instrumentForget @forgetDirective @gating
6 +
7 +function Bar(props) {
8 + "use forget";
9 + return <div>{props.bar}</div>;
10 +}
11 +
12 +function NoForget(props) {
13 + return <Bar>{props.noForget}</Bar>;
14 +}
15 +
16 +function Foo(props) {
17 + "use forget";
18 + return <Foo>{props.bar}</Foo>;
19 +}
20 +
21 +```
22 +
23 +## Code
24 +
25 +```javascript
26 +import { useRenderCounter } from "react-forget-runtime";
27 +import { isInstrumentForgetEnabled_Fixtures } from "ReactInstrumentForgetFeatureFlag";
28 +import { isForgetEnabled_Fixtures } from "ReactForgetFeatureFlag";
29 +import { unstable_useMemoCache as useMemoCache } from "react"; // @instrumentForget @forgetDirective @gating
30 +
31 +function Bar_uncompiled(props) {
32 + "use forget";
33 + if (__DEV__ && isInstrumentForgetEnabled_Fixtures) useRenderCounter("Bar");
34 + return <div>{props.bar}</div>;
35 +}
36 +function Bar_forget(props) {
37 + if (__DEV__ && isInstrumentForgetEnabled_Fixtures) useRenderCounter("Bar");
38 + const $ = useMemoCache(2);
39 + const c_0 = $[0] !== props.bar;
40 + let t0;
41 + if (c_0) {
42 + t0 = <div>{props.bar}</div>;
43 + $[0] = props.bar;
44 + $[1] = t0;
45 + } else {
46 + t0 = $[1];
47 + }
48 + return t0;
49 +}
50 +const Bar = isForgetEnabled_Fixtures() ? Bar_forget : Bar_uncompiled;
51 +
52 +function NoForget(props) {
53 + return <Bar>{props.noForget}</Bar>;
54 +}
55 +
56 +function Foo_uncompiled(props) {
57 + "use forget";
58 + if (__DEV__ && isInstrumentForgetEnabled_Fixtures) useRenderCounter("Foo");
59 + return <Foo>{props.bar}</Foo>;
60 +}
61 +function Foo_forget(props) {
62 + if (__DEV__ && isInstrumentForgetEnabled_Fixtures) useRenderCounter("Foo");
63 + const $ = useMemoCache(2);
64 + const c_0 = $[0] !== props.bar;
65 + let t0;
66 + if (c_0) {
67 + t0 = <Foo>{props.bar}</Foo>;
68 + $[0] = props.bar;
69 + $[1] = t0;
70 + } else {
71 + t0 = $[1];
72 + }
73 + return t0;
74 +}
75 +const Foo = isForgetEnabled_Fixtures() ? Foo_forget : Foo_uncompiled;
76 +
77 +```
78 +
\ No newline at end of file
compiler/forget/src/__tests__/fixtures/compiler/instrument-forget-gating-test.js new
+15
@@ -0,0 +1,15 @@
1 +// @instrumentForget @forgetDirective @gating
2 +
3 +function Bar(props) {
4 + "use forget";
5 + return <div>{props.bar}</div>;
6 +}
7 +
8 +function NoForget(props) {
9 + return <Bar>{props.noForget}</Bar>;
10 +}
11 +
12 +function Foo(props) {
13 + "use forget";
14 + return <Foo>{props.bar}</Foo>;
15 +}
compiler/forget/src/__tests__/fixtures/compiler/instrument-forget-test.expect.md new
+65
@@ -0,0 +1,65 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @instrumentForget @forgetDirective
6 +
7 +function Bar(props) {
8 + "use forget";
9 + return <div>{props.bar}</div>;
10 +}
11 +
12 +function NoForget(props) {
13 + return <Bar>{props.noForget}</Bar>;
14 +}
15 +
16 +function Foo(props) {
17 + "use forget";
18 + return <Foo>{props.bar}</Foo>;
19 +}
20 +
21 +```
22 +
23 +## Code
24 +
25 +```javascript
26 +import { useRenderCounter } from "react-forget-runtime";
27 +import { isInstrumentForgetEnabled_Fixtures } from "ReactInstrumentForgetFeatureFlag";
28 +import { unstable_useMemoCache as useMemoCache } from "react"; // @instrumentForget @forgetDirective
29 +
30 +function Bar(props) {
31 + if (__DEV__ && isInstrumentForgetEnabled_Fixtures) useRenderCounter("Bar");
32 + const $ = useMemoCache(2);
33 + const c_0 = $[0] !== props.bar;
34 + let t0;
35 + if (c_0) {
36 + t0 = <div>{props.bar}</div>;
37 + $[0] = props.bar;
38 + $[1] = t0;
39 + } else {
40 + t0 = $[1];
41 + }
42 + return t0;
43 +}
44 +
45 +function NoForget(props) {
46 + return <Bar>{props.noForget}</Bar>;
47 +}
48 +
49 +function Foo(props) {
50 + if (__DEV__ && isInstrumentForgetEnabled_Fixtures) useRenderCounter("Foo");
51 + const $ = useMemoCache(2);
52 + const c_0 = $[0] !== props.bar;
53 + let t0;
54 + if (c_0) {
55 + t0 = <Foo>{props.bar}</Foo>;
56 + $[0] = props.bar;
57 + $[1] = t0;
58 + } else {
59 + t0 = $[1];
60 + }
61 + return t0;
62 +}
63 +
64 +```
65 +
\ No newline at end of file
compiler/forget/src/__tests__/fixtures/compiler/instrument-forget-test.js new
+15
@@ -0,0 +1,15 @@
1 +// @instrumentForget @forgetDirective
2 +
3 +function Bar(props) {
4 + "use forget";
5 + return <div>{props.bar}</div>;
6 +}
7 +
8 +function NoForget(props) {
9 + return <Bar>{props.noForget}</Bar>;
10 +}
11 +
12 +function Foo(props) {
13 + "use forget";
14 + return <Foo>{props.bar}</Foo>;
15 +}
compiler/forget/src/__tests__/test-utils/generateTestsFromFixtures.ts
+20 -2
@@ -12,7 +12,11 @@
12 import fs from "fs";
13 import glob from "glob";
14 import path from "path";
15 -import { GatingOptions, PluginOptions } from "../../Babel/PluginOptions";
15 +import {
16 + ExternalFunction,
17 + InstrumentForgetOptions,
18 + PluginOptions,
19 +} from "../../Babel/PluginOptions";
20
21 const EXPECT_SUFFIX = ".expect.md";
22
@@ -93,7 +97,8 @@ export default function generateTestsFromFixtures(
97 let input: string | null = null;
98 let debug = false;
99 let enableOnlyOnUseForgetDirective = false;
96 - let gating: GatingOptions | null = null;
100 + let gating: ExternalFunction | null = null;
101 + let instrumentForget: InstrumentForgetOptions | null = null;
102 let inlineUseMemo = true;
103 let panicOnBailout = true;
104 let memoizeJsxElements = true;
@@ -119,6 +124,18 @@ export default function generateTestsFromFixtures(
124 importSpecifierName: "isForgetEnabled_Fixtures",
125 };
126 }
127 + if (lines[0]!.indexOf("@instrumentForget") !== -1) {
128 + instrumentForget = {
129 + gating: {
130 + source: "ReactInstrumentForgetFeatureFlag",
131 + importSpecifierName: "isInstrumentForgetEnabled_Fixtures",
132 + },
133 + instrumentFn: {
134 + source: "react-forget-runtime",
135 + importSpecifierName: "useRenderCounter",
136 + },
137 + };
138 + }
139 if (lines[0]!.indexOf("@inlineUseMemo false") !== -1) {
140 inlineUseMemo = false;
141 }
@@ -139,6 +156,7 @@ export default function generateTestsFromFixtures(
156 debug,
157 enableOnlyOnUseForgetDirective,
158 gating,
159 + instrumentForget,
160 language: parseLanguage(input),
161 panicOnBailout,
162 isDev: true,