[compiler] ReactiveIR: refactor Branch/Fallthrough

Continuing the exploration of the sea-of-nodes ReactiveIR. Three main changes: * ReactiveGraph has both entry and exit. Entry is mostly for completeness since it would be the first entry in the nodes map. * Moves the 'terminal' data from JoinNode to BranchNode. Each branch has both an entry and exit, but these are intended for use with reconstruction back to ReactiveFunction and for interpreting the IR or forward data flow analysis. * Renames JoinNode => FallthroughNode, removing the terminal. Combined with the previous changes — ensuring that every block ends in a terminal that has controls transitively through to its entry node — it's now always possible to reach all nodes from the exit and to walk forward for forward data flow analysis or abstract/actual interpretation of the instructions. [ghstack-poisoned]
facebook · Jan 8, 2025 · e465ec7 · e465ec7
1 parent d9441e4
commit e465ec7
Show file tree

Hide file tree

Showing 2 changed files with 92 additions and 47 deletions.
diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveIR/BuildReactiveGraph.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveIR/BuildReactiveGraph.ts
@@ -31,7 +31,6 @@ import {
   ControlNode,
   EntryNode,
   InstructionNode,
-  JoinNode,
   LoadArgumentNode,
   LoadNode,
   makeReactiveId,
@@ -45,6 +44,8 @@ import {
   reversePostorderReactiveGraph,
   StoreNode,
   eachNodeDependency,
+  FallthroughNode,
+  printReactiveGraph,
 } from './ReactiveIR';
 
 export function buildReactiveGraph(fn: HIRFunction): ReactiveGraph {
@@ -73,7 +74,10 @@ export function buildReactiveGraph(fn: HIRFunction): ReactiveGraph {
 
   const exitNode = buildBlockScope(fn, context, fn.body.entry, entryNode.id);
 
-  const graph = builder.build(fn, exitNode);
+  const graph = builder.build(fn, entryNode.id, exitNode);
+
+  console.log();
+  console.log(printReactiveGraph(graph));
 
   populateReactiveGraphNodeOutputs(graph);
   reversePostorderReactiveGraph(graph);
@@ -85,11 +89,12 @@ class Builder {
   #environment: Map<IdentifierId, ReactiveId> = new Map();
   #nodes: Map<ReactiveId, ReactiveNode> = new Map();
 
-  build(fn: HIRFunction, exit: ReactiveId): ReactiveGraph {
+  build(fn: HIRFunction, entry: ReactiveId, exit: ReactiveId): ReactiveGraph {
     const graph: ReactiveGraph = {
       async: fn.async,
       directives: fn.directives,
       env: fn.env,
+      entry,
       exit,
       fnType: fn.fnType,
       generator: fn.generator,
@@ -494,24 +499,16 @@ function buildBlockScope(
           from: {...terminal.test},
           as: {...terminal.test},
         };
-        const branch: BranchNode = {
-          kind: 'Branch',
-          control,
-          dependencies: [],
-          id: context.nextReactiveId,
-          loc: terminal.loc,
-          outputs: [],
-        };
-        context.putNode(branch);
-        const joinNodeId = context.nextReactiveId;
+        const branchNodeId = context.nextReactiveId;
+        const fallthroughNodeId = context.nextReactiveId;
         const joinFallthrough = {
           kind: 'If',
           block: terminal.fallthrough,
-          fallthrough: joinNodeId,
+          fallthrough: fallthroughNodeId,
         } as const;
         const consequentContext = context.fork(joinFallthrough);
         const consequentControl = consequentContext.controlNode(
-          branch.id,
+          branchNodeId,
           terminal.loc,
         );
         const consequent = buildBlockScope(
@@ -522,7 +519,7 @@ function buildBlockScope(
         );
         const alternateContext = context.fork(joinFallthrough);
         const alternateControl = alternateContext.controlNode(
-          branch.id,
+          branchNodeId,
           terminal.loc,
         );
         const alternate =
@@ -534,19 +531,37 @@ function buildBlockScope(
                 alternateControl,
               )
             : alternateControl;
-        const ifNode: JoinNode = {
-          kind: 'Join',
-          control: branch.id,
-          id: joinNodeId,
+
+        const branch: BranchNode = {
+          kind: 'Branch',
+          control,
+          dependencies: [],
+          id: branchNodeId,
           loc: terminal.loc,
           outputs: [],
+          fallthrough: fallthroughNodeId,
           terminal: {
             kind: 'If',
             test,
-            consequent,
-            alternate,
+            consequent: {
+              entry: consequentControl,
+              exit: consequent,
+            },
+            alternate: {
+              entry: alternateControl,
+              exit: alternate,
+            },
           },
         };
+        context.putNode(branch);
+        const ifNode: FallthroughNode = {
+          kind: 'Fallthrough',
+          control: branch.id,
+          id: fallthroughNodeId,
+          loc: terminal.loc,
+          outputs: [],
+          branches: [consequent, alternate],
+        };
 
         const predecessors: Array<{
           enter: ReactiveId;

diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveIR/ReactiveIR.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveIR/ReactiveIR.ts
@@ -22,6 +22,7 @@ import {assertExhaustive} from '../Utils/utils';
 export type ReactiveGraph = {
   nodes: Map<ReactiveId, ReactiveNode>;
   nextNodeId: number;
+  entry: ReactiveId;
   exit: ReactiveId;
   loc: SourceLocation;
   id: string | null;
@@ -57,7 +58,7 @@ export type ReactiveNode =
   | LoadArgumentNode
   | InstructionNode
   | BranchNode
-  | JoinNode
+  | FallthroughNode
   | ControlNode
   | ReturnNode
   | GotoNode;
@@ -146,24 +147,26 @@ export type BranchNode = {
   outputs: Array<ReactiveId>;
   dependencies: Array<ReactiveId>; // values/scopes depended on by more than one branch, or by the terminal
   control: ReactiveId;
+  fallthrough: ReactiveId;
+  terminal: BranchTerminal;
 };
 
-export type JoinNode = {
-  kind: 'Join';
-  id: ReactiveId;
-  loc: SourceLocation;
-  outputs: Array<ReactiveId>;
-  terminal: NodeTerminal;
-  control: ReactiveId; // join node always has a control, which is the corresponding Branch node
-};
-
-export type NodeTerminal = IfBranch;
+export type BranchTerminal = IfBranch;
 
 export type IfBranch = {
   kind: 'If';
   test: NodeReference;
-  consequent: ReactiveId;
-  alternate: ReactiveId;
+  consequent: {entry: ReactiveId; exit: ReactiveId};
+  alternate: {entry: ReactiveId; exit: ReactiveId};
+};
+
+export type FallthroughNode = {
+  kind: 'Fallthrough';
+  id: ReactiveId;
+  loc: SourceLocation;
+  outputs: Array<ReactiveId>;
+  control: ReactiveId; // always the corresponding branch node
+  branches: Array<ReactiveId>; // the other control-flow paths that reach the fallthrough
 };
 
 export type ControlNode = {
@@ -238,22 +241,34 @@ export function reversePostorderReactiveGraph(graph: ReactiveGraph): void {
   graph.nodes = nodes;
 }
 
+export function* eachBranchTerminalDependency(
+  terminal: BranchTerminal,
+): Iterable<ReactiveId> {
+  switch (terminal.kind) {
+    case 'If': {
+      yield terminal.test.node;
+    }
+  }
+}
+
 export function* eachNodeDependency(node: ReactiveNode): Iterable<ReactiveId> {
   switch (node.kind) {
     case 'Entry':
     case 'LoadArgument': {
       break;
     }
     case 'Goto':
-    case 'Control':
+    case 'Control': {
+      yield* node.dependencies;
+      break;
+    }
     case 'Branch': {
       yield* node.dependencies;
+      yield* eachBranchTerminalDependency(node.terminal);
       break;
     }
-    case 'Join': {
-      yield node.terminal.test.node;
-      yield node.terminal.consequent;
-      yield node.terminal.alternate;
+    case 'Fallthrough': {
+      yield* node.branches;
       break;
     }
     case 'Load': {
@@ -282,6 +297,17 @@ export function* eachNodeDependency(node: ReactiveNode): Iterable<ReactiveId> {
   }
 }
 
+export function* eachBranchTerminalReference(
+  terminal: BranchTerminal,
+): Iterable<NodeReference> {
+  switch (terminal.kind) {
+    case 'If': {
+      yield terminal.test;
+      break;
+    }
+  }
+}
+
 export function* eachNodeReference(
   node: ReactiveNode,
 ): Iterable<NodeReference> {
@@ -305,10 +331,10 @@ export function* eachNodeReference(
       break;
     }
     case 'Branch': {
+      yield* eachBranchTerminalReference(node.terminal);
       break;
     }
-    case 'Join': {
-      yield node.terminal.test;
+    case 'Fallthrough': {
       break;
     }
     case 'Value': {
@@ -415,14 +441,12 @@ function writeReactiveNodes(
         buffer.push(
           `£${id} Branch deps=[${node.dependencies.map(id => `£${id}`).join(', ')}]${control}`,
         );
-        break;
-      }
-      case 'Join': {
-        buffer.push(`£${id} Join${control}`);
         switch (node.terminal.kind) {
           case 'If': {
             buffer.push(
-              `  If test=${printNodeReference(node.terminal.test)} consequent=£${node.terminal.consequent} alternate=£${node.terminal.alternate}${control}`,
+              `  If test=${printNodeReference(node.terminal.test)} ` +
+                `consequent=£${node.terminal.consequent.entry}:${node.terminal.consequent.exit} ` +
+                `alternate=£${node.terminal.alternate.entry}:${node.terminal.alternate.exit}`,
             );
             break;
           }
@@ -432,6 +456,12 @@ function writeReactiveNodes(
         }
         break;
       }
+      case 'Fallthrough': {
+        buffer.push(
+          `£${id} Fallthrough${control} branches=[${node.branches.map(id => `£${id}`).join(', ')}]`,
+        );
+        break;
+      }
       case 'Value': {
         const deps = [...eachNodeReference(node)]
           .map(id => printNodeReference(id))