[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/Entrypoint/Pipeline.ts

@@ -123,6 +123,7 @@ function run(
     fnType,
     config,
     contextIdentifiers,
+    func,
     logger,
     filename,
     code,

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

@@ -68,12 +68,14 @@ import {BuiltInArrayId} from './ObjectShape';
 export function lower(
   func: NodePath<t.Function>,
   env: Environment,
+  // Bindings captured from the outer function, in case lower() is called recursively (for lambdas)
   bindings: Bindings | null = null,
   capturedRefs: Array<t.Identifier> = [],
-  // 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, {
+    bindings,
+    context: capturedRefs,
+  });
   const context: HIRFunction['context'] = [];
 
   for (const ref of capturedRefs ?? []) {
@@ -213,7 +215,7 @@ export function lower(
   return Ok({
     id,
     params,
-    fnType: parent == null ? env.fnType : 'Other',
+    fnType: bindings == null ? env.fnType : 'Other',
     returnTypeAnnotation: null, // TODO: extract the actual return type node if present
     returnType: makeType(),
     body: builder.build(),
@@ -3375,7 +3377,7 @@ function lowerFunction(
     | t.ObjectMethod
   >,
 ): LoweredFunction | null {
-  const componentScope: Scope = builder.parentFunction.scope;
+  const componentScope: Scope = builder.environment.parentFunction.scope;
   const capturedContext = gatherCapturedContext(expr, componentScope);
 
   /*
@@ -3386,13 +3388,10 @@ function lowerFunction(
    * This isn't a problem in practice because use Babel's scope analysis to
    * identify the correct references.
    */
-  const lowering = lower(
-    expr,
-    builder.environment,
-    builder.bindings,
-    [...builder.context, ...capturedContext],
-    builder.parentFunction,
-  );
+  const lowering = lower(expr, builder.environment, builder.bindings, [
+    ...builder.context,
+    ...capturedContext,
+  ]);
   let loweredFunc: HIRFunction;
   if (lowering.isErr()) {
     lowering
@@ -3414,7 +3413,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

@@ -194,7 +194,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 +211,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 +240,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 +256,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 +278,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 +345,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 +576,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

@@ -24,8 +24,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
@@ -34,9 +35,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',
       );
@@ -69,7 +71,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
@@ -78,9 +81,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: '[DeriveMinimalDependenciesHIR] Conflicting reactive root flag',
+        description: `Identifier ${printIdentifier(identifier)}`,
+        loc: GeneratedSource,
+      });
     }
     return rootNode;
   }
@@ -91,9 +101,10 @@ export class ReactiveScopeDependencyTreeHIR {
    * safe-to-evaluate subpath
    */
   addDependency(dep: ReactiveScopeDependency): void {
-    const {identifier, path} = dep;
+    const {identifier, reactive, path} = dep;
     let depCursor = ReactiveScopeDependencyTreeHIR.#getOrCreateRoot(
       identifier,
+      reactive,
       this.#deps,
       PropertyAccessType.UnconditionalAccess,
     );
@@ -171,7 +182,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;
@@ -293,25 +310,24 @@ type HoistableNode = TreeNode<'Optional' | 'NonNull'>;
 type DependencyNode = TreeNode<PropertyAccessType>;
 
 /**
- * TODO: this is directly pasted from DeriveMinimalDependencies. Since we no
- * longer have conditionally accessed nodes, we can simplify
- *
  * Recursively calculates minimal dependencies in a subtree.
  * @param node DependencyNode representing a dependency subtree.
  * @returns a minimal list of dependencies in this subtree.
  */
 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,