[hir] Wire up object method lowering and codegen

Object methods are lowered to functions and added to ObjectExpression. The 
codegen is interesting because we shouldn't emit code that lowers the object 
method into a separate statement and then stores it into an object expression. 

An shorthand object method has different semantics than an object method using 
the function syntax, so we need to preserve the shorthand object method syntax 
in the generated code. 

To do this, we don't immediately generate an AST node for the ObjectMethod but 
instead store it in a side table during codegen. Only when emitting code for an 
ObjectExpression, we lookup this side table and emit the object method inline in 
the body.

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

@@ -1277,6 +1277,23 @@ function lowerStatement(
   }
 }
 
+function lowerObjectMethod(
+  builder: HIRBuilder,
+  property: NodePath<t.ObjectMethod>
+): InstructionValue {
+  const loc = property.node.loc ?? GeneratedSource;
+  const loweredFunc = lowerFunction(builder, property);
+  if (!loweredFunc) {
+    return { kind: "UnsupportedNode", node: property.node, loc: loc };
+  }
+
+  return {
+    kind: "ObjectMethod",
+    loc,
+    loweredFunc,
+  };
+}
+
 function lowerObjectPropertyKey(
   builder: HIRBuilder,
   key: t.PrivateName | t.Expression
@@ -1379,6 +1396,22 @@ function lowerExpression(
             kind: "Spread",
             place,
           });
