From 98bc4f9d3ecff4741b70e4f5d418584e89bf3954 Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Fri, 2 May 2025 15:55:03 +0900 Subject: [PATCH] [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] --- .../src/HIR/BuildHIR.ts | 10 ++- .../src/HIR/HIR.ts | 5 ++ .../src/HIR/PrintHIR.ts | 4 + .../src/HIR/visitors.ts | 2 + .../src/Inference/InferReferenceEffects.ts | 5 +- .../src/Optimization/DeadCodeElimination.ts | 1 + .../src/Optimization/OutlineJsx.ts | 1 + .../ReactiveScopes/CodegenReactiveFunction.ts | 6 ++ .../InferReactiveScopeVariables.ts | 1 + .../ReactiveScopes/PruneNonEscapingScopes.ts | 1 + .../src/TypeInference/InferTypes.ts | 5 ++ .../compiler/non-null-assertion.expect.md | 15 ++-- .../compiler/non-null-expression.expect.md | 83 +++++++++++++++++++ .../fixtures/compiler/non-null-expression.tsx | 17 ++++ 14 files changed, 146 insertions(+), 10 deletions(-) create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/non-null-expression.expect.md create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/non-null-expression.tsx diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildHIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildHIR.ts index b9f82eea18e9f..0f0a8a1dccfce 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildHIR.ts +++ b/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; + const value = lowerExpressionToTemporary(builder, expr.get('expression')); + return { + kind: 'NonNullExpression', + value, + loc: exprLoc, + }; + } + case 'TSInstantiationExpression': { let expr = exprPath as NodePath; return lowerExpression(builder, expr.get('expression')); } diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts index 611b5bd210226..32e14604e562f 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts +++ b/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'; diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/PrintHIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/PrintHIR.ts index c8182c9e72a7c..be178063e63cc 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/PrintHIR.ts +++ b/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, diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/visitors.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/visitors.ts index 49ff3c256e016..cbc2dc33970e3 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/visitors.ts +++ b/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; diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts index d1546038edcbe..ea9f6a1861a54 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts +++ b/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 diff --git a/compiler/packages/babel-plugin-react-compiler/src/Optimization/DeadCodeElimination.ts b/compiler/packages/babel-plugin-react-compiler/src/Optimization/DeadCodeElimination.ts index 2b752c6dfd28e..b46e361a5be13 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Optimization/DeadCodeElimination.ts +++ b/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 diff --git a/compiler/packages/babel-plugin-react-compiler/src/Optimization/OutlineJsx.ts b/compiler/packages/babel-plugin-react-compiler/src/Optimization/OutlineJsx.ts index d35c4d77362db..bb010a0fcfe21 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Optimization/OutlineJsx.ts +++ b/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': { diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/CodegenReactiveFunction.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/CodegenReactiveFunction.ts index 790a64b407316..71bb36da24cea 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/CodegenReactiveFunction.ts +++ b/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': diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts index 0c1fd759bd5b8..75dfbbcb317f0 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts +++ b/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': diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts index 5ae4c7dfc72f9..33ed822bb1d76 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts +++ b/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 diff --git a/compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts b/compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts index 69812fc130ded..ffb453039f6e1 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts +++ b/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); diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/non-null-assertion.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/non-null-assertion.expect.md index bf8590ec3a63c..34c54c618cc85 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/non-null-assertion.expect.md +++ b/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 = { diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/non-null-expression.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/non-null-expression.expect.md new file mode 100644 index 0000000000000..9246326d3ad3b --- /dev/null +++ b/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 ( + <> +