[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/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,19 @@ 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:
+          '[HoistablePropertyLoads] Found inconsistencies in `reactive` flag after ',
+        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

@@ -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,

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

@@ -62,12 +62,14 @@ export type ReactiveFunction = {
 
 export type ReactiveScopeBlock = {
   kind: 'scope';
+  dependencyInstructions: Array<ReactiveInstructionStatement>;
   scope: ReactiveScope;
   instructions: ReactiveBlock;
 };
 
 export type PrunedReactiveScopeBlock = {
   kind: 'pruned-scope';
+  dependencyInstructions: Array<ReactiveInstructionStatement>;
   scope: ReactiveScope;
   instructions: ReactiveBlock;
 };
@@ -614,6 +616,7 @@ export type MaybeThrowTerminal = {
 export type ReactiveScopeTerminal = {
   kind: 'scope';
   fallthrough: BlockId;
+  dependencies: BlockId;
   block: BlockId;
   scope: ReactiveScope;
   id: InstructionId;
@@ -623,6 +626,7 @@ export type ReactiveScopeTerminal = {
 export type PrunedScopeTerminal = {
   kind: 'pruned-scope';
   fallthrough: BlockId;
+  dependencies: BlockId;
   block: BlockId;
   scope: ReactiveScope;
   id: InstructionId;
@@ -1453,9 +1457,10 @@ export type ReactiveScope = {
   range: MutableRange;
 
   /**
-   * The inputs to this reactive scope
+   * Note the dependencies of a reactive scope are tracked in HIR and
+   * ReactiveFunction
    */
-  dependencies: ReactiveScopeDependencies;
+  dependencies: Array<Place>;
 
   /**
    * The set of values produced by this scope. This may be empty
@@ -1506,6 +1511,18 @@ export type DependencyPathEntry = {property: string; optional: boolean};
 export type DependencyPath = Array<DependencyPathEntry>;
 export type ReactiveScopeDependency = {
   identifier: Identifier;
+  /**
+   * Reflects whether the base identifier is reactive. Note that some reactive
+   * objects may have non-reactive properties, but we do not currently track
+   * this.
+   *
+   * ```js
+   * // Technically, result[0] is reactive and result[1] is not.
+   * // Currently, both dependencies would be marked as reactive.
+   * const result = useState();
+   * ```
+   */
+  reactive: boolean;
   path: DependencyPath;
 };
 

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

@@ -6,7 +6,6 @@
  */
 
 import generate from '@babel/generator';
-import {printReactiveFunction} from '..';
 import {CompilerError} from '../CompilerError';
 import {printReactiveScopeSummary} from '../ReactiveScopes/PrintReactiveFunction';
 import DisjointSet from '../Utils/DisjointSet';
@@ -287,13 +286,13 @@ export function printTerminal(terminal: Terminal): Array<string> | string {
     case 'scope': {
       value = `[${terminal.id}] Scope ${printReactiveScopeSummary(
         terminal.scope,
-      )} block=bb${terminal.block} fallthrough=bb${terminal.fallthrough}`;
+      )} dependencies=bb${terminal.dependencies} block=bb${terminal.block} fallthrough=bb${terminal.fallthrough}`;
       break;
     }
     case 'pruned-scope': {
       value = `[${terminal.id}] <pruned> Scope ${printReactiveScopeSummary(
         terminal.scope,
-      )} block=bb${terminal.block} fallthrough=bb${terminal.fallthrough}`;
+      )} dependencies=bb${terminal.dependencies} block=bb${terminal.block} fallthrough=bb${terminal.fallthrough}`;
       break;
     }
     case 'try': {