@samitouri / QOS-React-2 / commits / 329809de81

Dramatically simplify InlineUseMemo

I realized a wayyyy simpler approach to inlining a lambda: wrap it in a labeled block. The transformation is roughly as follows: ```javascript // Before const x = useMemo(() => { if (a) { return b; } return c; }, [a, b, c]); return x; // After let x; label: { if (a) { x = b; break label; } x = c; break label; } return x; ``` The key to making this work is fixing up some edge cases in labeled blocks, hence the previous PRs.

Joe Savona committed May 17, 2023 at 15:55 UTC 329809de818366cfccf0a6a69e00191b2639efeb

5 files changed +65 -179

compiler/forget/src/Inference/InlineUseMemo.ts

+8 -133

@@ -13,20 +13,19 @@ import {

13 Environment,

14 FunctionExpression,

15 GeneratedSource,

16 - GotoTerminal,

16 GotoVariant,

18 - HIR,

17 HIRFunction,

18 Identifier,

19 IdentifierId,

20 InstructionKind,

21 + LabelTerminal,

22 Place,

23 makeInstructionId,

24 makeType,

25 reversePostorderBlocks,

26 } from "../HIR";

27 import { markInstructionIds, markPredecessors } from "../HIR/HIRBuilder";

29 -import { assertExhaustive, retainWhere } from "../Utils/utils";

28 +import { retainWhere } from "../Utils/utils";

