Update on "[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-poisoned]

Author: josephsavona

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;
 
@@ -481,6 +489,40 @@ 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/MergeReactiveScopesThatInvalidateTogether.ts

@@ -8,7 +8,6 @@
 import {CompilerError} from '..';
 import {
   DeclarationId,
-  DependencyPath,
   InstructionId,
   InstructionKind,
   Place,

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(