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.

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

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) {

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(" ");
 }