[compiler] Add fallthrough to branch terminal

Branch terminals didn't have a fallthrough because they correspond to an outer terminal (optional, logical, etc) that has the "real" fallthrough. But understanding how branch terminals correspond to these outer terminals requires knowing the branch fallthrough. For example, `foo?.bar?.baz` creates terminals along the lines of:

```
bb0:
  optional fallthrough=bb4
bb1:
  optional fallthrough=bb3
bb2:
  ...
  branch ... (fallthrough=bb3)

...

bb3:
  ...
  branch ... (fallthrough=bb4)

...

bb4:
  ...
```

Without a fallthrough on `branch` terminals, it's unclear that the optional from bb0 has its branch node in bb3. With the fallthroughs, we can see look for a branch with the same fallthrough as the outer optional terminal to match them up.

ghstack-source-id: d48c623
Pull Request resolved: #30814

Author: josephsavona

compiler/packages/babel-plugin-react-compiler/src/HIR/BuildHIR.ts

@@ -607,6 +607,7 @@ function lowerStatement(
             ),
             consequent: bodyBlock,
             alternate: continuationBlock.id,
+            fallthrough: continuationBlock.id,
             id: makeInstructionId(0),
             loc: stmt.node.loc ?? GeneratedSource,
           },
@@ -656,16 +657,13 @@ function lowerStatement(
         },
         conditionalBlock,
       );
-      /*
-       * The conditional block is empty and exists solely as conditional for
-       * (re)entering or exiting the loop
-       */
       const test = lowerExpressionToTemporary(builder, stmt.get('test'));
       const terminal: BranchTerminal = {
         kind: 'branch',
         test,
         consequent: loopBlock,
         alternate: continuationBlock.id,
+        fallthrough: conditionalBlock.id,
         id: makeInstructionId(0),
         loc: stmt.node.loc ?? GeneratedSource,
       };
@@ -975,6 +973,7 @@ function lowerStatement(
         test,
         consequent: loopBlock,
         alternate: continuationBlock.id,
+        fallthrough: conditionalBlock.id,
         id: makeInstructionId(0),
         loc,
       };
@@ -1118,6 +1117,7 @@ function lowerStatement(
           consequent: loopBlock,
           alternate: continuationBlock.id,
           loc: stmt.node.loc ?? GeneratedSource,
+          fallthrough: continuationBlock.id,
         },
         continuationBlock,
       );
@@ -1203,6 +1203,7 @@ function lowerStatement(
           test,
           consequent: loopBlock,
           alternate: continuationBlock.id,
+          fallthrough: continuationBlock.id,
           loc: stmt.node.loc ?? GeneratedSource,
         },
         continuationBlock,
@@ -1800,6 +1801,7 @@ function lowerExpression(
           test: {...testPlace},
           consequent: consequentBlock,
           alternate: alternateBlock,
+          fallthrough: continuationBlock.id,
           id: makeInstructionId(0),
           loc: exprLoc,
         },
@@ -1878,6 +1880,7 @@ function lowerExpression(
           test: {...leftPlace},
           consequent,
           alternate,
+          fallthrough: continuationBlock.id,
           id: makeInstructionId(0),
           loc: exprLoc,
         },
@@ -2611,6 +2614,7 @@ function lowerOptionalMemberExpression(
       test: {...object},
       consequent: consequent.id,
       alternate,
+      fallthrough: continuationBlock.id,
       id: makeInstructionId(0),
       loc,
     };
@@ -2750,6 +2754,7 @@ function lowerOptionalCallExpression(
       test: {...testPlace},
       consequent: consequent.id,
       alternate,
+      fallthrough: continuationBlock.id,
       id: makeInstructionId(0),
       loc,
     };
@@ -4025,6 +4030,7 @@ function lowerAssignment(
           test: {...test},
           consequent,
           alternate,
+          fallthrough: continuationBlock.id,
           id: makeInstructionId(0),
           loc,
         },

compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts

