Codegen for early returns in reactive scopes
Implements codegen for reactive scopes with early returns, though we don't ever construct such a case yet. See comments in the code. There is a slightly more optimal output that would require a larger refactor (also described in code comments), for now i'm starting with the simpler approach since this is relatively rare so we don't need to optimize code size / runtime as much.
Joe Savona committed
Dec 20, 2023 at 13:52 UTC
3175056935ed76ea3790c6e956fedd354587ca5d
3 files changed
+93
-13
compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts
+5
-1
@@ -1080,7 +1080,11 @@ export type ReactiveScope = {
1080
*
1081
* This value is null for scopes that do not contain early returns.
1082
*/
1083
- earlyReturnValue: { value: IdentifierId; loc: SourceLocation } | null;
1083
+ earlyReturnValue: {
1084
+ value: IdentifierId;
1085
+ loc: SourceLocation;
1086
+ label: string;
1087
+ } | null;
1088
1089
/*
1090
* Some passes may merge scopes together. The merged set contains the
compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts
+85
-11
@@ -448,22 +448,96 @@ function codegenReactiveScope(
448
);
449
}
450
451
- const computationBlock = codegenBlock(cx, block);
451
+ let computationBlock = codegenBlock(cx, block);
452
computationBlock.body.push(...cacheStoreStatements);
453
const memoBlock = t.blockStatement(cacheLoadStatements);
454
455
- let memoStatement;
455
if (scope.earlyReturnValue !== null) {
457
- // Has early return
458
- CompilerError.throwTodo({
459
- reason: `Codegen support for reactive scopes with early return`,
460
- loc: scope.earlyReturnValue.loc,
461
- description: null,
462
- suggestions: null,
463
- });
464
- } else {
465
- memoStatement = t.ifStatement(testCondition, computationBlock, memoBlock);
456
+ /**
457
+ * Handle early return. PropagateEarlyReturns should already have
458
+ * converted the actual return statements within the block into
459
+ * the appropriate form, so we just have to add the appropriate
460
+ * wrapping code.
461
+ *
462
+ * Example:
463
+ *
464
+ * ```
465
+ * if (input !== $[0]) {
466
+ * let t0 = Symbol.for('react.memo_cache_sentinel');
467
+ * label: {
468
+ * ... memo block ...
469
+ * if (cond) {
470
+ * // this part is already rewritten by PropagateEarlyReturns
471
+ * t0 = ...; // save the early return value
472
+ * break label;
473
+ * }
474
+ * ... more memo block...
475
+ * }
476
+ * $[1] = t0;
477
+ * if (t0 !== Symbol.for('react.memo_cache_sentinel')) {
478
+ * return t0;
479
+ * }
480
+ * } else {
481
+ * ...
482
+ * const t0 = $[1];
483
+ * if (t0 !== Symbol.for('react.memo_cache_sentinel')) {
484
+ * return t0;
485
+ * }
486
+ * }
487
+ * ```
488
+ *
489
+ * TODO: factor out the common if-return from the if/else branches
490
+ * We can lift the temporary (here `t0`) to the outer block,
491
+ * then move the `if (t0 !== sentinel) { return t0 }` to after the
492
+ * memo block if/else, since both branches need to execute that check.
493
+ */
494
+ const index = cx.nextCacheIndex;
495
+ const identifier = t.identifier(`t${cx.env.nextIdentifierId}`);
496
+ const sentinel = t.callExpression(
497
+ t.memberExpression(t.identifier("Symbol"), t.identifier("for")),
498
+ [t.stringLiteral("react.memo_cache_sentinel")]
499
+ );
500
+ computationBlock = t.blockStatement([
501
+ t.variableDeclaration("let", [
502
+ t.variableDeclarator(identifier, sentinel),
503
+ ]),
504
+ t.labeledStatement(
505
+ t.identifier(scope.earlyReturnValue.label),
506
+ computationBlock
507
+ ),
508
+ t.expressionStatement(
509
+ t.assignmentExpression(
510
+ "=",
511
+ t.memberExpression(t.identifier("$"), t.numericLiteral(index), true),
512
+ identifier
513
+ )
514
+ ),
515
+ t.ifStatement(
516
+ t.binaryExpression("!==", identifier, sentinel),
517
+ t.blockStatement([t.returnStatement(identifier)])
518
+ ),
519
+ ]);
520
+
521
+ memoBlock.body.push(
522
+ ...[
523
+ t.variableDeclaration("const", [
524
+ t.variableDeclarator(
525
+ identifier,
526
+ t.memberExpression(t.identifier("$"), t.numericLiteral(index), true)
527
+ ),
528
+ ]),
529
+ t.ifStatement(
530
+ t.binaryExpression("!==", identifier, sentinel),
531
+ t.blockStatement([t.returnStatement(identifier)])
532
+ ),
533
+ ]
534
+ );
535
}
536
+ const memoStatement = t.ifStatement(
537
+ testCondition,
538
+ computationBlock,
539
+ memoBlock
540
+ );
541
542
if (cx.env.config.enableMemoizationComments) {
543
if (changeExpressionComments.length) {
compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/PrintReactiveFunction.ts
+3
-1
@@ -65,7 +65,9 @@ export function printReactiveScopeSummary(scope: ReactiveScope): string {
65
)}]`
66
);
67
if (scope.earlyReturnValue !== null) {
68
- items.push(`earlyReturn=${scope.earlyReturnValue.value}`);
68
+ items.push(
69
+ `earlyReturn={id: ${scope.earlyReturnValue.value}, label: ${scope.earlyReturnValue.label}}`
70
+ );
71
}
72
return items.join(" ");
73
}