[hir] Put useMemo inlining behind a flag

It's still a WIP so disable behind a flag for now.

compiler/forget/src/CompilerPipeline.ts

@@ -61,8 +61,10 @@ export function* run(
   const hir = lower(func, env).unwrap();
   yield log({ kind: "hir", name: "HIR", value: hir });
 
-  inlineUseMemo(hir);
-  yield log({ kind: "hir", name: "RewriteUseMemo", value: hir });
+  if (config?.inlineUseMemo) {
+    inlineUseMemo(hir);
+    yield log({ kind: "hir", name: "RewriteUseMemo", value: hir });
+  }
 
   mergeConsecutiveBlocks(hir);
   yield log({ kind: "hir", name: "MergeConsecutiveBlocks", value: hir });

compiler/forget/src/HIR/Environment.ts

@@ -40,6 +40,7 @@ const HOOK_PATTERN = /^_?use/;
 export type EnvironmentConfig = Partial<{
   customHooks: Map<string, Hook>;
   memoizeJsxElements: boolean;
+  inlineUseMemo: boolean;
 }>;
 
 export class Environment {

compiler/forget/src/__tests__/compiler-test.ts

@@ -54,6 +54,7 @@ describe("React Forget", () => {
                   },
                 ],
               ]),
+              inlineUseMemo: options.environment?.inlineUseMemo ?? false,
             },
             logger: null,
             gating: options.gating,

compiler/forget/src/__tests__/disableMemoizeJsxElements-test.ts

@@ -50,6 +50,7 @@ describe("React Forget (Disable memoization of JSX elements)", () => {
                   },
                 ],
               ]),
+              inlineUseMemo: options.environment?.inlineUseMemo ?? false,
             },
             logger: null,
             gating: options.gating,

compiler/forget/src/__tests__/fixtures/compiler/useMemo-if-else-multiple-return.expect.md

@@ -2,6 +2,7 @@
 ## Input
 
 ```javascript
+// @inlineUseMemo
 function Component(props) {
   const x = useMemo(() => {
     if (props.cond) {
@@ -17,6 +18,7 @@ function Component(props) {
 ## Code
 
 ```javascript
