[compiler] Allow refs to be lazily initialized during render

Summary:
The official guidance for useRef notes an exception to the rule that refs cannot be accessed during render: to avoid recreating the ref's contents, you can test that the ref is uninitialized and then initialize it using an if statement:

```
if (ref.current == null) {
  ref.current = SomeExpensiveOperation()
}
```

The compiler didn't recognize this exception, however, leading to code that obeyed all the official guidance for refs being rejected by the compiler. This PR fixes that, by extending the ref validation machinery with an awareness of guard operations that allow lazy initialization. We now understand `== null` and similar operations, when applied to a ref and consumed by an if terminal, as marking the consequent of the if as a block in which the ref can be safely written to. In order to do so we need to create a notion of ref ids, which link different usages of the same ref via both the ref and the ref value.

[ghstack-poisoned]

Author: mvitousek

compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateNoRefAccesInRender.ts

@@ -7,8 +7,8 @@
 
 import {CompilerError, ErrorSeverity} from '../CompilerError';
 import {
+  BlockId,
   HIRFunction,
-  Identifier,
   IdentifierId,
   Place,
   SourceLocation,
@@ -17,6 +17,7 @@ import {
   isUseRefType,
 } from '../HIR';
 import {
+  eachInstructionOperand,
   eachInstructionValueOperand,
   eachPatternOperand,
   eachTerminalOperand,
@@ -44,11 +45,32 @@ import {Err, Ok, Result} from '../Utils/Result';
  * or based on property name alone (`foo.current` might be a ref).
  */
 
-type RefAccessType = {kind: 'None'} | RefAccessRefType;
+const opaqueRefId = Symbol();
+type RefId = number & {[opaqueRefId]: 'RefId'};
+
+function makeRefId(id: number): RefId {
+  CompilerError.invariant(id >= 0 && Number.isInteger(id), {
+    reason: 'Expected identifier id to be a non-negative integer',
+    description: null,
+    loc: null,
+    suggestions: null,
+  });
+  return id as RefId;
+}
+let _refId = 0;
+function nextRefId(): RefId {
+  return makeRefId(_refId++);
+}
+
+type RefAccessType =
+  | {kind: 'None'}
+  | {kind: 'Nullable'}
+  | {kind: 'Guard'; refId: RefId}
+  | RefAccessRefType;
 
 type RefAccessRefType =
-  | {kind: 'Ref'}
-  | {kind: 'RefValue'; loc?: SourceLocation}
+  | {kind: 'Ref'; refId: RefId}
+  | {kind: 'RefValue'; loc?: SourceLocation; refId?: RefId}
   | {kind: 'Structure'; value: null | RefAccessRefType; fn: null | RefFnType};
 
 type RefFnType = {readRefEffect: boolean; returnType: RefAccessType};
@@ -82,11 +104,11 @@ export function validateNoRefAccessInRender(fn: HIRFunction): void {
   validateNoRefAccessInRenderImpl(fn, env).unwrap();
 }
 
-function refTypeOfType(identifier: Identifier): RefAccessType {
-  if (isRefValueType(identifier)) {
+function refTypeOfType(place: Place): RefAccessType {
+  if (isRefValueType(place.identifier)) {
     return {kind: 'RefValue'};
-  } else if (isUseRefType(identifier)) {
-    return {kind: 'Ref'};
+  } else if (isUseRefType(place.identifier)) {
+    return {kind: 'Ref', refId: nextRefId()};
   } else {
     return {kind: 'None'};
   }
@@ -101,6 +123,14 @@ function tyEqual(a: RefAccessType, b: RefAccessType): boolean {
       return true;
     case 'Ref':
       return true;
+    case 'Nullable':
+      return true;
+    case 'Guard':
+      CompilerError.invariant(b.kind === 'Guard', {
+        reason: 'Expected ref value',
+        loc: null,
+      });
+      return a.refId === b.refId;
     case 'RefValue':
       CompilerError.invariant(b.kind === 'RefValue', {
         reason: 'Expected ref value',
@@ -133,11 +163,17 @@ function joinRefAccessTypes(...types: Array<RefAccessType>): RefAccessType {
     b: RefAccessRefType,
   ): RefAccessRefType {
     if (a.kind === 'RefValue') {
-      return a;
+      if (b.kind === 'RefValue' && a.refId === b.refId) {
+        return a;
+      }
+      return {kind: 'RefValue'};
     } else if (b.kind === 'RefValue') {
       return b;
     } else if (a.kind === 'Ref' || b.kind === 'Ref') {
-      return {kind: 'Ref'};
+      if (a.kind === 'Ref' && b.kind === 'Ref' && a.refId === b.refId) {
+        return a;
+      }
+      return {kind: 'Ref', refId: nextRefId()};
     } else {
       CompilerError.invariant(
         a.kind === 'Structure' && b.kind === 'Structure',
@@ -178,6 +214,16 @@ function joinRefAccessTypes(...types: Array<RefAccessType>): RefAccessType {
         return b;
       } else if (b.kind === 'None') {
         return a;
+      } else if (a.kind === 'Guard' || b.kind === 'Guard') {
+        if (a.kind === 'Guard' && b.kind === 'Guard' && a.refId === b.refId) {
+          return a;
+        }
+        return {kind: 'None'};
+      } else if (a.kind === 'Nullable' || b.kind === 'Nullable') {
+        if (a.kind === 'Nullable' && b.kind === 'Nullable') {
+          return a;
+        }
+        return {kind: 'None'};
       } else {
         return joinRefAccessRefTypes(a, b);
       }
@@ -198,13 +244,14 @@ function validateNoRefAccessInRenderImpl(
     } else {
       place = param.place;
     }
-    const type = refTypeOfType(place.identifier);
+    const type = refTypeOfType(place);
     env.set(place.identifier.id, type);
   }
 
   for (let i = 0; (i == 0 || env.hasChanged()) && i < 10; i++) {
     env.resetChanged();
     returnValues = [];
+    const safeBlocks = new Map<BlockId, RefId>();
     const errors = new CompilerError();
     for (const [, block] of fn.body.blocks) {
       for (const phi of block.phis) {
@@ -238,11 +285,15 @@ function validateNoRefAccessInRenderImpl(
             if (objType?.kind === 'Structure') {
               lookupType = objType.value;
             } else if (objType?.kind === 'Ref') {
-              lookupType = {kind: 'RefValue', loc: instr.loc};
+              lookupType = {
+                kind: 'RefValue',
+                loc: instr.loc,
+                refId: objType.refId,
+              };
             }
             env.set(
               instr.lvalue.identifier.id,
-              lookupType ?? refTypeOfType(instr.lvalue.identifier),
+              lookupType ?? refTypeOfType(instr.lvalue),
             );
             break;
           }
@@ -251,7 +302,7 @@ function validateNoRefAccessInRenderImpl(
             env.set(
               instr.lvalue.identifier.id,
               env.get(instr.value.place.identifier.id) ??
-                refTypeOfType(instr.lvalue.identifier),
+                refTypeOfType(instr.lvalue),
             );
             break;
           }
@@ -260,12 +311,12 @@ function validateNoRefAccessInRenderImpl(
             env.set(
               instr.value.lvalue.place.identifier.id,
               env.get(instr.value.value.identifier.id) ??
-                refTypeOfType(instr.value.lvalue.place.identifier),
+                refTypeOfType(instr.value.lvalue.place),
             );
             env.set(
               instr.lvalue.identifier.id,
               env.get(instr.value.value.identifier.id) ??
-                refTypeOfType(instr.lvalue.identifier),
+                refTypeOfType(instr.lvalue),
             );
             break;
           }
@@ -277,13 +328,10 @@ function validateNoRefAccessInRenderImpl(
             }
             env.set(
               instr.lvalue.identifier.id,
-              lookupType ?? refTypeOfType(instr.lvalue.identifier),
+              lookupType ?? refTypeOfType(instr.lvalue),
             );
             for (const lval of eachPatternOperand(instr.value.lvalue.pattern)) {
-              env.set(
-                lval.identifier.id,
-                lookupType ?? refTypeOfType(lval.identifier),
-              );
+              env.set(lval.identifier.id, lookupType ?? refTypeOfType(lval));
             }
             break;
           }
@@ -354,7 +402,11 @@ function validateNoRefAccessInRenderImpl(
               types.push(env.get(operand.identifier.id) ?? {kind: 'None'});
             }
             const value = joinRefAccessTypes(...types);
-            if (value.kind === 'None') {
+            if (
+              value.kind === 'None' ||
+              value.kind === 'Guard' ||
+              value.kind === 'Nullable'
+            ) {
               env.set(instr.lvalue.identifier.id, {kind: 'None'});
             } else {
               env.set(instr.lvalue.identifier.id, {
@@ -369,7 +421,18 @@ function validateNoRefAccessInRenderImpl(
           case 'PropertyStore':
           case 'ComputedDelete':
           case 'ComputedStore': {
-            validateNoRefAccess(errors, env, instr.value.object, instr.loc);
+            const safe = safeBlocks.get(block.id);
+            const target = env.get(instr.value.object.identifier.id);
+            if (
+              instr.value.kind === 'PropertyStore' &&
+              safe != null &&
+              target?.kind === 'Ref' &&
+              target.refId === safe
+            ) {
+              safeBlocks.delete(block.id);
+            } else {
+              validateNoRefAccess(errors, env, instr.value.object, instr.loc);
+            }
             for (const operand of eachInstructionValueOperand(instr.value)) {
               if (operand === instr.value.object) {
                 continue;
@@ -381,23 +444,67 @@ function validateNoRefAccessInRenderImpl(
           case 'StartMemoize':
           case 'FinishMemoize':
             break;
+          case 'Primitive': {
+            if (instr.value.value == null) {
+              env.set(instr.lvalue.identifier.id, {kind: 'Nullable'});
+            }
+            break;
+          }
+          case 'BinaryExpression': {
+            const left = env.get(instr.value.left.identifier.id);
+            const right = env.get(instr.value.right.identifier.id);
+            let nullish: boolean = false;
+            let refId: RefId | null = null;
+            if (left?.kind === 'RefValue' && left.refId != null) {
+              refId = left.refId;
+            } else if (right?.kind === 'RefValue' && right.refId != null) {
+              refId = right.refId;
+            }
+
+            if (left?.kind === 'Nullable') {
+              nullish = true;
+            } else if (right?.kind === 'Nullable') {
+              nullish = true;
+            }
+
+            if (refId !== null && nullish) {
+              env.set(instr.lvalue.identifier.id, {kind: 'Guard', refId});
+            } else {
+              for (const operand of eachInstructionValueOperand(instr.value)) {
+                validateNoRefValueAccess(errors, env, operand);
+              }
+            }
+            break;
+          }
           default: {
             for (const operand of eachInstructionValueOperand(instr.value)) {
               validateNoRefValueAccess(errors, env, operand);
             }
             break;
           }
         }
-        if (isUseRefType(instr.lvalue.identifier)) {
+
+        // Guard values are derived from ref.current, so they can only be used in if statement targets
+        for (const operand of eachInstructionOperand(instr)) {
+          guardCheck(errors, operand, env);
+        }
+
+        if (
+          isUseRefType(instr.lvalue.identifier) &&
+          env.get(instr.lvalue.identifier.id)?.kind !== 'Ref'
+        ) {
           env.set(
             instr.lvalue.identifier.id,
             joinRefAccessTypes(
               env.get(instr.lvalue.identifier.id) ?? {kind: 'None'},
-              {kind: 'Ref'},
+              {kind: 'Ref', refId: nextRefId()},
             ),
           );
         }
-        if (isRefValueType(instr.lvalue.identifier)) {
+        if (
+          isRefValueType(instr.lvalue.identifier) &&
+          env.get(instr.lvalue.identifier.id)?.kind !== 'RefValue'
+        ) {
           env.set(
             instr.lvalue.identifier.id,
             joinRefAccessTypes(
@@ -407,12 +514,24 @@ function validateNoRefAccessInRenderImpl(
           );
         }
       }
+
+      if (block.terminal.kind === 'if') {
+        const test = env.get(block.terminal.test.identifier.id);
+        if (test?.kind === 'Guard') {
+          safeBlocks.set(block.terminal.consequent, test.refId);
+        }
+      }
+
       for (const operand of eachTerminalOperand(block.terminal)) {
         if (block.terminal.kind !== 'return') {
           validateNoRefValueAccess(errors, env, operand);
+          if (block.terminal.kind !== 'if') {
+            guardCheck(errors, operand, env);
+          }
         } else {
           // Allow functions containing refs to be returned, but not direct ref values
           validateNoDirectRefValueAccess(errors, operand, env);
+          guardCheck(errors, operand, env);
           returnValues.push(env.get(operand.identifier.id));
         }
       }
@@ -444,6 +563,23 @@ function destructure(
   return type;
 }
 
+function guardCheck(errors: CompilerError, operand: Place, env: Env): void {
+  if (env.get(operand.identifier.id)?.kind === 'Guard') {
+    errors.push({
+      severity: ErrorSeverity.InvalidReact,
+      reason:
+        'Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef)',
+      loc: operand.loc,
+      description:
+        operand.identifier.name !== null &&
+        operand.identifier.name.kind === 'named'
+          ? `Cannot access ref value \`${operand.identifier.name.value}\``
+          : null,
+      suggestions: null,
+    });
+  }
+}
+
 function validateNoRefValueAccess(
   errors: CompilerError,
   env: Env,

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-initialization.expect.md

@@ -0,0 +1,42 @@
+
+## Input
+
+```javascript
+//@flow
+import {useRef} from 'react';
+
+component C() {
+  const r = useRef(null);
+  if (r.current == null) {
+    r.current = 1;
+  }
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: C,
+  params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+import { useRef } from "react";
+
+function C() {
+  const r = useRef(null);
+  if (r.current == null) {
+    r.current = 1;
+  }
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: C,
+  params: [{}],
+};
+
+```
+      
+### Eval output
+(kind: ok)