[hir] Add a DropMemoCall pass

This drops the memo hook calls from the IR

compiler/forget/src/CompilerPipeline.ts

@@ -14,6 +14,7 @@ import {
 } from "./HIR";
 import {
   analyseFunctions,
+  dropMemoCalls,
   inferMutableRanges,
   inferReferenceEffects,
 } from "./Inference";
@@ -65,6 +66,9 @@ export function* run(
   inferTypes(hir);
   yield log({ kind: "hir", name: "InferTypes", value: hir });
 
+  dropMemoCalls(hir);
+  yield log({ kind: "hir", name: "DropMemoCalls", value: hir });
+
   analyseFunctions(hir);
   yield log({ kind: "hir", name: "AnalyseFunctions", value: hir });
 

compiler/forget/src/HIR/HIR.ts

@@ -745,3 +745,7 @@ export function isObjectType(id: Identifier): boolean {
 export function isPrimitiveType(id: Identifier): boolean {
   return id.type.kind === "Primitive";
 }
+
+export function isHookType(id: Identifier): boolean {
+  return id.type.kind === "Hook";
+}

compiler/forget/src/Inference/DropMemoCalls.ts

@@ -0,0 +1,60 @@
+import invariant from "invariant";
+import {
+  Effect,
+  HIRFunction,
+  HookType,
+  InstructionValue,
+  isHookType,
+} from "../HIR";
+
+export default function (func: HIRFunction) {
+  for (const [_, block] of func.body.blocks) {
+    for (const instr of block.instructions) {
+      switch (instr.value.kind) {
+        case "CallExpression": {
+          if (isHookType(instr.value.callee.identifier)) {
+            const name = (instr.value.callee.identifier.type as HookType).name;
+            if (name === "useMemo") {
+              const [fn] = instr.value.args;
+
+              // TODO(gsn): Consider inlining the function passed to useMemo,
+              // rather than just calling it directly.
+              //
+              // Replace the hook callee with the fn arg.
+              //
+              // before:
+              //   foo = Call useMemo$2($9, $10)
+              //
+              // after:
+              //   foo = Call $9()
+              instr.value = {
+                kind: "CallExpression",
+                callee: fn,
+                // Drop the args, including the deps array which DCE will remove
+                // later.
+                args: [],
+                loc: instr.value.loc,
+              };
+            } else if (name === "useCallback") {
+              const [fn] = instr.value.args;
+
+              // Instead of a Call, just alias the callback directly.
+              //
+              // before:
+              //   foo = Call useCallback$8($19)
+              //
+              // after:
+              //   foo = $19
+              instr.value = {
+                kind: "Identifier",
+                identifier: fn.identifier,
+                effect: Effect.Unknown,
+                loc: instr.value.loc,
+              };
+            }
+          }
+        }
+      }
+    }
+  }
+}

compiler/forget/src/Inference/index.ts

@@ -5,6 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+export { default as dropMemoCalls } from "./DropMemoCalls";
 export { inferMutableRanges } from "./InferMutableRanges";
 export { default as analyseFunctions } from "./AnalyseFunctions";
 export { default as inferReferenceEffects } from "./InferReferenceEffects";

compiler/forget/src/__tests__/fixtures/hir/use-callback-simple.expect.md

@@ -0,0 +1,47 @@
+
+## Input
+
+```javascript
+function component() {
+  const [count, setCount] = useState(0);
+  const increment = useCallback(() => setCount(count + 1));
+
+  return <Foo onClick={increment}></Foo>;
+}
+
+```
+
+## Code
+
+```javascript
+function component() {
+  const $ = React.unstable_useMemoCache();
+  const t2 = useState(0);
+  const count = t2[0];
+  const setCount = t2[1];
+  const c_0 = $[0] !== setCount;
+  const c_1 = $[1] !== count;
+  let t0;
+  if (c_0 || c_1) {
+    t0 = () => setCount(count + 1);
+    $[0] = setCount;
+    $[1] = count;
+    $[2] = t0;
+  } else {
+    t0 = $[2];
+  }
+  const increment = t0;
+  const c_3 = $[3] !== increment;
+  let t1;
+  if (c_3) {
+    t1 = <Foo onClick={increment}></Foo>;
+    $[3] = increment;
+    $[4] = t1;
+  } else {
+    t1 = $[4];
+  }
+  return t1;
+}
+
+```
+      

compiler/forget/src/__tests__/fixtures/hir/use-callback-simple.js

@@ -0,0 +1,6 @@
+function component() {
+  const [count, setCount] = useState(0);
+  const increment = useCallback(() => setCount(count + 1));
+
+  return <Foo onClick={increment}></Foo>;
+}

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

@@ -0,0 +1,39 @@
+
+## Input
+
+```javascript
+function component(a) {
+  let x = useMemo(() => [a], [a]);
+  return <Foo x={x}></Foo>;
+}
+
+```
+
+## Code
+
+```javascript
+function component(a) {
+  const $ = React.unstable_useMemoCache();
+  const c_0 = $[0] !== a;
+  let x;
+  if (c_0) {
+    x = (() => [a])();
+    $[0] = a;
+    $[1] = x;
+  } else {
+    x = $[1];
+  }
+  const c_2 = $[2] !== x;
+  let t0;
+  if (c_2) {
+    t0 = <Foo x={x}></Foo>;
+    $[2] = x;
+    $[3] = t0;
+  } else {
+    t0 = $[3];
+  }
+  return t0;
+}
+
+```
+      

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

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