+// @inlineUseMemo
 function Component(props) {
   const $ = React.unstable_useMemoCache(5);
   if (props.cond) {

compiler/forget/src/__tests__/fixtures/compiler/useMemo-if-else-multiple-return.js

@@ -1,3 +1,4 @@
+// @inlineUseMemo
 function Component(props) {
   const x = useMemo(() => {
     if (props.cond) {

compiler/forget/src/__tests__/fixtures/compiler/useMemo-independently-memoizeable.expect.md

@@ -2,6 +2,7 @@
 ## Input
 
 ```javascript
+// @inlineUseMemo
 function Component(props) {
   const [a, b] = useMemo(() => {
     const items = [];
@@ -17,6 +18,7 @@ function Component(props) {
 ## Code
 
 ```javascript
+// @inlineUseMemo
 function Component(props) {
   const $ = React.unstable_useMemoCache(10);
   const c_0 = $[0] !== props.a;

compiler/forget/src/__tests__/fixtures/compiler/useMemo-independently-memoizeable.js

@@ -1,3 +1,4 @@
+// @inlineUseMemo
 function Component(props) {
   const [a, b] = useMemo(() => {
     const items = [];

compiler/forget/src/__tests__/fixtures/compiler/useMemo-labeled-statement-unconditional-return.expect.md

@@ -2,6 +2,7 @@
 ## Input
 
 ```javascript
+// @inlineUseMemo
 function Component(props) {
   const x = useMemo(() => {
     label: {
@@ -16,6 +17,7 @@ function Component(props) {
 ## Code
 
 ```javascript
+// @inlineUseMemo
 function Component(props) {
   const t19 = props.value;
   const x = t19;

compiler/forget/src/__tests__/fixtures/compiler/useMemo-labeled-statement-unconditional-return.js

@@ -1,3 +1,4 @@
+// @inlineUseMemo
 function Component(props) {
   const x = useMemo(() => {
     label: {

compiler/forget/src/__tests__/fixtures/compiler/useMemo-logical.expect.md

@@ -2,6 +2,7 @@
 ## Input
 
 ```javascript
+// @inlineUseMemo
 function Component(props) {
   const x = useMemo(() => props.a && props.b);
   return x;
@@ -12,6 +13,7 @@ function Component(props) {
 ## Code
 
 ```javascript
+// @inlineUseMemo
 function Component(props) {
   const t32 = props.a && props.b;
   const x = t32;

compiler/forget/src/__tests__/fixtures/compiler/useMemo-logical.js

@@ -1,3 +1,4 @@
+// @inlineUseMemo
 function Component(props) {
   const x = useMemo(() => props.a && props.b);
   return x;

compiler/forget/src/__tests__/fixtures/compiler/useMemo-multiple-if-else.expect.md

@@ -2,6 +2,7 @@
 ## Input
 
 ```javascript
+// @inlineUseMemo
 function Component(props) {
   const x = useMemo(() => {
     let y = [];
@@ -22,6 +23,7 @@ function Component(props) {
 ## Code
 
 ```javascript
+// @inlineUseMemo
 function Component(props) {
   const $ = React.unstable_useMemoCache(2);
   const c_0 = $[0] !== props;

compiler/forget/src/__tests__/fixtures/compiler/useMemo-multiple-if-else.js

@@ -1,3 +1,4 @@
+// @inlineUseMemo
 function Component(props) {
   const x = useMemo(() => {
     let y = [];

compiler/forget/src/__tests__/fixtures/compiler/useMemo-simple.expect.md

@@ -2,6 +2,7 @@
 ## Input
 
 ```javascript
+// @inlineUseMemo
 function component(a) {
   let x = useMemo(() => [a], [a]);
   return <Foo x={x}></Foo>;
@@ -12,6 +13,7 @@ function component(a) {
 ## Code
 
 ```javascript
+// @inlineUseMemo
 function component(a) {
   const $ = React.unstable_useMemoCache(4);
   const c_0 = $[0] !== a;

compiler/forget/src/__tests__/fixtures/compiler/useMemo-simple.js

@@ -1,3 +1,4 @@
+// @inlineUseMemo
 function component(a) {
   let x = useMemo(() => [a], [a]);
   return <Foo x={x}></Foo>;

compiler/forget/src/__tests__/fixtures/compiler/useMemo-switch-no-fallthrough.expect.md

@@ -2,6 +2,7 @@
 ## Input
 
 ```javascript
+// @inlineUseMemo
 function Component(props) {
   const x = useMemo(() => {
     switch (props.key) {
@@ -21,6 +22,7 @@ function Component(props) {
 ## Code
 
 ```javascript
+// @inlineUseMemo
 function Component(props) {
   bb8: switch (props.key) {
     case "key": {

compiler/forget/src/__tests__/fixtures/compiler/useMemo-switch-no-fallthrough.js

@@ -1,3 +1,4 @@
+// @inlineUseMemo
 function Component(props) {
   const x = useMemo(() => {
     switch (props.key) {

compiler/forget/src/__tests__/hir-test.ts

@@ -45,6 +45,7 @@ describe("React Forget (HIR version)", () => {
             },
           ],
         ]),
+        inlineUseMemo: options.environment?.inlineUseMemo ?? false,
       });
 
       if (compileResult.isErr()) {

compiler/forget/src/__tests__/test-utils/generateTestsFromFixtures.ts

@@ -91,6 +91,7 @@ export default function generateTestsFromFixtures(
         let debug = false;
         let enableOnlyOnUseForgetDirective = false;
         let gating: GatingOptions | null = null;
+        let inlineUseMemo = true;
 
         if (inputFile != null) {
           input = fs.readFileSync(inputFile, "utf8");
@@ -111,13 +112,16 @@ export default function generateTestsFromFixtures(
               importSpecifierName: "isForgetEnabled_Fixtures",
             };
           }
+          if (lines[0]!.indexOf("@inlineUseMemo") !== -1) {
+            inlineUseMemo = true;
+          }
         }
 
         testCommand(basename, () => {
           let receivedOutput;
           if (input !== null) {
             receivedOutput = transform(input, basename, {
-              environment: null,
+              environment: { inlineUseMemo },
               logger: null,
               debug,
               enableOnlyOnUseForgetDirective,