[hir] Do not memoize object methods separately

Object methods must not be cached independently, so this PR flattens the 
reactive scope to prevent memoization. 

In the future, we can combine the FlattenScopesWithObjectMethods and 
FlattedScopesWithHooks passes by making them more modular. But this works for 
now.

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

@@ -40,6 +40,7 @@ import {
   extractScopeDeclarationsFromDestructuring,
   flattenReactiveLoops,
   flattenScopesWithHooks,
+  flattenScopesWithObjectMethods,
   inferReactiveScopeVariables,
   memoizeFbtOperandsInSameScope,
   mergeConsecutiveScopes,
@@ -230,6 +231,13 @@ export function* run(
     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/HIR/HIR.ts

@@ -1083,6 +1083,10 @@ export function makeInstructionId(id: number): InstructionId {
   return id as InstructionId;
 }
 
+export function isObjectMethodType(id: Identifier): boolean {
+  return id.type.kind == "ObjectMethod";
+}
+
 export function isObjectType(id: Identifier): boolean {
   return id.type.kind === "Object";
 }

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

@@ -9,7 +9,13 @@ import { CompilerError } from "../CompilerError";
 
 export type BuiltInType = PrimitiveType | FunctionType | ObjectType;
 
-export type Type = BuiltInType | PhiType | TypeVar | PolyType | PropType;
+export type Type =
+  | BuiltInType
+  | PhiType
+  | TypeVar
+  | PolyType
+  | PropType
+  | ObjectMethod;
 export type PrimitiveType = { kind: "Primitive" };
 
 /**
@@ -54,6 +60,11 @@ export type PropType = {
   object: Type;
   propertyName: string;
 };
+
+export type ObjectMethod = {
+  kind: "ObjectMethod";
+};
+
 /**
  * Simulated opaque type for TypeId to prevent using normal numbers as ids
  * accidentally.
@@ -87,7 +98,8 @@ export function typeEquals(tA: Type, tB: Type): boolean {
     objectTypeEquals(tA, tB) ||
     primitiveTypeEquals(tA, tB) ||
     polyTypeEquals(tA, tB) ||
-    phiTypeEquals(tA, tB)
+    phiTypeEquals(tA, tB) ||
+    objectMethodTypeEquals(tA, tB)
   );
 }
 
@@ -102,6 +114,10 @@ function typeKindCheck(tA: Type, tb: Type, type: string): boolean {
   return tA.kind === type && tb.kind === type;
 }
 
+function objectMethodTypeEquals(tA: Type, tB: Type): boolean {
+  return typeKindCheck(tA, tB, "ObjectMethod");
+}
+
 function primitiveTypeEquals(tA: Type, tB: Type): boolean {
   return typeKindCheck(tA, tB, "Primitive");
 }

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

@@ -0,0 +1,58 @@
+/**
+ * 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/PropagateScopeDependencies.ts

@@ -11,6 +11,7 @@ import {
   IdentifierId,
   InstructionId,
   InstructionKind,
+  isObjectMethodType,
   isRefValueType,
   isUseRefType,
   makeInstructionId,
@@ -330,6 +331,12 @@ class Context {
       return false;
     }
 
+    // object methods are not deps because they will be codegen'ed back in to
+    // the object literal.
+    if (isObjectMethodType(maybeDependency.identifier)) {
+      return false;
+    }
+
     const identifier = maybeDependency.identifier;
     // If this operand is used in a scope, has a dynamic value, and was defined
     // before this scope, then its a dependency of the scope.

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

@@ -15,6 +15,7 @@ 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 { mergeConsecutiveScopes } from "./MergeConsecutiveScopes";

compiler/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts

@@ -255,7 +255,6 @@ function* generateInstructionTypes(
       break;
     }
 
-    case "ObjectMethod":
     case "FunctionExpression": {
       yield* generate(value.loweredFunc.func);
       break;
@@ -266,6 +265,11 @@ function* generateInstructionTypes(
       break;
     }
 
+    case "ObjectMethod": {
+      yield equation(left, { kind: "ObjectMethod" });
+      break;
+    }
+
     case "DeclareLocal":
     case "NewExpression":
     case "JsxExpression":

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

@@ -0,0 +1,47 @@
+
+## Input
+
+```javascript
+import { mutate } from "shared-runtime";
+
+function Component(a) {
+  const x = { a };
+  let obj = {
+    method() {
+      mutate(x);
+      return x.a;
+    },
+  };
+  return obj.method();
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ x: 1 }],
+};
+
+```
+
+## Code
+
+```javascript
+import { mutate } from "shared-runtime";
+
+function Component(a) {
+  const x = { a };
+  const obj = {
+    method() {
+      mutate(x);
+      return x.a;
+    },
+  };
+  return obj.method();
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ x: 1 }],
+};
+
+```
+      

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-3.expect.md~051f3e57 ([hir] Do not memoize object methods separately)

@@ -0,0 +1,47 @@
+
+## Input
+
+```javascript
+import { mutate } from "shared-runtime";
+
+function Component(a) {
+  const x = { a };
+  let obj = {
+    method() {
+      mutate(x);
+      return x;
+    },
+  };
+  return obj.method();
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ x: 1 }, { a: 2 }, { b: 2 }],
+};
+
+```
+
+## Code
+
+```javascript
+import { mutate } from "shared-runtime";
+
+function Component(a) {
+  const x = { a };
+  const obj = {
+    method() {
+      mutate(x);
+      return x;
+    },
+  };
+  return obj.method();
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ x: 1 }, { a: 2 }, { b: 2 }],
+};
+
+```
+      

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

@@ -0,0 +1,17 @@
+import { mutate } from "shared-runtime";
+
+function Component(a) {
+  const x = { a };
+  let obj = {
+    method() {
+      mutate(x);
+      return x.a;
+    },
+  };
+  return obj.method();
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ x: 1 }],
+};

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

@@ -0,0 +1,49 @@
+
+## Input
+
+```javascript
+function Component() {
+  let obj = {
+    method() {
+      return 1;
+    },
+  };
+  return obj.method();
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ x: 1 }, { a: 2 }, { b: 2 }],
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Component() {
+  const $ = useMemoCache(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    const obj = {
+      method() {
+        return 1;
+      },
+    };
+
+    t0 = obj.method();
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ x: 1 }, { a: 2 }, { b: 2 }],
+};
+
+```
+      

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand.js

@@ -0,0 +1,13 @@
+function Component() {
+  let obj = {
+    method() {
+      return 1;
+    },
+  };
+  return obj.method();
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ x: 1 }, { a: 2 }, { b: 2 }],
+};

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

@@ -2,6 +2,7 @@
 ## Input
 
 ```javascript
+// @debug
 function Component({ a, b }) {
   return {
     x: function () {
@@ -15,7 +16,7 @@ function Component({ a, b }) {
 
 export const FIXTURE_ENTRYPOINT = {
   fn: Component,
-  params: [{ x: 1 }, { a: 2 }, { b: 2 }],
+  params: [{ x: 1 }, { a: 2 }],
 };
 
 ```
@@ -23,9 +24,9 @@ export const FIXTURE_ENTRYPOINT = {
 ## Code
 
 ```javascript
-import { unstable_useMemoCache as useMemoCache } from "react";
+import { unstable_useMemoCache as useMemoCache } from "react"; // @debug
 function Component(t16) {
-  const $ = useMemoCache(7);
+  const $ = useMemoCache(4);
   const { a, b } = t16;
   const c_0 = $[0] !== a;
   let t0;
@@ -38,36 +39,26 @@ function Component(t16) {
   } else {
     t0 = $[1];
   }
-  const c_2 = $[2] !== b;
+  const c_2 = $[2] !== t0;
   let t1;
   if (c_2) {
-    $[2] = b;
-    $[3] = t1;
-  } else {
-    t1 = $[3];
-  }
-  const c_4 = $[4] !== t0;
-  const c_5 = $[5] !== t1;
-  let t2;
-  if (c_4 || c_5) {
-    t2 = {
+    t1 = {
       x: t0,
       y() {
         return [b];
       },
     };
-    $[4] = t0;
-    $[5] = t1;
-    $[6] = t2;
+    $[2] = t0;
+    $[3] = t1;
   } else {
-    t2 = $[6];
+    t1 = $[3];
   }
-  return t2;
+  return t1;
 }
 
 export const FIXTURE_ENTRYPOINT = {
   fn: Component,
-  params: [{ x: 1 }, { a: 2 }, { b: 2 }],
+  params: [{ x: 1 }, { a: 2 }],
 };
 
 ```

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

@@ -1,3 +1,4 @@
+// @debug
 function Component({ a, b }) {
   return {
     x: function () {
@@ -11,5 +12,5 @@ function Component({ a, b }) {
 
 export const FIXTURE_ENTRYPOINT = {
   fn: Component,
-  params: [{ x: 1 }, { a: 2 }, { b: 2 }],
+  params: [{ x: 1 }, { a: 2 }],
 };

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

@@ -24,7 +24,7 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { unstable_useMemoCache as useMemoCache } from "react";
 function Component(t16) {
-  const $ = useMemoCache(10);
+  const $ = useMemoCache(7);
   const { a, b, c } = t16;
   const c_0 = $[0] !== a;
   let t0;
@@ -35,43 +35,33 @@ function Component(t16) {
   } else {
     t0 = $[1];
   }
-  const c_2 = $[2] !== b;
+  const c_2 = $[2] !== c;
   let t1;
   if (c_2) {
-    $[2] = b;
+    t1 = { c };
+    $[2] = c;
     $[3] = t1;
   } else {
     t1 = $[3];
   }
-  const c_4 = $[4] !== c;
+  const c_4 = $[4] !== t0;
+  const c_5 = $[5] !== t1;
   let t2;
-  if (c_4) {
-    t2 = { c };
-    $[4] = c;
-    $[5] = t2;
-  } else {
-    t2 = $[5];
-  }
-  const c_6 = $[6] !== t0;
-  const c_7 = $[7] !== t1;
-  const c_8 = $[8] !== t2;
-  let t3;
-  if (c_6 || c_7 || c_8) {
-    t3 = {
+  if (c_4 || c_5) {
+    t2 = {
       x: t0,
       y() {
         return [b];
       },
-      z: t2,
+      z: t1,
     };
-    $[6] = t0;
-    $[7] = t1;
-    $[8] = t2;
-    $[9] = t3;
+    $[4] = t0;
+    $[5] = t1;
+    $[6] = t2;
   } else {
-    t3 = $[9];
+    t2 = $[6];
   }
-  return t3;
+  return t2;
 }
 
 export const FIXTURE_ENTRYPOINT = {