From 3603ee00a65f292d9d08cf36c31a9cf5f8843467 Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Wed, 17 May 2023 15:55:54 -0700
Subject: [PATCH] Retain original structure of labeled blocks

Handles some edge-cases where we previously flattened away some of the structure
of a labeled block, instead ensuring that we retain the original shape. See the
output.

## Test Plan

Tested on the internal app we're focused on (w useMemo inlining enabled), it
works fine.
---
 .../forget/src/HIR/MergeConsecutiveBlocks.ts  |  4 ++++
 .../ReactiveScopes/BuildReactiveFunction.ts   | 22 +++++++++++--------
 .../fixtures/compiler/dominator.expect.md     |  3 ++-
 .../fixtures/compiler/early-return.expect.md  |  2 --
 .../compiler/useMemo-inverted-if.expect.md    | 11 +++++++---
 5 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/compiler/forget/src/HIR/MergeConsecutiveBlocks.ts b/compiler/forget/src/HIR/MergeConsecutiveBlocks.ts
index dc1351eac7..4ed4487e06 100644
--- a/compiler/forget/src/HIR/MergeConsecutiveBlocks.ts
+++ b/compiler/forget/src/HIR/MergeConsecutiveBlocks.ts
@@ -14,6 +14,7 @@ import {
   Instruction,
 } from "./HIR";
 import { markPredecessors, removeUnreachableFallthroughs } from "./HIRBuilder";
+import { mapOptionalFallthroughs } from "./visitors";
 
 /**
  * Merges sequences of blocks that will always execute consecutively —
@@ -86,6 +87,9 @@ export function mergeConsecutiveBlocks(fn: HIRFunction): void {
     fn.body.blocks.delete(block.id);
   }
   markPredecessors(fn.body);
+  for (const [, block] of fn.body.blocks) {
+    mapOptionalFallthroughs(block.terminal, (blockId) => merged.get(blockId));
+  }
   removeUnreachableFallthroughs(fn.body);
 }
 
diff --git a/compiler/forget/src/ReactiveScopes/BuildReactiveFunction.ts b/compiler/forget/src/ReactiveScopes/BuildReactiveFunction.ts
index 5d7b6d2b78..e2aa716365 100644
--- a/compiler/forget/src/ReactiveScopes/BuildReactiveFunction.ts
+++ b/compiler/forget/src/ReactiveScopes/BuildReactiveFunction.ts
@@ -573,15 +573,18 @@ class Driver {
       case "optional":
       case "ternary":
       case "logical": {
-        const fallthroughId = terminal.fallthrough;
-        invariant(
-          !this.cx.isScheduled(fallthroughId),
-          "Logical terminal fallthrough cannot have been scheduled"
-        );
-        const scheduleId = this.cx.schedule(fallthroughId, "if");
-        scheduleIds.push(scheduleId);
+        const fallthroughId =
+          terminal.fallthrough !== null &&
+          !this.cx.isScheduled(terminal.fallthrough)
+            ? terminal.fallthrough
+            : null;
+        if (fallthroughId !== null) {
+          const scheduleId = this.cx.schedule(fallthroughId, "if");
+          scheduleIds.push(scheduleId);
+        }
 
         const { place, value } = this.visitValueBlockTerminal(terminal);
+        this.cx.unscheduleAll(scheduleIds);
         blockValue.push({
           kind: "instruction",
           instruction: {
@@ -592,8 +595,9 @@ class Driver {
           },
         });
 
-        this.cx.unschedule(scheduleId);
-        this.visitBlock(this.cx.ir.blocks.get(fallthroughId)!, blockValue);
+        if (fallthroughId !== null) {
+          this.visitBlock(this.cx.ir.blocks.get(fallthroughId)!, blockValue);
+        }
         break;
       }
       case "goto": {
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/dominator.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/dominator.expect.md
index e9035aa273..3f659b2ea2 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/dominator.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/dominator.expect.md
@@ -47,10 +47,11 @@ function Component(props) {
     x = 1;
   } else {
     if (props.b) {
-      x = 3;
     } else {
       break bb1;
     }
+
+    x = 3;
   }
   bb10: bb12: switch (props.c) {
     case "a": {
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/early-return.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/early-return.expect.md
index fbe27fb900..e6c3acef6e 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/early-return.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/early-return.expect.md
@@ -19,8 +19,6 @@ function MyApp(props) {
 function MyApp(props) {
   if (props.cond) {
     return;
-  } else {
-    return;
   }
 }
 
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-inverted-if.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-inverted-if.expect.md
index ebe949ec3b..a1ffe4a99c 100644
--- a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-inverted-if.expect.md
+++ b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-inverted-if.expect.md
@@ -23,11 +23,16 @@ function Component(props) {
 function Component(props) {
   let t17 = undefined;
   bb10: {
-    if (props.cond) {
-      t17 = props.b;
+    bb5: {
+      if (props.cond) {
+        break bb5;
+      }
+
+      t17 = props.a;
       break bb10;
     }
-    t17 = props.a;
+
+    t17 = props.b;
   }
   const x = t17;
   return x;