diff --git a/compiler/forget/src/CompilerPipeline.ts b/compiler/forget/src/CompilerPipeline.ts
index 9807f6c212..cbd44e26b7 100644
--- a/compiler/forget/src/CompilerPipeline.ts
+++ b/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 });
diff --git a/compiler/forget/src/HIR/HIR.ts b/compiler/forget/src/HIR/HIR.ts
index 9da418c3cf..6126f6c5cc 100644
--- a/compiler/forget/src/HIR/HIR.ts
+++ b/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";
+}
diff --git a/compiler/forget/src/Inference/DropMemoCalls.ts b/compiler/forget/src/Inference/DropMemoCalls.ts
new file mode 100644
index 0000000000..75b8fd5cd5
--- /dev/null
+++ b/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,
+ };
+ }
+ }
+ }
+ }
+ }
+ }
+}
diff --git a/compiler/forget/src/Inference/index.ts b/compiler/forget/src/Inference/index.ts
index 907516d3c9..a77861f352 100644
--- a/compiler/forget/src/Inference/index.ts
+++ b/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";
diff --git a/compiler/forget/src/__tests__/fixtures/hir/use-callback-simple.expect.md b/compiler/forget/src/__tests__/fixtures/hir/use-callback-simple.expect.md
new file mode 100644
index 0000000000..456784061e
--- /dev/null
+++ b/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 ;
+}
+
+```
+
+## 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 = ;
+ $[3] = increment;
+ $[4] = t1;
+ } else {
+ t1 = $[4];
+ }
+ return t1;
+}
+
+```
+
\ No newline at end of file
diff --git a/compiler/forget/src/__tests__/fixtures/hir/use-callback-simple.js b/compiler/forget/src/__tests__/fixtures/hir/use-callback-simple.js
new file mode 100644
index 0000000000..132fcb30a2
--- /dev/null
+++ b/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 ;
+}
diff --git a/compiler/forget/src/__tests__/fixtures/hir/useMemo-simple.expect.md b/compiler/forget/src/__tests__/fixtures/hir/useMemo-simple.expect.md
new file mode 100644
index 0000000000..76ee9099ad
--- /dev/null
+++ b/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 ;
+}
+
+```
+
+## 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 = ;
+ $[2] = x;
+ $[3] = t0;
+ } else {
+ t0 = $[3];
+ }
+ return t0;
+}
+
+```
+
\ No newline at end of file
diff --git a/compiler/forget/src/__tests__/fixtures/hir/useMemo-simple.js b/compiler/forget/src/__tests__/fixtures/hir/useMemo-simple.js
new file mode 100644
index 0000000000..a680d099be
--- /dev/null
+++ b/compiler/forget/src/__tests__/fixtures/hir/useMemo-simple.js
@@ -0,0 +1,4 @@
+function component(a) {
+ let x = useMemo(() => [a], [a]);
+ return ;
+}