From 8f4c47b3b792281081d8133667142ae7c2a7b2c8 Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Mon, 26 Aug 2024 12:20:22 -0700
Subject: [PATCH] [compiler] Infer optional manual memo deps

When inferring dependencies of manual memoization in DropManualMemo, we now infer which parts of a dependency path were from optional member expressions.

[ghstack-poisoned]
---
 .../src/Inference/DropManualMemoization.ts    | 46 ++++++++++++++++++-
 .../ValidatePreservedManualMemoization.ts     | 11 ++++-
 2 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts
index 153dd7b1773ba..fdd9bcc968aef 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts
+++ b/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;
diff --git a/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidatePreservedManualMemoization.ts b/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidatePreservedManualMemoization.ts
index 5a9a947d88bfd..6d948fad9711a 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidatePreservedManualMemoization.ts
+++ b/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(