diff --git a/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts b/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts index 181e585093..97df903b4d 100644 --- a/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts +++ b/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts @@ -1080,7 +1080,11 @@ export type ReactiveScope = { * * This value is null for scopes that do not contain early returns. */ - earlyReturnValue: { value: IdentifierId; loc: SourceLocation } | null; + earlyReturnValue: { + value: IdentifierId; + loc: SourceLocation; + label: string; + } | null; /* * Some passes may merge scopes together. The merged set contains the diff --git a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts index 8c69ecfdf3..6f9361bd97 100644 --- a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts +++ b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts @@ -448,22 +448,96 @@ function codegenReactiveScope( ); } - const computationBlock = codegenBlock(cx, block); + let computationBlock = codegenBlock(cx, block); computationBlock.body.push(...cacheStoreStatements); const memoBlock = t.blockStatement(cacheLoadStatements); - let memoStatement; if (scope.earlyReturnValue !== null) { - // Has early return - CompilerError.throwTodo({ - reason: `Codegen support for reactive scopes with early return`, - loc: scope.earlyReturnValue.loc, - description: null, - suggestions: null, - }); - } else { - memoStatement = t.ifStatement(testCondition, computationBlock, memoBlock); + /** + * Handle early return. PropagateEarlyReturns should already have + * converted the actual return statements within the block into + * the appropriate form, so we just have to add the appropriate + * wrapping code. + * + * Example: + * + * ``` + * if (input !== $[0]) { + * let t0 = Symbol.for('react.memo_cache_sentinel'); + * label: { + * ... memo block ... + * if (cond) { + * // this part is already rewritten by PropagateEarlyReturns + * t0 = ...; // save the early return value + * break label; + * } + * ... more memo block... + * } + * $[1] = t0; + * if (t0 !== Symbol.for('react.memo_cache_sentinel')) { + * return t0; + * } + * } else { + * ... + * const t0 = $[1]; + * if (t0 !== Symbol.for('react.memo_cache_sentinel')) { + * return t0; + * } + * } + * ``` + * + * TODO: factor out the common if-return from the if/else branches + * We can lift the temporary (here `t0`) to the outer block, + * then move the `if (t0 !== sentinel) { return t0 }` to after the + * memo block if/else, since both branches need to execute that check. + */ + const index = cx.nextCacheIndex; + const identifier = t.identifier(`t${cx.env.nextIdentifierId}`); + const sentinel = t.callExpression( + t.memberExpression(t.identifier("Symbol"), t.identifier("for")), + [t.stringLiteral("react.memo_cache_sentinel")] + ); + computationBlock = t.blockStatement([ + t.variableDeclaration("let", [ + t.variableDeclarator(identifier, sentinel), + ]), + t.labeledStatement( + t.identifier(scope.earlyReturnValue.label), + computationBlock + ), + t.expressionStatement( + t.assignmentExpression( + "=", + t.memberExpression(t.identifier("$"), t.numericLiteral(index), true), + identifier + ) + ), + t.ifStatement( + t.binaryExpression("!==", identifier, sentinel), + t.blockStatement([t.returnStatement(identifier)]) + ), + ]); + + memoBlock.body.push( + ...[ + t.variableDeclaration("const", [ + t.variableDeclarator( + identifier, + t.memberExpression(t.identifier("$"), t.numericLiteral(index), true) + ), + ]), + t.ifStatement( + t.binaryExpression("!==", identifier, sentinel), + t.blockStatement([t.returnStatement(identifier)]) + ), + ] + ); } + const memoStatement = t.ifStatement( + testCondition, + computationBlock, + memoBlock + ); if (cx.env.config.enableMemoizationComments) { if (changeExpressionComments.length) { diff --git a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/PrintReactiveFunction.ts b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/PrintReactiveFunction.ts index 13b8be8267..ef859fc403 100644 --- a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/PrintReactiveFunction.ts +++ b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/PrintReactiveFunction.ts @@ -65,7 +65,9 @@ export function printReactiveScopeSummary(scope: ReactiveScope): string { )}]` ); if (scope.earlyReturnValue !== null) { - items.push(`earlyReturn=${scope.earlyReturnValue.value}`); + items.push( + `earlyReturn={id: ${scope.earlyReturnValue.value}, label: ${scope.earlyReturnValue.label}}` + ); } return items.join(" "); }