[compiler][ez] Add shape for global Object.keys

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

@@ -87,6 +87,21 @@ const UNTYPED_GLOBALS: Set<string> = new Set([
 ]);
 
 const TYPED_GLOBALS: Array<[string, BuiltInType]> = [
+  [
+    'Object',
+    addObject(DEFAULT_SHAPES, 'Object', [
+      [
+        'keys',
+        addFunction(DEFAULT_SHAPES, [], {
+          positionalParams: [Effect.Read],
+          restParam: null,
+          returnType: {kind: 'Object', shapeId: BuiltInArrayId},
+          calleeEffect: Effect.Read,
+          returnValueKind: ValueKind.Mutable,
+        }),
+      ],
+    ]),
+  ],
   [
     'Array',
     addObject(DEFAULT_SHAPES, 'Array', [

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

@@ -840,6 +840,11 @@ export type LoadLocal = {
   place: Place;
   loc: SourceLocation;
 };
+export type LoadContext = {
+  kind: 'LoadContext';
+  place: Place;
+  loc: SourceLocation;
+};
 
 /*
  * The value of a given instruction. Note that values are not recursive: complex
@@ -852,11 +857,7 @@ export type LoadLocal = {
 
 export type InstructionValue =
   | LoadLocal
-  | {
-      kind: 'LoadContext';
-      place: Place;
-      loc: SourceLocation;
-    }
+  | LoadContext
   | {
       kind: 'DeclareLocal';
       lvalue: LValue;

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

@@ -17,6 +17,11 @@ import {
   areEqualPaths,
   IdentifierId,
   Terminal,
+  InstructionValue,
+  LoadContext,
+  TInstruction,
+  FunctionExpression,
+  ObjectMethod,
 } from './HIR';
 import {
   collectHoistablePropertyLoads,
@@ -223,11 +228,25 @@ export function collectTemporariesSidemap(
     fn,
     usedOutsideDeclaringScope,
     temporaries,
-    false,
+    null,
   );
   return temporaries;
 }
 
+function isLoadContextMutable(
+  instrValue: InstructionValue,
+  id: InstructionId,
+): instrValue is LoadContext {
+  if (instrValue.kind === 'LoadContext') {
+    CompilerError.invariant(instrValue.place.identifier.scope != null, {
+      reason:
+        '[PropagateScopeDependencies] Expected all context variables to be assigned a scope',
+      loc: instrValue.loc,
+    });
+    return id >= instrValue.place.identifier.scope.range.end;
+  }
+  return false;
+}
 /**
  * Recursive collect a sidemap of all `LoadLocal` and `PropertyLoads` with a
  * function and all nested functions.
@@ -239,17 +258,21 @@ function collectTemporariesSidemapImpl(
   fn: HIRFunction,
   usedOutsideDeclaringScope: ReadonlySet<DeclarationId>,
   temporaries: Map<IdentifierId, ReactiveScopeDependency>,
-  isInnerFn: boolean,
+  innerFnContext: {instrId: InstructionId} | null,
 ): void {
   for (const [_, block] of fn.body.blocks) {
-    for (const instr of block.instructions) {
-      const {value, lvalue} = instr;
+    for (const {value, lvalue, id: origInstrId} of block.instructions) {
+      const instrId =
+        innerFnContext != null ? innerFnContext.instrId : origInstrId;
       const usedOutside = usedOutsideDeclaringScope.has(
         lvalue.identifier.declarationId,
       );
 
       if (value.kind === 'PropertyLoad' && !usedOutside) {
-        if (!isInnerFn || temporaries.has(value.object.identifier.id)) {
+        if (
+          innerFnContext == null ||
+          temporaries.has(value.object.identifier.id)
+        ) {
           /**
            * All dependencies of a inner / nested function must have a base
            * identifier from the outermost component / hook. This is because the
@@ -265,13 +288,13 @@ function collectTemporariesSidemapImpl(
           temporaries.set(lvalue.identifier.id, property);
         }
       } else if (
-        value.kind === 'LoadLocal' &&
+        (value.kind === 'LoadLocal' || isLoadContextMutable(value, instrId)) &&
         lvalue.identifier.name == null &&
         value.place.identifier.name !== null &&
         !usedOutside
       ) {
         if (
-          !isInnerFn ||
+          innerFnContext == null ||
           fn.context.some(
             context => context.identifier.id === value.place.identifier.id,
           )
@@ -289,7 +312,7 @@ function collectTemporariesSidemapImpl(
           value.loweredFunc.func,
           usedOutsideDeclaringScope,
           temporaries,
-          true,
+          innerFnContext ?? {instrId},
         );
       }
     }
@@ -358,19 +381,22 @@ class Context {
 
   #temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>;
   #temporariesUsedOutsideScope: ReadonlySet<DeclarationId>;
+  #processedInstrsInOptional: ReadonlySet<Instruction | Terminal>;
 
   /**
    * Tracks the traversal state. See Context.declare for explanation of why this
    * is needed.
    */
-  inInnerFn: boolean = false;
+  #innerFnContext: {outerInstrId: InstructionId} | null = null;
 
   constructor(
     temporariesUsedOutsideScope: ReadonlySet<DeclarationId>,
     temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
+    processedInstrsInOptional: ReadonlySet<Instruction | Terminal>,
   ) {
     this.#temporariesUsedOutsideScope = temporariesUsedOutsideScope;
     this.#temporaries = temporaries;
+    this.#processedInstrsInOptional = processedInstrsInOptional;
   }
 
   enterScope(scope: ReactiveScope): void {
@@ -431,7 +457,7 @@ class Context {
    *     by root identifier mutable ranges).
    */
   declare(identifier: Identifier, decl: Decl): void {
-    if (this.inInnerFn) return;
+    if (this.#innerFnContext != null) return;
     if (!this.#declarations.has(identifier.declarationId)) {
       this.#declarations.set(identifier.declarationId, decl);
     }
@@ -574,22 +600,52 @@ class Context {
       currentScope.reassignments.add(place.identifier);
     }
   }
+  enterInnerFn<T>(
+    innerFn: TInstruction<FunctionExpression> | TInstruction<ObjectMethod>,
+    cb: () => T,
+  ): T {
+    const prevContext = this.#innerFnContext;
+    this.#innerFnContext = this.#innerFnContext ?? {outerInstrId: innerFn.id};
+    const result = cb();
+    this.#innerFnContext = prevContext;
+    return result;
+  }
+
+  /**
+   * Skip dependencies that are subexpressions of other dependencies. e.g. if a
+   * dependency is tracked in the temporaries sidemap, it can be added at
+   * site-of-use
+   */
+  isDeferredDependency(
+    instr:
+      | {kind: HIRValue.Instruction; value: Instruction}
+      | {kind: HIRValue.Terminal; value: Terminal},
+  ): boolean {
+    return (
+      this.#processedInstrsInOptional.has(instr.value) ||
+      (instr.kind === HIRValue.Instruction &&
+        this.#temporaries.has(instr.value.lvalue.identifier.id))
+    );
+  }
+}
+enum HIRValue {
+  Instruction = 1,
+  Terminal,
 }
 
 function handleInstruction(instr: Instruction, context: Context): void {
   const {id, value, lvalue} = instr;
-  if (value.kind === 'LoadLocal') {
-    if (
-      value.place.identifier.name === null ||
-      lvalue.identifier.name !== null ||
-      context.isUsedOutsideDeclaringScope(lvalue)
-    ) {
-      context.visitOperand(value.place);
-    }
-  } else if (value.kind === 'PropertyLoad') {
-    if (context.isUsedOutsideDeclaringScope(lvalue)) {
-      context.visitProperty(value.object, value.property, false);
-    }
+  context.declare(lvalue.identifier, {
+    id,
+    scope: context.currentScope,
+  });
+  if (
+    context.isDeferredDependency({kind: HIRValue.Instruction, value: instr})
+  ) {
+    return;
+  }
+  if (value.kind === 'PropertyLoad') {
+    context.visitProperty(value.object, value.property, false);
   } else if (value.kind === 'StoreLocal') {
     context.visitOperand(value.value);
     if (value.lvalue.kind === InstructionKind.Reassign) {
@@ -632,11 +688,6 @@ function handleInstruction(instr: Instruction, context: Context): void {
       context.visitOperand(operand);
     }
   }
-
-  context.declare(lvalue.identifier, {
-    id,
-    scope: context.currentScope,
-  });
 }
 
 function collectDependencies(
@@ -645,7 +696,11 @@ function collectDependencies(
   temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
   processedInstrsInOptional: ReadonlySet<Instruction | Terminal>,
 ): Map<ReactiveScope, Array<ReactiveScopeDependency>> {
-  const context = new Context(usedOutsideDeclaringScope, temporaries);
+  const context = new Context(
+    usedOutsideDeclaringScope,
+    temporaries,
+    processedInstrsInOptional,
+  );
 
   for (const param of fn.params) {
     if (param.kind === 'Identifier') {
@@ -694,16 +749,26 @@ function collectDependencies(
           /**
            * Recursively visit the inner function to extract dependencies there
            */
-          const wasInInnerFn = context.inInnerFn;
-          context.inInnerFn = true;
-          handleFunction(instr.value.loweredFunc.func);
-          context.inInnerFn = wasInInnerFn;
-        } else if (!processedInstrsInOptional.has(instr)) {
+          const innerFn = instr.value.loweredFunc.func;
+          context.enterInnerFn(
+            instr as
+              | TInstruction<FunctionExpression>
+              | TInstruction<ObjectMethod>,
+            () => {
+              handleFunction(innerFn);
+            },
+          );
+        } else {
           handleInstruction(instr, context);
         }
       }
 
-      if (!processedInstrsInOptional.has(block.terminal)) {
+      if (
+        !context.isDeferredDependency({
+          kind: HIRValue.Terminal,
+          value: block.terminal,
+        })
+      ) {
         for (const place of eachTerminalOperand(block.terminal)) {
           context.visitOperand(place);
         }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-functiondecl-hoisting.expect.md

@@ -58,18 +58,16 @@ function Foo(t0) {
     bar = $[1];
     result = $[2];
   }
-
-  const t1 = bar;
-  let t2;
-  if ($[3] !== result || $[4] !== t1) {
-    t2 = <Stringify result={result} fn={t1} shouldInvokeFns={true} />;
-    $[3] = result;
-    $[4] = t1;
-    $[5] = t2;
+  let t1;
+  if ($[3] !== bar || $[4] !== result) {
+    t1 = <Stringify result={result} fn={bar} shouldInvokeFns={true} />;
+    $[3] = bar;
+    $[4] = result;
+    $[5] = t1;
   } else {
-    t2 = $[5];
+    t1 = $[5];
   }
-  return t2;
+  return t1;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/destructure-array-assignment-to-context-var.expect.md

@@ -43,16 +43,15 @@ function Component(props) {
   } else {
     x = $[1];
   }
-  const t0 = x;
-  let t1;
-  if ($[2] !== t0) {
-    t1 = { x: t0 };
-    $[2] = t0;
-    $[3] = t1;
+  let t0;
+  if ($[2] !== x) {
+    t0 = { x };
+    $[2] = x;
+    $[3] = t0;
   } else {
-    t1 = $[3];
+    t0 = $[3];
   }
-  return t1;
+  return t0;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/destructure-array-declaration-to-context-var.expect.md

@@ -42,16 +42,15 @@ function Component(props) {
   } else {
     x = $[1];
   }
-  const t0 = x;
-  let t1;
-  if ($[2] !== t0) {
-    t1 = <div>{t0}</div>;
-    $[2] = t0;
-    $[3] = t1;
+  let t0;
+  if ($[2] !== x) {
+    t0 = <div>{x}</div>;
+    $[2] = x;
+    $[3] = t0;
   } else {
-    t1 = $[3];
+    t0 = $[3];
   }
-  return t1;
+  return t0;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/destructure-object-assignment-to-context-var.expect.md

@@ -43,16 +43,15 @@ function Component(props) {
   } else {
     x = $[1];
   }
-  const t0 = x;
-  let t1;
-  if ($[2] !== t0) {
-    t1 = { x: t0 };
-    $[2] = t0;
-    $[3] = t1;
+  let t0;
+  if ($[2] !== x) {
+    t0 = { x };
+    $[2] = x;
+    $[3] = t0;
   } else {
-    t1 = $[3];
+    t0 = $[3];
   }
-  return t1;
+  return t0;
 }
 
 export const FIXTURE_ENTRYPOINT = {