@@ -491,7 +491,7 @@ export type BranchTerminal = {
   alternate: BlockId;
   id: InstructionId;
   loc: SourceLocation;
-  fallthrough?: never;
+  fallthrough: BlockId;
 };
 
 export type SwitchTerminal = {

compiler/packages/babel-plugin-react-compiler/src/HIR/visitors.ts

@@ -660,11 +660,13 @@ export function mapTerminalSuccessors(
     case 'branch': {
       const consequent = fn(terminal.consequent);
       const alternate = fn(terminal.alternate);
+      const fallthrough = fn(terminal.fallthrough);
       return {
         kind: 'branch',
         test: terminal.test,
         consequent,
         alternate,
+        fallthrough,
         id: makeInstructionId(0),
         loc: terminal.loc,
       };
@@ -883,7 +885,6 @@ export function terminalHasFallthrough<
 >(terminal: T): terminal is U {
   switch (terminal.kind) {
     case 'maybe-throw':
-    case 'branch':
     case 'goto':
     case 'return':
     case 'throw':
@@ -892,6 +893,7 @@ export function terminalHasFallthrough<
       const _: undefined = terminal.fallthrough;
       return false;
     }
+    case 'branch':
     case 'try':
     case 'do-while':
     case 'for-of':

compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts

@@ -42,6 +42,7 @@ type IdentifierSidemap = {
   react: Set<IdentifierId>;
   maybeDepsLists: Map<IdentifierId, Array<Place>>;
   maybeDeps: Map<IdentifierId, ManualMemoDependency>;
+  optionals: Set<IdentifierId>;
 };
 
 /**
@@ -52,6 +53,7 @@ type IdentifierSidemap = {
 export function collectMaybeMemoDependencies(
   value: InstructionValue,
   maybeDeps: Map<IdentifierId, ManualMemoDependency>,
+  optional: boolean,
 ): ManualMemoDependency | null {
   switch (value.kind) {
     case 'LoadGlobal': {
@@ -69,7 +71,7 @@ export function collectMaybeMemoDependencies(
         return {
           root: object.root,
           // TODO: determine if the access is optional
-          path: [...object.path, {property: value.property, optional: false}],
+          path: [...object.path, {property: value.property, optional}],
         };
       }
       break;
@@ -162,7 +164,11 @@ function collectTemporaries(
       break;
     }
   }
-  const maybeDep = collectMaybeMemoDependencies(value, sidemap.maybeDeps);
+  const maybeDep = collectMaybeMemoDependencies(
+    value,
+    sidemap.maybeDeps,
+    sidemap.optionals.has(lvalue.identifier.id),
+  );
   // We don't expect named lvalues during this pass (unlike ValidatePreservingManualMemo)
   if (maybeDep != null) {
     sidemap.maybeDeps.set(lvalue.identifier.id, maybeDep);
@@ -338,12 +344,14 @@ export function dropManualMemoization(func: HIRFunction): void {
     func.env.config.validatePreserveExistingMemoizationGuarantees ||
     func.env.config.validateNoSetStateInRender ||
     func.env.config.enablePreserveExistingMemoizationGuarantees;
+  const optionals = findOptionalPlaces(func);
   const sidemap: IdentifierSidemap = {
     functions: new Map(),
     manualMemos: new Map(),
     react: new Set(),
     maybeDeps: new Map(),
     maybeDepsLists: new Map(),
+    optionals,
   };
   let nextManualMemoId = 0;
 
@@ -476,3 +484,46 @@ export function dropManualMemoization(func: HIRFunction): void {
     }
   }
 }
+
+function findOptionalPlaces(fn: HIRFunction): Set<IdentifierId> {
+  const optionals = new Set<IdentifierId>();
+  for (const [, block] of fn.body.blocks) {
+    if (block.terminal.kind === 'optional') {
+      const optionalTerminal = block.terminal;
+      let testBlock = fn.body.blocks.get(block.terminal.test)!;
+      loop: while (true) {
+        const terminal = testBlock.terminal;
+        switch (terminal.kind) {
+          case 'branch': {
+            if (terminal.fallthrough === optionalTerminal.fallthrough) {
+              // found it
+              const consequent = fn.body.blocks.get(terminal.consequent)!;
+              const last = consequent.instructions.at(-1);
+              if (last !== undefined && last.value.kind === 'StoreLocal') {
+                optionals.add(last.value.value.identifier.id);
+              }
+              break loop;
+            } else {
+              testBlock = fn.body.blocks.get(terminal.fallthrough)!;
+            }
+            break;
+          }
+          case 'optional':
+          case 'logical':
+          case 'sequence':
+          case 'ternary': {
+            testBlock = fn.body.blocks.get(terminal.fallthrough)!;
+            break;
+          }
+          default: {
+            CompilerError.invariant(false, {
+              reason: `Unexpected terminal in optional`,
+              loc: terminal.loc,
+            });
+          }
+        }
+      }
+    }
+  }
+  return optionals;
+}

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/AlignReactiveScopesToBlockScopesHIR.ts

@@ -140,7 +140,7 @@ export function alignReactiveScopesToBlockScopesHIR(fn: HIRFunction): void {
     }
 
     const fallthrough = terminalFallthrough(terminal);
-    if (fallthrough !== null) {
+    if (fallthrough !== null && terminal.kind !== 'branch') {
       /*
        * Any currently active scopes that overlaps the block-fallthrough range
        * need their range extended to at least the first instruction of the

compiler/packages/babel-plugin-react-compiler/src/Validation/ValidatePreservedManualMemoization.ts

@@ -167,7 +167,10 @@ function compareDeps(
 
   let isSubpath = true;
   for (let i = 0; i < Math.min(inferred.path.length, source.path.length); i++) {
-    if (inferred.path[i].property !== source.path[i].property) {
+    if (
+      inferred.path[i].property !== source.path[i].property ||
+      inferred.path[i].optional !== source.path[i].optional
+    ) {
       isSubpath = false;
       break;
     }
@@ -339,7 +342,11 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
         return null;
       }
       default: {
-        const dep = collectMaybeMemoDependencies(value, this.temporaries);
+        const dep = collectMaybeMemoDependencies(
+          value,
+          this.temporaries,
+          false,
+        );
         if (value.kind === 'StoreLocal' || value.kind === 'StoreContext') {
           const storeTarget = value.lvalue.place;
           state.manualMemoState?.decls.add(