From fb8f293c321050817cb1c1fdb201cc33954d089b Mon Sep 17 00:00:00 2001 From: Sathya Gunasekaran Date: Tue, 14 Feb 2023 23:07:19 +0000 Subject: [PATCH] [hir] Add a DropMemoCall pass This drops the memo hook calls from the IR --- compiler/forget/src/CompilerPipeline.ts | 4 ++ compiler/forget/src/HIR/HIR.ts | 4 ++ .../forget/src/Inference/DropMemoCalls.ts | 60 +++++++++++++++++++ compiler/forget/src/Inference/index.ts | 1 + .../hir/use-callback-simple.expect.md | 47 +++++++++++++++ .../fixtures/hir/use-callback-simple.js | 6 ++ .../fixtures/hir/useMemo-simple.expect.md | 39 ++++++++++++ .../__tests__/fixtures/hir/useMemo-simple.js | 4 ++ 8 files changed, 165 insertions(+) create mode 100644 compiler/forget/src/Inference/DropMemoCalls.ts create mode 100644 compiler/forget/src/__tests__/fixtures/hir/use-callback-simple.expect.md create mode 100644 compiler/forget/src/__tests__/fixtures/hir/use-callback-simple.js create mode 100644 compiler/forget/src/__tests__/fixtures/hir/useMemo-simple.expect.md create mode 100644 compiler/forget/src/__tests__/fixtures/hir/useMemo-simple.js 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 ; +}