[compiler] Preserve TSNonNullExpressions

This is a first step toward teaching the compiler that TS non-null expressions mean a variable may be nullable and that we can't take dependencies past the non-null assertion.

To start, we add a new instruction type and add support from lowering through codegen.

[ghstack-poisoned]

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

@@ -2563,8 +2563,16 @@ function lowerExpression(
         loc: expr.node.loc ?? GeneratedSource,
       };
     }
-    case 'TSInstantiationExpression':
     case 'TSNonNullExpression': {
+      let expr = exprPath as NodePath<t.TSNonNullExpression>;
+      const value = lowerExpressionToTemporary(builder, expr.get('expression'));
+      return {
+        kind: 'NonNullExpression',
+        value,
+        loc: exprLoc,
+      };
+    }
+    case 'TSInstantiationExpression': {
       let expr = exprPath as NodePath<t.TSNonNullExpression>;
       return lowerExpression(builder, expr.get('expression'));
     }

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

@@ -958,6 +958,11 @@ export type InstructionValue =
           typeAnnotationKind: 'as' | 'satisfies';
         }
     ))
+  | {
+      kind: 'NonNullExpression';
+      value: Place;
+      loc: SourceLocation;
+    }
   | JsxExpression
   | {
       kind: 'ObjectExpression';

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

@@ -703,6 +703,10 @@ export function printInstructionValue(instrValue: ReactiveValue): string {
       value = `FinishMemoize decl=${printPlace(instrValue.decl)}`;
       break;
     }
+    case 'NonNullExpression': {
+      value = `NonNullExpression ${printPlace(instrValue.value)}`;
+      break;
+    }
     default: {
       assertExhaustive(
         instrValue,

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

@@ -200,6 +200,7 @@ export function* eachInstructionValueOperand(
       yield instrValue.tag;
       break;
     }
+    case 'NonNullExpression':
     case 'TypeCastExpression': {
       yield instrValue.value;
       break;
@@ -526,6 +527,7 @@ export function mapInstructionValueOperands(
       instrValue.tag = fn(instrValue.tag);
       break;
     }
+    case 'NonNullExpression':
     case 'TypeCastExpression': {
       instrValue.value = fn(instrValue.value);
       break;

compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts

@@ -1422,10 +1422,11 @@ function inferBlock(
         continuation = {kind: 'funeffects'};
         break;
       }
+      case 'NonNullExpression':
       case 'TypeCastExpression': {
         /*
-         * A type cast expression has no effect at runtime, so it's equivalent to a raw
-         * identifier:
+         * A non-null expression or type cast expression has no effect at runtime,
+         * so it's equivalent to a raw identifier:
          * ```
          * x = (y: type)  // is equivalent to...
          * x = y

compiler/packages/babel-plugin-react-compiler/src/Optimization/DeadCodeElimination.ts

@@ -371,6 +371,7 @@ function pruneableValue(value: InstructionValue, state: State): boolean {
     case 'Primitive':
     case 'PropertyLoad':
     case 'TemplateLiteral':
+    case 'NonNullExpression':
     case 'TypeCastExpression':
     case 'UnaryExpression': {
       // Definitely safe to prune since they are read-only

compiler/packages/babel-plugin-react-compiler/src/Optimization/OutlineJsx.ts

@@ -148,6 +148,7 @@ function outlineJsxImpl(
         case 'StoreLocal':
         case 'TaggedTemplateExpression':
         case 'TemplateLiteral':
+        case 'NonNullExpression':
         case 'TypeCastExpression':
         case 'UnsupportedNode':
         case 'UnaryExpression': {

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/CodegenReactiveFunction.ts

@@ -2292,6 +2292,12 @@ function codegenInstructionValue(
       );
       break;
     }
+    case 'NonNullExpression': {
+      value = t.tsNonNullExpression(
+        codegenPlaceToExpression(cx, instrValue.value),
+      );
+      break;
+    }
     case 'StartMemoize':
     case 'FinishMemoize':
     case 'Debugger':

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts

@@ -220,6 +220,7 @@ function mayAllocate(_env: Environment, instruction: Instruction): boolean {
     case 'StoreLocal':
     case 'LoadGlobal':
     case 'MetaProperty':
+    case 'NonNullExpression':
     case 'TypeCastExpression':
     case 'LoadLocal':
     case 'LoadContext':

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts

@@ -544,6 +544,7 @@ class CollectDependenciesVisitor extends ReactiveFunctionVisitor<
         };
       }
       case 'Await':
+      case 'NonNullExpression':
       case 'TypeCastExpression': {
         return {
           // Indirection for the inner value, memoized if the value is

compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts

@@ -386,6 +386,11 @@ function* generateInstructionTypes(
       break;
     }
 
+    case 'NonNullExpression': {
+      yield equation(left, value.value.identifier.type);
+      break;
+    }
+
     case 'TypeCastExpression': {
       if (env.config.enableUseTypeAnnotations) {
         yield equation(value.type, value.value.identifier.type);

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/non-null-assertion.expect.md

@@ -27,15 +27,16 @@ interface ComponentProps {
 
 function Component(props) {
   const $ = _c(2);
-  let t0;
-  if ($[0] !== props.name) {
-    t0 = props.name.toUpperCase();
-    $[0] = props.name;
-    $[1] = t0;
+  const t0 = props.name!;
+  let t1;
+  if ($[0] !== t0) {
+    t1 = t0.toUpperCase();
+    $[0] = t0;
+    $[1] = t1;
   } else {
-    t0 = $[1];
+    t1 = $[1];
   }
-  return t0;
+  return t1;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/non-null-expression.expect.md

@@ -0,0 +1,83 @@
+
+## Input
+
+```javascript
+import {useState} from 'react';
+
+function Component() {
+  const [value, setValue] = useState(null);
+  const createValue = () => {
+    setValue({value: 42});
+  };
+  const logValue = () => {
+    console.log(value!.value);
+  };
+  return (
+    <>
+      <button onClick={createValue} />
+      <button disabled={value == null} onClick={logValue} />
+    </>
+  );
+}
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { useState } from "react";
+
+function Component() {
+  const $ = _c(7);
+  const [value, setValue] = useState(null);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = () => {
+      setValue({ value: 42 });
+    };
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const createValue = t0;
+  let t1;
+  if ($[1] !== value) {
+    t1 = () => {
+      console.log(value!.value);
+    };
+    $[1] = value;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  const logValue = t1;
+  let t2;
+  if ($[3] === Symbol.for("react.memo_cache_sentinel")) {
+    t2 = <button onClick={createValue} />;
+    $[3] = t2;
+  } else {
+    t2 = $[3];
+  }
+  const t3 = value == null;
+  let t4;
+  if ($[4] !== logValue || $[5] !== t3) {
+    t4 = (
+      <>
+        {t2}
+        <button disabled={t3} onClick={logValue} />
+      </>
+    );
+    $[4] = logValue;
+    $[5] = t3;
+    $[6] = t4;
+  } else {
+    t4 = $[6];
+  }
+  return t4;
+}
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/non-null-expression.tsx

@@ -0,0 +1,17 @@
+import {useState} from 'react';
+
+function Component() {
+  const [value, setValue] = useState(null);
+  const createValue = () => {
+    setValue({value: 42});
+  };
+  const logValue = () => {
+    console.log(value!.value);
+  };
+  return (
+    <>
+      <button onClick={createValue} />
+      <button disabled={value == null} onClick={logValue} />
+    </>
+  );
+}