+        } else if (propertyPath.isObjectMethod()) {
+          const method = lowerObjectMethod(builder, propertyPath);
+          const place = lowerValueToTemporary(builder, method);
+          const loweredKey = lowerObjectPropertyKey(
+            builder,
+            propertyPath.node.key
+          );
+          if (!loweredKey) {
+            continue;
+          }
+          properties.push({
+            kind: "ObjectProperty",
+            type: "method",
+            place,
+            key: loweredKey,
+          });
         } else {
           builder.errors.push({
             reason: `(BuildHIR::lowerExpression) Handle ${propertyPath.type} properties in ObjectExpression`,

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

@@ -9,6 +9,7 @@ type FindContextIdentifierState = {
     | NodePath<t.FunctionDeclaration>
     | NodePath<t.FunctionExpression>
     | NodePath<t.ArrowFunctionExpression>
+    | NodePath<t.ObjectMethod>
   >;
   reassigned: Set<t.Identifier>;
   referenced: Set<t.Identifier>;
@@ -75,6 +76,12 @@ export function findContextIdentifiers(
         const left = path.get("left");
         handleAssignment(state.reassigned, left);
       },
+      ObjectMethod(
+        fn: NodePath<t.ObjectMethod>,
+        state: FindContextIdentifierState
+      ): void {
+        state.currentLambda.push(fn);
+      },
       Identifier(
         path: NodePath<t.Identifier>,
         state: FindContextIdentifierState
@@ -93,7 +100,8 @@ function handleIdentifier(
   currentLambda:
     | NodePath<t.FunctionDeclaration>
     | NodePath<t.FunctionExpression>
-    | NodePath<t.ArrowFunctionExpression>,
+    | NodePath<t.ArrowFunctionExpression>
+    | NodePath<t.ObjectMethod>,
   referenced: Set<t.Identifier>,
   path: NodePath<t.Identifier>
 ): void {

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

@@ -732,6 +732,7 @@ export type InstructionValue =
       properties: Array<ObjectProperty | SpreadPattern>;
       loc: SourceLocation;
     }
+  | ObjectMethod
   | ArrayExpression
   | { kind: "JsxFragment"; children: Array<Place>; loc: SourceLocation }
   | {

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

@@ -32,7 +32,10 @@ export function mergeConsecutiveBlocks(fn: HIRFunction): void {
   const merged = new MergedBlocks();
   for (const [, block] of fn.body.blocks) {
     for (const instr of block.instructions) {
-      if (instr.value.kind === "FunctionExpression") {
+      if (
+        instr.value.kind === "FunctionExpression" ||
+        instr.value.kind === "ObjectMethod"
+      ) {
         mergeConsecutiveBlocks(instr.value.loweredFunc.func);
       }
     }

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

@@ -462,6 +462,7 @@ export function printInstructionValue(instrValue: ReactiveValue): string {
       )}]`;
       break;
     }
+    case "ObjectMethod":
     case "FunctionExpression": {
       const fn = printFunction(instrValue.loweredFunc.func)
         .split("\n")

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

@@ -169,6 +169,7 @@ export function* eachInstructionValueOperand(
       }
       break;
     }
+    case "ObjectMethod":
     case "FunctionExpression": {
       yield* instrValue.loweredFunc.dependencies;
       break;
@@ -456,6 +457,7 @@ export function mapInstructionOperands(
       instrValue.children = instrValue.children.map((e) => fn(e));
       break;
     }
+    case "ObjectMethod":
     case "FunctionExpression": {
       instrValue.loweredFunc.dependencies =
         instrValue.loweredFunc.dependencies.map((d) => fn(d));

compiler/packages/babel-plugin-react-forget/src/Inference/AnalyseFunctions.ts

@@ -8,12 +8,12 @@
 import { CompilerError } from "../CompilerError";
 import {
   Effect,
-  FunctionExpression,
   HIRFunction,
   Identifier,
   isRefValueType,
   isSetStateType,
   isUseRefType,
+  LoweredFunction,
   Place,
   ReactiveScopeDependency,
 } from "../HIR";
@@ -70,9 +70,10 @@ export default function analyseFunctions(func: HIRFunction): void {
   for (const [_, block] of func.body.blocks) {
     for (const instr of block.instructions) {
       switch (instr.value.kind) {
+        case "ObjectMethod":
         case "FunctionExpression": {
           lower(instr.value.loweredFunc.func);
-          infer(instr.value, state, func.context);
+          infer(instr.value.loweredFunc, state, func.context);
           break;
         }
         case "PropertyLoad": {
@@ -113,12 +114,12 @@ function lower(func: HIRFunction): void {
 }
 
 function infer(
-  value: FunctionExpression,
+  loweredFunc: LoweredFunction,
   state: IdentifierState,
   context: Place[]
 ): void {
   const mutations = new Map<string, Effect>();
-  for (const operand of value.loweredFunc.func.context) {
+  for (const operand of loweredFunc.func.context) {
     if (
       isMutatedOrReassigned(operand.identifier) &&
       operand.identifier.name !== null
@@ -128,7 +129,7 @@ function infer(
     operand.identifier.mutableRange.end = operand.identifier.mutableRange.start;
   }
 
-  for (const dep of value.loweredFunc.dependencies) {
+  for (const dep of loweredFunc.dependencies) {
     let name: string | null = null;
 
     if (state.properties.has(dep.identifier)) {
@@ -174,7 +175,7 @@ function infer(
     const effect = mutations.get(place.identifier.name);
     if (effect !== undefined) {
       place.effect = effect === Effect.Unknown ? Effect.Capture : effect;
-      value.loweredFunc.dependencies.push(place);
+      loweredFunc.dependencies.push(place);
     }
   }
 }

compiler/packages/babel-plugin-react-forget/src/Inference/InferReferenceEffects.ts

@@ -714,6 +714,7 @@ function inferBlock(
         valueKind = ValueKind.Immutable;
         break;
       }
+      case "ObjectMethod":
       case "FunctionExpression": {
         let hasMutableOperand = false;
         for (const operand of eachInstructionOperand(instr)) {

compiler/packages/babel-plugin-react-forget/src/Optimization/ConstantPropagation.ts

@@ -106,7 +106,10 @@ function applyConstantPropagation(
   const functionDependencies = new Set<IdentifierId>();
   for (const [, block] of fn.body.blocks) {
     for (const instr of block.instructions) {
-      if (instr.value.kind === "FunctionExpression") {
+      if (
+        instr.value.kind === "FunctionExpression" ||
+        instr.value.kind === "ObjectMethod"
+      ) {
         for (const operand of instr.value.loweredFunc.dependencies) {
           functionDependencies.add(operand.identifier.id);
         }

compiler/packages/babel-plugin-react-forget/src/Optimization/DeadCodeElimination.ts

@@ -234,6 +234,7 @@ function pruneableValue(value: InstructionValue, state: State): boolean {
     case "ArrayExpression":
     case "BinaryExpression":
     case "ComputedLoad":
+    case "ObjectMethod":
     case "FunctionExpression":
     case "LoadLocal":
     case "JsxExpression":

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

@@ -16,6 +16,7 @@ import {
   IdentifierId,
   InstructionKind,
   JsxAttribute,
+  ObjectMethod,
   ObjectPropertyKey,
   Pattern,
   Place,
@@ -118,6 +119,8 @@ class Context {
   #declarations: Set<IdentifierId> = new Set();
   temp: Temporaries = new Map();
   errors: CompilerError = new CompilerError();
+  objectMethods: Map<IdentifierId, ObjectMethod> = new Map();
+
   constructor(env: Environment, fnName: string) {
     this.env = env;
     this.fnName = fnName;
@@ -707,6 +710,14 @@ function codegenInstructionNullable(
     }
   } else if (instr.value.kind === "Debugger") {
     return t.debuggerStatement();
+  } else if (instr.value.kind === "ObjectMethod") {
+    CompilerError.invariant(instr.lvalue, {
+      reason: "Expected object methods to have a temp lvalue",
+      loc: null,
+      suggestions: null,
+    });
+    cx.objectMethods.set(instr.lvalue.identifier.id, instr.value);
+    return null;
   } else {
     const value = codegenInstructionValue(cx, instr.value);
     const statement = codegenInstruction(cx, instr, value);
@@ -969,15 +980,54 @@ function codegenInstructionValue(
       for (const property of instrValue.properties) {
         if (property.kind === "ObjectProperty") {
           const key = codegenObjectPropertyKey(property.key);
-          const value = codegenPlace(cx, property.place);
-          properties.push(
-            t.objectProperty(
-              key,
-              value,
-              false,
-              value.type === "Identifier" && value.name === property.key.name
-            )
-          );
+
+          switch (property.type) {
+            case "property": {
+              const value = codegenPlace(cx, property.place);
+              properties.push(
+                t.objectProperty(
+                  key,
+                  value,
+                  false,
+                  value.type === "Identifier" &&
+                    value.name === property.key.name
+                )
+              );
+              break;
+            }
+            case "method": {
+              const method = cx.objectMethods.get(property.place.identifier.id);
+              CompilerError.invariant(method, {
+                reason: "Expected ObjectMethod instruction",
+                loc: null,
+                suggestions: null,
+              });
+              const loweredFunc = method.loweredFunc;
+              const reactiveFunction = buildReactiveFunction(loweredFunc.func);
+              pruneUnusedLabels(reactiveFunction);
+              pruneUnusedLValues(reactiveFunction);
+              renameVariables(reactiveFunction);
+              const fn = codegenReactiveFunction(reactiveFunction).unwrap();
+
+              properties.push(
+                t.objectMethod(
+                  "method",
+                  key,
+                  fn.params,
+                  fn.body,
+                  false,
+                  fn.generator,
+                  fn.async
+                )
+              );
+              break;
+            }
+            default:
+              assertExhaustive(
+                property.type,
+                `Unexpected property type: ${property.type}`
+              );
+          }
         } else {
           properties.push(t.spreadElement(codegenPlace(cx, property.place)));
         }
@@ -1296,6 +1346,7 @@ function codegenInstructionValue(
     case "DeclareContext":
     case "Destructure":
     case "StoreLocal":
+    case "ObjectMethod":
     case "StoreContext": {
       CompilerError.invariant(false, {
         reason: `Unexpected ${instrValue.kind} in codegenInstructionValue`,

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

@@ -264,6 +264,7 @@ function mayAllocate(env: Environment, instruction: Instruction): boolean {
     case "NewExpression":
     case "ObjectExpression":
     case "UnsupportedNode":
+    case "ObjectMethod":
     case "FunctionExpression":
     case "TaggedTemplateExpression": {
       return true;

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

@@ -661,6 +661,7 @@ function computeMemoizationInputs(
       };
     }
     case "RegExpLiteral":
+    case "ObjectMethod":
     case "FunctionExpression":
     case "TaggedTemplateExpression":
     case "ArrayExpression":

compiler/packages/babel-plugin-react-forget/src/SSA/EliminateRedundantPhi.ts

@@ -106,7 +106,10 @@ export function eliminateRedundantPhi(
           rewritePlace(place, rewrites);
         }
 
-        if (instr.value.kind === "FunctionExpression") {
+        if (
+          instr.value.kind === "FunctionExpression" ||
+          instr.value.kind === "ObjectMethod"
+        ) {
           const { context } = instr.value.loweredFunc.func;
           for (const place of context) {
             rewritePlace(place, rewrites);

compiler/packages/babel-plugin-react-forget/src/SSA/EnterSSA.ts

@@ -275,7 +275,10 @@ function enterSSAImpl(
       mapInstructionOperands(instr, (place) => builder.getPlace(place));
       mapInstructionLValues(instr, (lvalue) => builder.definePlace(lvalue));
 
-      if (instr.value.kind === "FunctionExpression") {
+      if (
+        instr.value.kind === "FunctionExpression" ||
+        instr.value.kind === "ObjectMethod"
+      ) {
         const loweredFunc = instr.value.loweredFunc.func;
         const entry = loweredFunc.body.blocks.get(loweredFunc.body.entry)!;
         CompilerError.invariant(entry.preds.size === 0, {

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

@@ -255,6 +255,7 @@ function* generateInstructionTypes(
       break;
     }
 
+    case "ObjectMethod":
     case "FunctionExpression": {
       yield* generate(value.loweredFunc.func);
       break;

compiler/packages/babel-plugin-react-forget/src/Validation/ValidateFrozenLambdas.ts

@@ -15,6 +15,7 @@ import {
   FunctionExpression,
   HIRFunction,
   IdentifierId,
+  ObjectMethod,
   Place,
   isRefValueType,
   isUseRefType,
@@ -66,6 +67,7 @@ export function validateFrozenLambdas(fn: HIRFunction): void {
     }
     for (const instr of block.instructions) {
       switch (instr.value.kind) {
+        case "ObjectMethod":
         case "FunctionExpression": {
           if (
             instr.value.loweredFunc.dependencies.some(
@@ -119,7 +121,7 @@ export function validateFrozenLambdas(fn: HIRFunction): void {
 }
 
 class State {
-  lambdas: Map<IdentifierId, FunctionExpression> = new Map();
+  lambdas: Map<IdentifierId, FunctionExpression | ObjectMethod> = new Map();
   temporaries: Map<IdentifierId, IdentifierId> = new Map();
 }
 
@@ -132,9 +134,13 @@ function validateOperand(
       state.temporaries.get(operand.identifier.id) ?? operand.identifier.id;
     const lambda = state.lambdas.get(operandId);
     if (lambda !== undefined) {
-      // TODO: these seem to always be null, we should try to preserve original names from source
+      // TODO: these seem to always be null, we should try to preserve original
+      // names from source
+      // TODO: figure out how to print object methods as they don't have names
       const description =
-        lambda.name !== null && operand.identifier.name !== null
+        lambda.kind === "FunctionExpression" &&
+        lambda.name !== null &&
+        operand.identifier.name !== null
           ? `\`${lambda.name}\` is a function that may mutate \`${operand.identifier.name}\`. If you must mutate \`${operand.identifier.name}\` try using a React API like useState and use its setter function instead`
           : null;
       return new CompilerErrorDetail({

compiler/packages/babel-plugin-react-forget/src/Validation/ValidateNoRefAccesInRender.ts

@@ -59,6 +59,7 @@ export function validateNoRefAccessInRender(fn: HIRFunction): void {
           }
           break;
         }
+        case "ObjectMethod":
         case "FunctionExpression": {
           // functions are allowed to capture refs, so long as the function is not called
           // during render. see AnalyzeFunctions for how we ensure that functions which

compiler/packages/babel-plugin-react-forget/src/Validation/ValidateNoSetStateInRender.ts

@@ -36,6 +36,7 @@ export function validateNoSetStateInRender(
     if (unconditionalBlocks.has(block.id)) {
       for (const instr of block.instructions) {
         switch (instr.value.kind) {
+          case "ObjectMethod":
           case "FunctionExpression": {
             /**
              * TODO: setState's return value is considered Frozen, so the lambda's mutable range

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.todo-kitchensink.expect.md

@@ -83,8 +83,6 @@ let moduleLocal = false;
 
 [ReactForget] Todo: (BuildHIR::lowerStatement) Handle ClassDeclaration statements (5:10)
 
-[ReactForget] Todo: (BuildHIR::lowerExpression) Handle ObjectMethod properties in ObjectExpression (12:12)
-
 [ReactForget] Todo: (BuildHIR::lowerStatement) Handle non-variable initialization in ForStatement (20:22)
 
 [ReactForget] Todo: (BuildHIR::lowerStatement) Handle non-variable initialization in ForStatement (23:25)

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

@@ -0,0 +1,74 @@
+
+## Input
+
+```javascript
+function Component({ a, b }) {
+  return {
+    x: function () {
+      return [a];
+    },
+    y() {
+      return [b];
+    },
+  };
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ x: 1 }, { a: 2 }, { b: 2 }],
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Component(t16) {
+  const $ = useMemoCache(7);
+  const { a, b } = t16;
+  const c_0 = $[0] !== a;
+  let t0;
+  if (c_0) {
+    t0 = function () {
+      return [a];
+    };
+    $[0] = a;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  const c_2 = $[2] !== b;
+  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 = {
+      x: t0,
+      y() {
+        return [b];
+      },
+    };
+    $[4] = t0;
+    $[5] = t1;
+    $[6] = t2;
+  } else {
+    t2 = $[6];
+  }
+  return t2;
+}
+
+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.js

@@ -0,0 +1,15 @@
+function Component({ a, b }) {
+  return {
+    x: function () {
+      return [a];
+    },
+    y() {
+      return [b];
+    },
+  };
+}
+
+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-2.expect.md

@@ -0,0 +1,83 @@
+
+## Input
+
+```javascript
+function Component({ a, b, c }) {
+  return {
+    x: [a],
+    y() {
+      return [b];
+    },
+    z: { c },
+  };
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ x: 1 }, { a: 2 }, { b: 2 }],
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Component(t16) {
+  const $ = useMemoCache(10);
+  const { a, b, c } = t16;
+  const c_0 = $[0] !== a;
+  let t0;
+  if (c_0) {
+    t0 = [a];
+    $[0] = a;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  const c_2 = $[2] !== b;
+  let t1;
+  if (c_2) {
+    $[2] = b;
+    $[3] = t1;
+  } else {
+    t1 = $[3];
+  }
+  const c_4 = $[4] !== c;
+  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 = {
+      x: t0,
+      y() {
+        return [b];
+      },
+      z: t2,
+    };
+    $[6] = t0;
+    $[7] = t1;
+    $[8] = t2;
+    $[9] = t3;
+  } else {
+    t3 = $[9];
+  }
+  return t3;
+}
+
+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-2.js

@@ -0,0 +1,14 @@
+function Component({ a, b, c }) {
+  return {
+    x: [a],
+    y() {
+      return [b];
+    },
+    z: { c },
+  };
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ x: 1 }, { a: 2 }, { b: 2 }],
+};