30 /**

31 * Rewrites `useMemo()` calls, rewriting so that the lambda body becomes part of the

@@ -133,70 +132,18 @@ export function inlineUseMemo(fn: HIRFunction): void {

132 // the useMemo

133 block.instructions.length = ii;

134

136 - // The block leading up to the useMemo needs to jump to the entry block of

137 - // the useMemo control flow graph. These will be merged into a single block

138 - // via MergeConsectuveBlocks

139 - const newTerminal: GotoTerminal = {

135 + // To account for complex control flow within the lambda, we treat the lambda

136 + // as if it were a single labeled statement, and replace all returns with gotos

137 + // to the label fallthrough.

138 + const newTerminal: LabelTerminal = {

139 block: body.loweredFunc.body.entry,

140 id: makeInstructionId(0),

142 - kind: "goto",

143 - variant: GotoVariant.Break,

141 + kind: "label",

142 + fallthrough: continuationBlockId,

143 loc: block.terminal.loc,

144 };

145 block.terminal = newTerminal;

146

148 - // If the final terminal type has a fallthrough, update it to point to the

149 - // continuation block

150 - const terminalBlock = getTerminalBlock(

151 - body.loweredFunc.body,

152 - body.loweredFunc.body.entry

153 - );

154 - switch (terminalBlock.terminal.kind) {

155 - case "if":

156 - case "switch":

157 - case "label": {

158 - // These terminals can all appear as the final top-level terminal

159 - // *and* have fallthroughs. If they are final, their fallthrough

160 - // must be updated to point to the continuation block to main

161 - // proper CFG structure (a block that succeeds all branches of a conditional

162 - // must be marked as that conditional's fallthrough)

163 - terminalBlock.terminal.fallthrough = continuationBlockId;

164 - break;

165 - }

166 - case "return":

167 - case "throw": {

168 - // These can appear as the final top-level terminal

169 - break;

170 - }

171 - // These all have non-nullable fallthroughs: there is always some code in the

172 - // CFG that succeeds them which we should find instead

173 - case "optional":

174 - case "ternary":

175 - case "logical":

176 - case "while":

177 - case "for":

178 - case "for-of":

179 - case "do-while":

180 - // These are invalid terminals for a top-level block

181 - case "branch":

182 - case "goto":

183 - case "unsupported": {

184 - CompilerError.invariant(

185 - `Unexpected final top-level terminal`,

186 - terminalBlock.terminal.loc,

187 - `Found ${terminalBlock.terminal.kind}, expected one of if, switch, label, return, or throw`

188 - );

189 - }

190 - default: {

191 - assertExhaustive(

192 - terminalBlock.terminal,

193 - `Unexpected terminal kind '${

194 - (terminalBlock.terminal as any).kind

195 - }'`

196 - );

197 - }

198 - }

199 -

147 // We store the result in the useMemo temporary

148 const result = instr.lvalue;

149

@@ -241,78 +188,6 @@ export function inlineUseMemo(fn: HIRFunction): void {

188 }

189 }

190

244 -// Finds the final top-level terminal node for a CFG, by following any

245 -// fallthrough nodes.

246 -function getTerminalBlock(cfg: HIR, start: BlockId): BasicBlock {

247 - let current = cfg.blocks.get(start)!;

248 - while (true) {

249 - const { terminal } = current;

250 - switch (terminal.kind) {

251 - case "if": {

252 - if (

253 - terminal.fallthrough !== null &&

254 - terminal.fallthrough === terminal.alternate

255 - ) {

256 - // Here we don't know if the fallthrough and alternate are the same because there was

257 - // no alternate or because both the alternate exists and the fallthrough is just unreachable

258 - // So we check if the fallthrough returns/throws (the if is the final top-level terminal)

259 - // or whether execution actually may continue.

260 - const fallthrough = getTerminalBlock(cfg, terminal.fallthrough);

261 - if (

262 - fallthrough.terminal.kind === "return" ||

263 - fallthrough.terminal.kind === "throw"

264 - ) {

265 - return current;

266 - } else {

267 - current = fallthrough;

268 - continue;

269 - }

270 - } else {

271 - return current;

272 - }

273 - }

274 - case "switch":

275 - case "label": {

276 - if (terminal.fallthrough !== null) {

277 - current = cfg.blocks.get(terminal.fallthrough)!;

278 - continue;

279 - } else {

280 - return current;

281 - }

282 - }

283 - case "optional":

284 - case "ternary":

285 - case "logical":

286 - case "while":

287 - case "for":

288 - case "for-of":

289 - case "do-while": {

290 - current = cfg.blocks.get(terminal.fallthrough)!;

291 - continue;

292 - }

293 - case "return":

294 - case "throw": {

295 - return current;

296 - }

297 - case "unsupported":

298 - case "branch":

299 - case "goto": {

300 - CompilerError.invariant(

301 - `Unexpected block terminal`,

302 - terminal.loc,

303 - `Top-level blocks may not end in a ${terminal.kind} terminal`

304 - );

305 - }

306 - default: {

307 - assertExhaustive(

308 - terminal,

309 - `Unexpected terminal kind '${(terminal as any).kind}'`

310 - );

311 - }

312 - }

313 - }

314 -}

315 -

191 /**

192 * Rewrites the block so that all `return` terminals are replaced:

193 * * Add a StoreLocal <returnValue> = <terminal.value>

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

+13 -11

@@ -21,18 +21,20 @@ import { unstable_useMemoCache as useMemoCache } from "react";

21 function Component(props) {

22 const $ = useMemoCache(4);

23 let t21 = undefined;

24 - if (props.cond) {

25 - const c_0 = $[0] !== props.a;

26 - let t0;

27 - if (c_0) {

28 - t0 = makeObject(props.a);

29 - $[0] = props.a;

30 - $[1] = t0;

31 - } else {

32 - t0 = $[1];

24 + bb7: {

25 + if (props.cond) {

26 + const c_0 = $[0] !== props.a;

27 + let t0;

28 + if (c_0) {

29 + t0 = makeObject(props.a);

30 + $[0] = props.a;

31 + $[1] = t0;

32 + } else {

33 + t0 = $[1];

34 + }

35 + t21 = t0;

36 + break bb7;

37 }

34 - t21 = t0;

35 - } else {

38 const c_2 = $[2] !== props.b;

39 let t1;

40 if (c_2) {

compiler/forget/src/__tests__/fixtures/compiler/useMemo-inlining-block-return.expect.md

+13 -11

@@ -20,18 +20,20 @@ import { unstable_useMemoCache as useMemoCache } from "react";

20 function component(a, b) {

21 const $ = useMemoCache(2);

22 let t14 = undefined;

23 - if (a) {

24 - const c_0 = $[0] !== b;

25 - let t0;

26 - if (c_0) {

27 - t0 = { b };

28 - $[0] = b;

29 - $[1] = t0;

30 - } else {

31 - t0 = $[1];

23 + bb6: {

24 + if (a) {

25 + const c_0 = $[0] !== b;

26 + let t0;

27 + if (c_0) {

28 + t0 = { b };

29 + $[0] = b;

30 + $[1] = t0;

31 + } else {

32 + t0 = $[1];

33 + }

34 + t14 = t0;

35 + break bb6;

36 }

33 - t14 = t0;

34 - } else {

37 t14 = undefined;

38 }

39 const x = t14;

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

+22 -17

@@ -24,25 +24,30 @@ function Component(props) {

24 ```javascript

25 import { unstable_useMemoCache as useMemoCache } from "react";

26 function Component(props) {

27 - const $ = useMemoCache(2);

28 - const c_0 = $[0] !== props;

29 - let t32;

30 - if (c_0) {

31 - const y = [];

32 - if (props.cond) {

33 - y.push(props.a);

34 - }

35 - t32 = undefined;

36 - if (props.cond2) {

37 - t32 = y;

38 - } else {

27 + const $ = useMemoCache(3);

28 + let t32 = undefined;

29 + bb9: {

30 + const c_0 = $[0] !== props;

31 + let y;

32 + if (c_0) {

33 + y = [];

34 + if (props.cond) {

35 + y.push(props.a);

36 + }

37 + if (props.cond2) {

38 + t32 = y;

39 + break bb9;

40 + }

41 +

42 y.push(props.b);

40 - t32 = y;

43 + $[0] = props;

44 + $[1] = y;

45 + $[2] = t32;

46 + } else {

47 + y = $[1];

48 + t32 = $[2];

49 }

42 - $[0] = props;

43 - $[1] = t32;

44 - } else {

45 - t32 = $[1];

50 + t32 = y;

51 }

52 const x = t32;

53 return x;

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

+9 -7

@@ -23,13 +23,15 @@ function Component(props) {

23 ```javascript

24 function Component(props) {

25 let t18 = undefined;

26 - bb8: switch (props.key) {

27 - case "key": {

28 - t18 = props.value;

29 - break bb8;

30 - }

31 - default: {

32 - t18 = props.defaultValue;

26 + bb8: {

27 + switch (props.key) {

28 + case "key": {

29 + t18 = props.value;

30 + break bb8;

31 + }

32 + default: {

33 + t18 = props.defaultValue;

34 + }

35 }

36 }

37 const x = t18;