[compiler][rewrite] Represent scope dependencies with value blocks

(needs cleanup)

- Scopes no longer store a flat list of their dependencies. Instead:
  - Scope terminals are effectively a `goto` for scope dependency instructions (represented as value blocks that terminate with a `goto scopeBlock` for HIR and a series of ReactiveInstructions for ReactiveIR)
  - Scopes themselves store `dependencies: Array<Place>`, which refer to temporaries written to by scope dependency instructions

Next steps:
- new pass to dedupe scope dependency instructions after all dependency and scope pruning passes, effectively 'hoisting' dependencies out
- more complex dependencies (unary ops like `Boolean` or `Not`, binary ops like `!==` or logical operators)

Author: mofeiZ

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

@@ -73,7 +73,10 @@ export function lower(
   // the outermost function being compiled, in case lower() is called recursively (for lambdas)
   parent: NodePath<t.Function> | null = null,
 ): Result<HIRFunction, CompilerError> {
-  const builder = new HIRBuilder(env, parent ?? func, bindings, capturedRefs);
+  const builder = new HIRBuilder(env, parent ?? func, {
+    bindings,
+    context: capturedRefs,
+  });
   const context: HIRFunction['context'] = [];
 
   for (const ref of capturedRefs ?? []) {
@@ -3414,7 +3417,7 @@ function lowerExpressionToTemporary(
   return lowerValueToTemporary(builder, value);
 }
 
-function lowerValueToTemporary(
+export function lowerValueToTemporary(
   builder: HIRBuilder,
   value: InstructionValue,
 ): Place {

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

@@ -184,6 +184,7 @@ type TerminalRewriteInfo =
   | {
       kind: 'StartScope';
       blockId: BlockId;
+      dependencyId: BlockId;
       fallthroughId: BlockId;
       instrId: InstructionId;
       scope: ReactiveScope;
@@ -208,12 +209,14 @@ function pushStartScopeTerminal(
   scope: ReactiveScope,
   context: ScopeTraversalContext,
 ): void {
+  const dependencyId = context.env.nextBlockId;
   const blockId = context.env.nextBlockId;
   const fallthroughId = context.env.nextBlockId;
   context.rewrites.push({
     kind: 'StartScope',
     blockId,
     fallthroughId,
+    dependencyId,
     instrId: scope.range.start,
     scope,
   });
@@ -255,10 +258,13 @@ type RewriteContext = {
  *     instr1, instr2, instr3, instr4, [[ original terminal ]]
  * Rewritten:
  *   bb0:
- *     instr1, [[ scope start block=bb1]]
+ *     instr1, [[ scope start dependencies=bb1 block=bb2]]
  *   bb1:
- *     instr2, instr3, [[ scope end goto=bb2 ]]
+ *     [[ empty, filled in in PropagateScopeDependenciesHIR ]]
+ *      goto bb2
  *   bb2:
+ *     instr2, instr3, [[ scope end goto=bb3 ]]
+ *   bb3:
  *     instr4, [[ original terminal ]]
  */
 function handleRewrite(
@@ -272,6 +278,7 @@ function handleRewrite(
       ? {
           kind: 'scope',
           fallthrough: terminalInfo.fallthroughId,
+          dependencies: terminalInfo.dependencyId,
           block: terminalInfo.blockId,
           scope: terminalInfo.scope,
           id: terminalInfo.instrId,
@@ -298,7 +305,28 @@ function handleRewrite(
   context.nextPreds = new Set([currBlockId]);
   context.nextBlockId =
     terminalInfo.kind === 'StartScope'
-      ? terminalInfo.blockId
+      ? terminalInfo.dependencyId
       : terminalInfo.fallthroughId;
   context.instrSliceIdx = idx;
+
+  if (terminalInfo.kind === 'StartScope') {
+    const currBlockId = context.nextBlockId;
+    context.rewrites.push({
+      kind: context.source.kind,
+      id: currBlockId,
+      instructions: [],
+      preds: context.nextPreds,
+      // Only the first rewrite should reuse source block phis
+      phis: new Set(),
+      terminal: {
+        kind: 'goto',
+        variant: GotoVariant.Break,
+        block: terminal.block,
+        id: terminalInfo.instrId,
+        loc: GeneratedSource,
+      },
+    });
+    context.nextPreds = new Set([currBlockId]);
+    context.nextBlockId = terminalInfo.blockId;
+  }
 }

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

@@ -1,3 +1,4 @@
+import {react} from '@babel/types';
 import {CompilerError} from '../CompilerError';
 import {inRange} from '../ReactiveScopes/InferReactiveScopeVariables';
 import {printDependency} from '../ReactiveScopes/PrintReactiveFunction';
@@ -194,7 +195,10 @@ type PropertyPathNode =
 class PropertyPathRegistry {
   roots: Map<IdentifierId, RootNode> = new Map();
 
-  getOrCreateIdentifier(identifier: Identifier): PropertyPathNode {
+  getOrCreateIdentifier(
+    identifier: Identifier,
+    reactive: boolean,
+  ): PropertyPathNode {
     /**
      * Reads from a statically scoped variable are always safe in JS,
      * with the exception of TDZ (not addressed by this pass).
@@ -208,12 +212,18 @@ class PropertyPathRegistry {
         optionalProperties: new Map(),
         fullPath: {
           identifier,
+          reactive,
           path: [],
         },
         hasOptional: false,
         parent: null,
       };
       this.roots.set(identifier.id, rootNode);
+    } else {
+      CompilerError.invariant(reactive === rootNode.fullPath.reactive, {
+        reason: 'Inconsistent reactive flag',
+        loc: GeneratedSource,
+      });
     }
     return rootNode;
   }
@@ -231,6 +241,7 @@ class PropertyPathRegistry {
         parent: parent,
         fullPath: {
           identifier: parent.fullPath.identifier,
+          reactive: parent.fullPath.reactive,
           path: parent.fullPath.path.concat(entry),
         },
         hasOptional: parent.hasOptional || entry.optional,
@@ -246,7 +257,7 @@ class PropertyPathRegistry {
      * so all subpaths of a PropertyLoad should already exist
      * (e.g. a.b is added before a.b.c),
      */
-    let currNode = this.getOrCreateIdentifier(n.identifier);
+    let currNode = this.getOrCreateIdentifier(n.identifier, n.reactive);
     if (n.path.length === 0) {
       return currNode;
     }
@@ -268,10 +279,11 @@ function getMaybeNonNullInInstruction(
   instr: InstructionValue,
   context: CollectHoistablePropertyLoadsContext,
 ): PropertyPathNode | null {
-  let path = null;
+  let path: ReactiveScopeDependency | null = null;
   if (instr.kind === 'PropertyLoad') {
     path = context.temporaries.get(instr.object.identifier.id) ?? {
       identifier: instr.object.identifier,
+      reactive: instr.object.reactive,
       path: [],
     };
   } else if (instr.kind === 'Destructure') {
@@ -334,7 +346,7 @@ function collectNonNullsInBlocks(
   ) {
     const identifier = fn.params[0].identifier;
     knownNonNullIdentifiers.add(
-      context.registry.getOrCreateIdentifier(identifier),
+      context.registry.getOrCreateIdentifier(identifier, true),
     );
   }
   const nodes = new Map<BlockId, BlockInfo>();
@@ -565,9 +577,11 @@ function reduceMaybeOptionalChains(
     changed = false;
 
     for (const original of optionalChainNodes) {
-      let {identifier, path: origPath} = original.fullPath;
-      let currNode: PropertyPathNode =
-        registry.getOrCreateIdentifier(identifier);
+      let {identifier, path: origPath, reactive} = original.fullPath;
+      let currNode: PropertyPathNode = registry.getOrCreateIdentifier(
+        identifier,
+        reactive,
+      );
       for (let i = 0; i < origPath.length; i++) {
         const entry = origPath[i];
         // If the base is known to be non-null, replace with a non-optional load

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

@@ -106,7 +106,7 @@ export type OptionalChainSidemap = {
   hoistableObjects: ReadonlyMap<BlockId, ReactiveScopeDependency>;
 };
 
-type OptionalTraversalContext = {
+export type OptionalTraversalContext = {
   currFn: HIRFunction;
   blocks: ReadonlyMap<BlockId, BasicBlock>;
 
@@ -227,7 +227,7 @@ function matchOptionalTestBlock(
  * property loads. If any part of the optional chain is not hoistable, returns
  * null.
  */
-function traverseOptionalBlock(
+export function traverseOptionalBlock(
   optional: TBasicBlock<OptionalTerminal>,
   context: OptionalTraversalContext,
   outerAlternate: BlockId | null,
@@ -282,6 +282,7 @@ function traverseOptionalBlock(
     );
     baseObject = {
       identifier: maybeTest.instructions[0].value.place.identifier,
+      reactive: maybeTest.instructions[0].value.place.reactive,
       path,
     };
     test = maybeTest.terminal;
@@ -383,6 +384,7 @@ function traverseOptionalBlock(
   );
   const load = {
     identifier: baseObject.identifier,
+    reactive: baseObject.reactive,
     path: [
       ...baseObject.path,
       {

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

@@ -13,7 +13,8 @@ import {
   ReactiveScopeDependency,
 } from '../HIR';
 import {printIdentifier} from '../HIR/PrintHIR';
-import {ReactiveScopePropertyDependency} from '../ReactiveScopes/DeriveMinimalDependencies';
+
+export type ReactiveScopePropertyDependency = ReactiveScopeDependency;
 
 /**
  * Simpler fork of DeriveMinimalDependencies, see PropagateScopeDependenciesHIR
@@ -25,8 +26,9 @@ export class ReactiveScopeDependencyTreeHIR {
    * `identifier.path`, or `identifier?.path` is in this map, it is safe to
    * evaluate (non-optional) PropertyLoads from.
    */
-  #hoistableObjects: Map<Identifier, HoistableNode> = new Map();
-  #deps: Map<Identifier, DependencyNode> = new Map();
+  #hoistableObjects: Map<Identifier, HoistableNode & {reactive: boolean}> =
+    new Map();
+  #deps: Map<Identifier, DependencyNode & {reactive: boolean}> = new Map();
 
   /**
    * @param hoistableObjects a set of paths from which we can safely evaluate
@@ -35,9 +37,10 @@ export class ReactiveScopeDependencyTreeHIR {
    * duplicates when traversing the CFG.
    */
   constructor(hoistableObjects: Iterable<ReactiveScopeDependency>) {
-    for (const {path, identifier} of hoistableObjects) {
+    for (const {path, identifier, reactive} of hoistableObjects) {
       let currNode = ReactiveScopeDependencyTreeHIR.#getOrCreateRoot(
         identifier,
+        reactive,
         this.#hoistableObjects,
         path.length > 0 && path[0].optional ? 'Optional' : 'NonNull',
       );
@@ -70,7 +73,8 @@ export class ReactiveScopeDependencyTreeHIR {
 
   static #getOrCreateRoot<T extends string>(
     identifier: Identifier,
-    roots: Map<Identifier, TreeNode<T>>,
+    reactive: boolean,
+    roots: Map<Identifier, TreeNode<T> & {reactive: boolean}>,
     defaultAccessType: T,
   ): TreeNode<T> {
     // roots can always be accessed unconditionally in JS
@@ -79,9 +83,16 @@ export class ReactiveScopeDependencyTreeHIR {
     if (rootNode === undefined) {
       rootNode = {
         properties: new Map(),
+        reactive,
         accessType: defaultAccessType,
       };
       roots.set(identifier, rootNode);
+    } else {
+      CompilerError.invariant(reactive === rootNode.reactive, {
+        reason: 'Conflicting reactive status',
+        description: `Identifier ${printIdentifier(identifier)}`,
+        loc: GeneratedSource,
+      });
     }
     return rootNode;
   }
@@ -92,9 +103,10 @@ export class ReactiveScopeDependencyTreeHIR {
    * safe-to-evaluate subpath
    */
   addDependency(dep: ReactiveScopePropertyDependency): void {
-    const {identifier, path} = dep;
+    const {identifier, reactive, path} = dep;
     let depCursor = ReactiveScopeDependencyTreeHIR.#getOrCreateRoot(
       identifier,
+      reactive,
       this.#deps,
       PropertyAccessType.UnconditionalAccess,
     );
@@ -172,7 +184,13 @@ export class ReactiveScopeDependencyTreeHIR {
   deriveMinimalDependencies(): Set<ReactiveScopeDependency> {
     const results = new Set<ReactiveScopeDependency>();
     for (const [rootId, rootNode] of this.#deps.entries()) {
-      collectMinimalDependenciesInSubtree(rootNode, rootId, [], results);
+      collectMinimalDependenciesInSubtree(
+        rootNode,
+        rootNode.reactive,
+        rootId,
+        [],
+        results,
+      );
     }
 
     return results;
@@ -303,16 +321,18 @@ type DependencyNode = TreeNode<PropertyAccessType>;
  */
 function collectMinimalDependenciesInSubtree(
   node: DependencyNode,
+  reactive: boolean,
   rootIdentifier: Identifier,
   path: Array<DependencyPathEntry>,
   results: Set<ReactiveScopeDependency>,
 ): void {
   if (isDependency(node.accessType)) {
-    results.add({identifier: rootIdentifier, path});
+    results.add({identifier: rootIdentifier, reactive, path});
   } else {
     for (const [childName, childNode] of node.properties) {
       collectMinimalDependenciesInSubtree(
         childNode,
+        reactive,
         rootIdentifier,
         [
           ...path,