[patch] ObjectMethods should have the same scope as their parent

ObjectExpressions 

--- 

Currently, we're removing all reactive scopes containing object methods. This 
could produce incorrect output as object method instructions may still be 
included in other reactive scopes (and will lose their dependencies).

compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts

@@ -35,6 +35,7 @@ import {
 } from "../Optimization";
 import {
   CodegenFunction,
+  alignObjectMethodScopes,
   alignReactiveScopesToBlockScopes,
   assertScopeInstructionsWithinScopes,
   buildReactiveBlocks,
@@ -43,7 +44,6 @@ import {
   extractScopeDeclarationsFromDestructuring,
   flattenReactiveLoops,
   flattenScopesWithHooks,
-  flattenScopesWithObjectMethods,
   inferReactiveScopeVariables,
   memoizeFbtOperandsInSameScope,
   mergeOverlappingReactiveScopes,
@@ -198,6 +198,13 @@ function* runWithEnvironment(
   inferReactiveScopeVariables(hir);
   yield log({ kind: "hir", name: "InferReactiveScopeVariables", value: hir });
 
+  alignObjectMethodScopes(hir);
+  yield log({
+    kind: "hir",
+    name: "AlignObjectMethodScopes",
+    value: hir,
+  });
+
   const reactiveFunction = buildReactiveFunction(hir);
   yield log({
     kind: "reactive",
@@ -265,13 +272,6 @@ function* runWithEnvironment(
     value: reactiveFunction,
   });
 
-  flattenScopesWithObjectMethods(reactiveFunction);
-  yield log({
-    kind: "reactive",
-    name: "FlattenScopesWithObjectMethods",
-    value: reactiveFunction,
-  });
-
   propagateScopeDependencies(reactiveFunction);
   yield log({
     kind: "reactive",

compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/AlignObjectMethodScopes.ts

@@ -0,0 +1,101 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import { CompilerError } from "..";
+import {
+  GeneratedSource,
+  HIRFunction,
+  Identifier,
+  ReactiveScope,
+  makeInstructionId,
+} from "../HIR";
+import { eachInstructionValueOperand } from "../HIR/visitors";
+import DisjointSet from "../Utils/DisjointSet";
+
+/**
+ * Align scopes of object method values to that of their enclosing object expressions.
+ * To produce a well-formed JS program in Codegen, object methods and object expressions
+ * must be in the same ReactiveBlock as object method definitions must be inlined.
+ */
+
+function findScopesToMerge(fn: HIRFunction): DisjointSet<ReactiveScope> {
+  const objectMethodDecls: Set<Identifier> = new Set();
+  const mergeScopesBuilder = new DisjointSet<ReactiveScope>();
+
+  for (const [_, block] of fn.body.blocks) {
+    for (const { lvalue, value } of block.instructions) {
+      if (value.kind === "ObjectMethod") {
+        objectMethodDecls.add(lvalue.identifier);
+      } else if (value.kind === "ObjectExpression") {
+        for (const operand of eachInstructionValueOperand(value)) {
+          if (objectMethodDecls.has(operand.identifier)) {
+            const operandScope = operand.identifier.scope;
+            const lvalueScope = lvalue.identifier.scope;
+
+            CompilerError.invariant(
+              operandScope != null && lvalueScope != null,
+              {
+                reason:
+                  "Internal error: Expected all ObjectExpressions and ObjectMethods to have non-null scope.",
+                suggestions: null,
+                loc: GeneratedSource,
+              }
+            );
+            mergeScopesBuilder.union([operandScope, lvalueScope]);
+          }
+        }
+      }
+    }
+  }
+  return mergeScopesBuilder;
+}
+
+export function alignObjectMethodScopes(fn: HIRFunction): void {
+  // Handle inner functions: we assume that Scopes are disjoint across functions
+  for (const [_, block] of fn.body.blocks) {
+    for (const { value } of block.instructions) {
+      if (
+        value.kind === "ObjectMethod" ||
+        value.kind === "FunctionExpression"
+      ) {
+        alignObjectMethodScopes(value.loweredFunc.func);
+      }
+    }
+  }
+
+  const scopeGroupsMap = findScopesToMerge(fn).canonicalize();
+  /**
+   * Step 1: Merge affected scopes to their canonical root.
+   */
+  for (const [scope, root] of scopeGroupsMap) {
+    if (scope !== root) {
+      root.range.start = makeInstructionId(
+        Math.min(scope.range.start, root.range.start)
+      );
+      root.range.end = makeInstructionId(
+        Math.max(scope.range.end, root.range.end)
+      );
+    }
+  }
+
+  /**
+   * Step 2: Repoint identifiers whose scopes were merged.
+   */
+  for (const [_, block] of fn.body.blocks) {
+    for (const {
+      lvalue: { identifier },
+    } of block.instructions) {
+      if (identifier.scope != null) {
+        const root = scopeGroupsMap.get(identifier.scope);
+        if (root != null) {
+          identifier.scope = root;
+        }
+        // otherwise, this identifier's scope was not affected by this pass
+      }
+    }
+  }
+}

compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/FlattenScopesWithObjectMethods.ts

@@ -1,58 +0,0 @@
-/*
- * Copyright (c) Meta Platforms, Inc. and affiliates.
- *
- * This source code is licensed under the MIT license found in the
- * LICENSE file in the root directory of this source tree.
- */
-
-import {
-  InstructionId,
-  ReactiveFunction,
-  ReactiveScopeBlock,
-  ReactiveStatement,
-  ReactiveValue,
-} from "../HIR";
-import {
-  ReactiveFunctionTransform,
-  Transformed,
-  visitReactiveFunction,
-} from "./visitors";
-
-export function flattenScopesWithObjectMethods(fn: ReactiveFunction): void {
-  visitReactiveFunction(fn, new Transform(), {
-    hasObjectMethod: false,
-  });
-}
-
-type State = {
-  hasObjectMethod: boolean;
-};
-
-class Transform extends ReactiveFunctionTransform<State> {
-  override transformScope(
-    scope: ReactiveScopeBlock,
-    outerState: State
-  ): Transformed<ReactiveStatement> {
-    const innerState: State = {
-      hasObjectMethod: false,
-    };
-    this.visitScope(scope, innerState);
-    outerState.hasObjectMethod ||= innerState.hasObjectMethod;
-    if (innerState.hasObjectMethod) {
-      return { kind: "replace-many", value: scope.instructions };
-    } else {
-      return { kind: "keep" };
-    }
-  }
-
-  override visitValue(
-    id: InstructionId,
-    value: ReactiveValue,
-    state: State
-  ): void {
-    this.traverseValue(id, value, state);
-    if (value.kind === "ObjectMethod") {
-      state.hasObjectMethod = true;
-    }
-  }
-}

compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/index.ts

@@ -5,6 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+export { alignObjectMethodScopes } from "./AlignObjectMethodScopes";
 export { alignReactiveScopesToBlockScopes } from "./AlignReactiveScopesToBlockScopes";
 export { assertScopeInstructionsWithinScopes } from "./AssertScopeInstructionsWithinScope";
 export { buildReactiveBlocks } from "./BuildReactiveBlocks";
@@ -16,7 +17,6 @@ export {
 export { extractScopeDeclarationsFromDestructuring } from "./ExtractScopeDeclarationsFromDestructuring";
 export { flattenReactiveLoops } from "./FlattenReactiveLoops";
 export { flattenScopesWithHooks } from "./FlattenScopesWithHooks";
-export { flattenScopesWithObjectMethods } from "./FlattenScopesWithObjectMethods";
 export { inferReactiveScopeVariables } from "./InferReactiveScopeVariables";
 export { memoizeFbtOperandsInSameScope } from "./MemoizeFbtOperandsInSameScope";
 export { mergeOverlappingReactiveScopes } from "./MergeOverlappingReactiveScopes";

compiler/packages/babel-plugin-react-forget/src/Utils/utils.ts

@@ -68,6 +68,12 @@ export function Set_union<T>(a: Set<T>, b: Set<T>): Set<T> {
   return union;
 }
 
+export function nonNull<T extends NonNullable<U>, U>(
+  value: T | null | undefined
+): value is T {
+  return value != null;
+}
+
 export function hasNode<T>(
   input: NodePath<T | null | undefined>
 ): input is NodePath<NonNullable<T>> {

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-literal-method-derived-in-ternary-consequent.expect.md

@@ -24,17 +24,28 @@ export const FIXTURE_ENTRYPOINT = {
 ## Code
 
 ```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
 import { identity, createHookWrapper } from "shared-runtime";
 
 function useHook(t17) {
+  const $ = useMemoCache(3);
   const { isCond, value } = t17;
-  return isCond
-    ? identity({
-        getValue() {
-          return value;
-        },
-      })
-    : 42;
+  let t0;
+  if ($[0] !== isCond || $[1] !== value) {
+    t0 = isCond
+      ? identity({
+          getValue() {
+            return value;
+          },
+        })
+      : 42;
+    $[0] = isCond;
+    $[1] = value;
+    $[2] = t0;
+  } else {
+    t0 = $[2];
+  }
+  return t0;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-literal-method-in-ternary-consequent.expect.md

@@ -24,17 +24,28 @@ export const FIXTURE_ENTRYPOINT = {
 ## Code
 
 ```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
 import { createHookWrapper } from "shared-runtime";
 
 function useHook(t15) {
+  const $ = useMemoCache(3);
   const { isCond, value } = t15;
-  return isCond
-    ? {
-        getValue() {
-          return value;
-        },
-      }
-    : 42;
+  let t0;
+  if ($[0] !== isCond || $[1] !== value) {
+    t0 = isCond
+      ? {
+          getValue() {
+            return value;
+          },
+        }
+      : 42;
+    $[0] = isCond;
+    $[1] = value;
+    $[2] = t0;
+  } else {
+    t0 = $[2];
+  }
+  return t0;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-maybe-alias.expect.md

@@ -27,20 +27,31 @@ export const FIXTURE_ENTRYPOINT = {
 ## Code
 
 ```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
 import { createHookWrapper, setProperty } from "shared-runtime";
 function useHook(props) {
-  const x = {
-    getX() {
-      return props;
-    },
-  };
+  const $ = useMemoCache(2);
+  let t0;
+  if ($[0] !== props) {
+    const x = {
+      getX() {
+        return props;
+      },
+    };
 
-  const y = {
-    getY() {
-      return "y";
-    },
-  };
-  return setProperty(x, y);
+    const y = {
+      getY() {
+        return "y";
+      },
+    };
+
+    t0 = setProperty(x, y);
+    $[0] = props;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-3.expect.md

@@ -25,17 +25,28 @@ export const FIXTURE_ENTRYPOINT = {
 ## Code
 
 ```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
 import { createHookWrapper, mutate } from "shared-runtime";
 
 function useHook(a) {
-  const x = { a };
-  const obj = {
-    method() {
-      mutate(x);
-      return x;
-    },
-  };
-  return obj.method();
+  const $ = useMemoCache(2);
+  let t0;
+  if ($[0] !== a) {
+    const x = { a };
+    const obj = {
+      method() {
+        mutate(x);
+        return x;
+      },
+    };
+
+    t0 = obj.method();
+    $[0] = a;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-aliased-mutate-after.expect.md

@@ -27,23 +27,23 @@ export const FIXTURE_ENTRYPOINT = {
 import { unstable_useMemoCache as useMemoCache } from "react";
 import { createHookWrapper, mutate, mutateAndReturn } from "shared-runtime";
 function useHook(t21) {
-  const $ = useMemoCache(1);
+  const $ = useMemoCache(2);
   const { value } = t21;
-  const x = mutateAndReturn({ value });
-  let t0;
-  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t0 = {
+  let obj;
+  if ($[0] !== value) {
+    const x = mutateAndReturn({ value });
+    obj = {
       getValue() {
         return value;
       },
     };
-    $[0] = t0;
-  } else {
-    t0 = $[0];
-  }
-  const obj = t0;
 
-  mutate(x);
+    mutate(x);
+    $[0] = value;
+    $[1] = obj;
+  } else {
+    obj = $[1];
+  }
   return obj;
 }
 

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-derived-value.expect.md

@@ -26,7 +26,7 @@ export const FIXTURE_ENTRYPOINT = {
 import { unstable_useMemoCache as useMemoCache } from "react";
 import { createHookWrapper, mutateAndReturn } from "shared-runtime";
 function useHook(t18) {
-  const $ = useMemoCache(3);
+  const $ = useMemoCache(4);
   const { value } = t18;
   let t0;
   if ($[0] !== value) {
@@ -38,15 +38,16 @@ function useHook(t18) {
   }
   const x = t0;
   let t1;
-  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+  if ($[2] !== x) {
     t1 = {
       getValue() {
         return x;
       },
     };
-    $[2] = t1;
+    $[2] = x;
+    $[3] = t1;
   } else {
-    t1 = $[2];
+    t1 = $[3];
   }
   const obj = t1;
   return obj;

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-hook-dep.expect.md

@@ -26,18 +26,19 @@ export const FIXTURE_ENTRYPOINT = {
 import { createHookWrapper } from "shared-runtime";
 import { useState, unstable_useMemoCache as useMemoCache } from "react";
 function useFoo() {
-  const $ = useMemoCache(1);
+  const $ = useMemoCache(2);
   const [state] = useState(false);
   let t0;
-  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+  if ($[0] !== state) {
     t0 = {
       func() {
         return state;
       },
     };
-    $[0] = t0;
+    $[0] = state;
+    $[1] = t0;
   } else {
-    t0 = $[0];
+    t0 = $[1];
   }
   return t0;
 }

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-mutated-after.expect.md

@@ -24,17 +24,26 @@ export const FIXTURE_ENTRYPOINT = {
 ## Code
 
 ```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
 import { createHookWrapper, mutate, mutateAndReturn } from "shared-runtime";
 function useHook(t21) {
+  const $ = useMemoCache(2);
   const { value } = t21;
-  const x = mutateAndReturn({ value });
-  const obj = {
-    getValue() {
-      return x;
-    },
-  };
+  let obj;
+  if ($[0] !== value) {
+    const x = mutateAndReturn({ value });
+    obj = {
+      getValue() {
+        return x;
+      },
+    };
 
-  mutate(obj);
+    mutate(obj);
+    $[0] = value;
+    $[1] = obj;
+  } else {
+    obj = $[1];
+  }
   return obj;
 }
 

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-shorthand-method-1.expect.md

@@ -27,7 +27,7 @@ export const FIXTURE_ENTRYPOINT = {
 import { unstable_useMemoCache as useMemoCache } from "react";
 import { createHookWrapper } from "shared-runtime";
 function useHook(t16) {
-  const $ = useMemoCache(4);
+  const $ = useMemoCache(5);
   const { a, b } = t16;
   let t0;
   if ($[0] !== a) {
@@ -40,17 +40,18 @@ function useHook(t16) {
     t0 = $[1];
   }
   let t1;
-  if ($[2] !== t0) {
+  if ($[2] !== b || $[3] !== t0) {
     t1 = {
       x: t0,
       y() {
         return [b];
       },
     };
-    $[2] = t0;
-    $[3] = t1;
+    $[2] = b;
+    $[3] = t0;
+    $[4] = t1;
   } else {
-    t1 = $[3];
+    t1 = $[4];
   }
   return t1;
 }

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-shorthand-method-2.expect.md

@@ -28,7 +28,7 @@ import { unstable_useMemoCache as useMemoCache } from "react";
 import { createHookWrapper } from "shared-runtime";
 
 function useHook(t16) {
-  const $ = useMemoCache(7);
+  const $ = useMemoCache(8);
   const { a, b, c } = t16;
   let t0;
   if ($[0] !== a) {
@@ -38,16 +38,16 @@ function useHook(t16) {
   } else {
     t0 = $[1];
   }
-  let t1;
-  if ($[2] !== c) {
-    t1 = { c };
-    $[2] = c;
-    $[3] = t1;
-  } else {
-    t1 = $[3];
-  }
   let t2;
-  if ($[4] !== t0 || $[5] !== t1) {
+  if ($[2] !== b || $[3] !== c || $[4] !== t0) {
+    let t1;
+    if ($[6] !== c) {
+      t1 = { c };
+      $[6] = c;
+      $[7] = t1;
+    } else {
+      t1 = $[7];
+    }
     t2 = {
       x: t0,
       y() {
@@ -55,11 +55,12 @@ function useHook(t16) {
       },
       z: t1,
     };
+    $[2] = b;
+    $[3] = c;
     $[4] = t0;
-    $[5] = t1;
-    $[6] = t2;
+    $[5] = t2;
   } else {
-    t2 = $[6];
+    t2 = $[5];
   }
   return t2;
 }

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-shorthand-method-nested.expect.md

@@ -35,11 +35,11 @@ import { useState, unstable_useMemoCache as useMemoCache } from "react";
 import { createHookWrapper } from "shared-runtime";
 
 function useHook(t21) {
-  const $ = useMemoCache(1);
+  const $ = useMemoCache(3);
   const { value } = t21;
   const [state] = useState(false);
   let t0;
-  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+  if ($[0] !== value || $[1] !== state) {
     t0 = {
       getX() {
         return {
@@ -51,9 +51,11 @@ function useHook(t21) {
         };
       },
     };
-    $[0] = t0;
+    $[0] = value;
+    $[1] = state;
+    $[2] = t0;
   } else {
-    t0 = $[0];
+    t0 = $[2];
   }
   return t0;
 }