[compiler] Allow reordering more types of instructions

Updates InstructionReordering (still off by default) to allow reordering more types of instructions:

* ArrayExpression and ObjectExpression can be reordered if they have no items
* PropertyLoad can be reordered if the object is not mutable at the given instruction, and if the object is single-use (this avoid reordering in cases where the propertyload is part of a methodcall)
* LoadLocal can be reordered if the temporary is used once (ie not as part of a methodcall) and if the value being loaded is not subsequently reassigned
* StoreLocal can be reordered if the variable being assigned is only used once, and the rvalue is not later reassigned.

We can further relax these rules but already it improves compilation output quite a bit, per the fixtures.

ghstack-source-id: 896807cb0a
Pull Request resolved: https://github.com/facebook/react/pull/30653

compiler/packages/babel-plugin-react-compiler/src/Optimization/InstructionReordering.ts

@@ -5,9 +5,11 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+import prettyFormat from 'pretty-format';
 import {CompilerError} from '..';
 import {
   BasicBlock,
+  DeclarationId,
   Environment,
   GeneratedSource,
   HIRFunction,
@@ -19,7 +21,7 @@ import {
   makeInstructionId,
   markInstructionIds,
 } from '../HIR';
-import {printInstruction} from '../HIR/PrintHIR';
+import {printInstruction, printPlace} from '../HIR/PrintHIR';
 import {
   eachInstructionLValue,
   eachInstructionValueLValue,
@@ -101,7 +103,7 @@ type References = {
   singleUseIdentifiers: SingleUseIdentifiers;
   lastAssignments: LastAssignments;
 };
-type LastAssignments = Map<string, InstructionId>;
+type LastAssignments = Map<DeclarationId, InstructionId>;
 type SingleUseIdentifiers = Set<IdentifierId>;
 enum ReferenceKind {
   Read,
@@ -110,28 +112,34 @@ enum ReferenceKind {
 function findReferencedRangeOfTemporaries(fn: HIRFunction): References {
   const singleUseIdentifiers = new Map<IdentifierId, number>();
   const lastAssignments: LastAssignments = new Map();
+
+  for (const param of fn.params) {
+    const place = param.kind === 'Identifier' ? param : param.place;
+    lastAssignments.set(place.identifier.declarationId, makeInstructionId(0));
+  }
+
   function reference(
     instr: InstructionId,
     place: Place,
     kind: ReferenceKind,
   ): void {
-    if (
-      place.identifier.name !== null &&
-      place.identifier.name.kind === 'named'
-    ) {
-      if (kind === ReferenceKind.Write) {
-        const name = place.identifier.name.value;
-        const previous = lastAssignments.get(name);
-        if (previous === undefined) {
-          lastAssignments.set(name, instr);
-        } else {
-          lastAssignments.set(
-            name,
-            makeInstructionId(Math.max(previous, instr)),
-          );
-        }
+    if (kind === ReferenceKind.Write) {
+      const declarationId = place.identifier.declarationId;
+      const previous = lastAssignments.get(declarationId);
+      if (previous === undefined) {
+        lastAssignments.set(declarationId, instr);
+      } else {
+        lastAssignments.set(
+          declarationId,
+          makeInstructionId(Math.max(previous, instr)),
+        );
       }
-      return;
+      CompilerError.invariant(!singleUseIdentifiers.has(place.identifier.id), {
+        reason: `Unexpected existing declaration for identifier`,
+        description: `[${instr}] ${printPlace(place)}`,
+        loc: place.loc,
+      });
+      singleUseIdentifiers.set(place.identifier.id, 0);
     } else if (kind === ReferenceKind.Read) {
       const previousCount = singleUseIdentifiers.get(place.identifier.id) ?? 0;
       singleUseIdentifiers.set(place.identifier.id, previousCount + 1);
@@ -139,7 +147,7 @@ function findReferencedRangeOfTemporaries(fn: HIRFunction): References {
   }
   for (const [, block] of fn.body.blocks) {
     for (const instr of block.instructions) {
-      for (const operand of eachInstructionValueLValue(instr.value)) {
+      for (const operand of eachInstructionValueOperand(instr.value)) {
         reference(instr.id, operand, ReferenceKind.Read);
       }
       for (const lvalue of eachInstructionLValue(instr)) {
@@ -153,7 +161,7 @@ function findReferencedRangeOfTemporaries(fn: HIRFunction): References {
   return {
     singleUseIdentifiers: new Set(
       [...singleUseIdentifiers]
-        .filter(([, count]) => count === 1)
+        .filter(([, count]) => count <= 1)
         .map(([id]) => id),
     ),
     lastAssignments,
@@ -167,7 +175,6 @@ function reorderBlock(
   references: References,
 ): void {
   const locals: Nodes = new Map();
-  const named: Map<string, IdentifierId> = new Map();
   let previous: IdentifierId | null = null;
   for (const instr of block.instructions) {
     const {lvalue, value} = instr;
@@ -198,15 +205,8 @@ function reorderBlock(
      * Establish dependencies on operands
      */
     for (const operand of eachInstructionValueOperand(value)) {
-      const {name, id} = operand.identifier;
-      if (name !== null && name.kind === 'named') {
-        // Serialize all accesses to named variables
-        const previous = named.get(name.value);
-        if (previous !== undefined) {
-          node.dependencies.add(previous);
-        }
-        named.set(name.value, lvalue.identifier.id);
-      } else if (locals.has(id) || shared.has(id)) {
+      const id = operand.identifier.id;
+      if (locals.has(id) || shared.has(id)) {
         node.dependencies.add(id);
       }
     }
@@ -228,22 +228,12 @@ function reorderBlock(
           }) as Node,
       );
       lvalueNode.dependencies.add(lvalue.identifier.id);
-      const name = lvalueOperand.identifier.name;
-      if (name !== null && name.kind === 'named') {
-        const previous = named.get(name.value);
-        if (previous !== undefined) {
-          node.dependencies.add(previous);
-        }
-        named.set(name.value, lvalue.identifier.id);
-      }
     }
   }
 
   const nextInstructions: Array<Instruction> = [];
   const seen = new Set<IdentifierId>();
 
-  DEBUG && console.log(`bb${block.id}`);
-
   /**
    * The ideal order for emitting instructions may change the final instruction,
    * but value blocks have special semantics for the final instruction of a block -
@@ -464,14 +454,24 @@ function emit(
 }
 
 enum Reorderability {
-  Reorderable,
-  Nonreorderable,
+  Reorderable = 'Reorderable',
+  Nonreorderable = 'Nonreorderable',
 }
 function getReorderability(
   instr: Instruction,
   references: References,
 ): Reorderability {
   switch (instr.value.kind) {
+    case 'ArrayExpression': {
+      return instr.value.elements.length === 0
+        ? Reorderability.Reorderable
+        : Reorderability.Nonreorderable;
+    }
+    case 'ObjectExpression': {
+      return instr.value.properties.length === 0
+        ? Reorderability.Reorderable
+        : Reorderability.Nonreorderable;
+    }
     case 'JsxExpression':
     case 'JsxFragment':
     case 'JSXText':
@@ -482,17 +482,52 @@ function getReorderability(
     case 'UnaryExpression': {
       return Reorderability.Reorderable;
     }
+    case 'PropertyLoad': {
+      const object = instr.value.object.identifier;
+      if (
+        object.scope !== null &&
+        instr.id >= object.scope.range.start &&
+        instr.id < object.scope.range.end
+      ) {
+        return Reorderability.Nonreorderable;
+      }
+      if (!references.singleUseIdentifiers.has(object.id)) {
+        return Reorderability.Nonreorderable;
+      }
+      return Reorderability.Reorderable;
+    }
     case 'LoadLocal': {
-      const name = instr.value.place.identifier.name;
-      if (name !== null && name.kind === 'named') {
-        const lastAssignment = references.lastAssignments.get(name.value);
-        if (
-          lastAssignment !== undefined &&
-          lastAssignment < instr.id &&
-          references.singleUseIdentifiers.has(instr.lvalue.identifier.id)
-        ) {
-          return Reorderability.Reorderable;
-        }
+      const lastAssignment = references.lastAssignments.get(
+        instr.value.place.identifier.declarationId,
+      );
+      if (
+        lastAssignment !== undefined &&
+        lastAssignment < instr.id &&
+        references.singleUseIdentifiers.has(instr.lvalue.identifier.id)
+      ) {
+        return Reorderability.Reorderable;
+      }
+      return Reorderability.Nonreorderable;
+    }
+    case 'StoreLocal': {
+      const lastAssignmentOfRvalue = references.lastAssignments.get(
+        instr.value.value.identifier.declarationId,
+      );
+      const isRValueNoLongerReassigned =
+        lastAssignmentOfRvalue === undefined ||
+        lastAssignmentOfRvalue < instr.id;
+      const isInstructionLvalueSingleUse = references.singleUseIdentifiers.has(
+        instr.lvalue.identifier.id,
+      );
+      const isAssignedVariableSingleUse = references.singleUseIdentifiers.has(
+        instr.value.lvalue.place.identifier.id,
+      );
+      if (
+        isRValueNoLongerReassigned &&
+        isInstructionLvalueSingleUse &&
+        isAssignedVariableSingleUse
+      ) {
+        return Reorderability.Reorderable;
       }
       return Reorderability.Nonreorderable;
     }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/instruction-reordering-store-local.expect.md

@@ -0,0 +1,72 @@
+
+## Input
+
+```javascript
+// @enableInstructionReordering
+
+function Component(props) {
+  const x = [];
+  const a = props.a;
+  x.push(props.a);
+  const b = a.b;
+  x.push(props.b);
+  const c = b.c;
+  x.push(props.c);
+  const d = c.d;
+  x.push(props.d);
+  return [d, x];
+}
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enableInstructionReordering
+
+function Component(props) {
+  const $ = _c(8);
+  let x;
+  if (
+    $[0] !== props.a ||
+    $[1] !== props.b ||
+    $[2] !== props.c ||
+    $[3] !== props.d
+  ) {
+    x = [];
+
+    x.push(props.a);
+
+    x.push(props.b);
+
+    x.push(props.c);
+
+    x.push(props.d);
+    $[0] = props.a;
+    $[1] = props.b;
+    $[2] = props.c;
+    $[3] = props.d;
+    $[4] = x;
+  } else {
+    x = $[4];
+  }
+  const a = props.a;
+  const b = a.b;
+  const c = b.c;
+  const d = c.d;
+  let t0;
+  if ($[5] !== d || $[6] !== x) {
+    t0 = [d, x];
+    $[5] = d;
+    $[6] = x;
+    $[7] = t0;
+  } else {
+    t0 = $[7];
+  }
+  return t0;
+}
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/instruction-reordering-store-local.js

@@ -0,0 +1,14 @@
+// @enableInstructionReordering
+
+function Component(props) {
+  const x = [];
+  const a = props.a;
+  x.push(props.a);
+  const b = a.b;
+  x.push(props.b);
+  const c = b.c;
+  x.push(props.c);
+  const d = c.d;
+  x.push(props.d);
+  return [d, x];
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-scopes-callback.expect.md

@@ -27,39 +27,33 @@ import { c as _c } from "react/compiler-runtime"; // @enableInstructionReorderin
 import { useState } from "react";
 
 function Component() {
-  const $ = _c(4);
+  const $ = _c(3);
   const [state, setState] = useState(0);
   let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t0 = () => {
+    const onClick = () => {
       setState((s) => s + 1);
     };
+
+    t0 = <button onClick={onClick}>Increment</button>;
     $[0] = t0;
   } else {
     t0 = $[0];
   }
-  const onClick = t0;
   let t1;
-  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
-    t1 = <button onClick={onClick}>Increment</button>;
-    $[1] = t1;
-  } else {
-    t1 = $[1];
-  }
-  let t2;
-  if ($[2] !== state) {
-    t2 = (
+  if ($[1] !== state) {
+    t1 = (
       <>
         <span>Count: {state}</span>
-        {t1}
+        {t0}
       </>
     );
-    $[2] = state;
-    $[3] = t2;
+    $[1] = state;
+    $[2] = t1;
   } else {
-    t2 = $[3];
+    t1 = $[2];
   }
-  return t2;
+  return t1;
 }
 
 ```