[compiler] Add missing source locations to statements, expressions (facebook#34406)

Adds missing locations to all the statement kinds that we produce in
codegenInstruction(), and adds generic handling of source locations for
the nodes produced by codegenInstructionValue(). There are definitely
some places where we are still missing a location, but this should
address some of the known issues we've seen such as missing location on
`throw`.

---
[//]: # (BEGIN SAPLING FOOTER)
Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with [ReviewStack](https://reviewstack.dev/facebook/react/pull/34406).
* facebook#34394
* __->__ facebook#34406
* facebook#34346

Author: josephsavona

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

@@ -982,7 +982,8 @@ function codegenTerminal(
       if (terminal.targetKind === 'implicit') {
         return null;
       }
-      return t.breakStatement(
+      return createBreakStatement(
+        terminal.loc,
         terminal.targetKind === 'labeled'
           ? t.identifier(codegenLabel(terminal.target))
           : null,
@@ -992,14 +993,16 @@ function codegenTerminal(
       if (terminal.targetKind === 'implicit') {
         return null;
       }
-      return t.continueStatement(
+      return createContinueStatement(
+        terminal.loc,
         terminal.targetKind === 'labeled'
           ? t.identifier(codegenLabel(terminal.target))
           : null,
       );
     }
     case 'for': {
-      return t.forStatement(
+      return createForStatement(
+        terminal.loc,
         codegenForInit(cx, terminal.init),
         codegenInstructionValueToExpression(cx, terminal.test),
         terminal.update !== null
@@ -1108,7 +1111,8 @@ function codegenTerminal(
             `Unhandled lvalue kind: ${iterableItem.value.lvalue.kind}`,
           );
       }
-      return t.forInStatement(
+      return createForInStatement(
+        terminal.loc,
         /*
          * Special handling here since we only want the VariableDeclarators without any inits
          * This needs to be updated when we handle non-trivial ForOf inits
@@ -1225,7 +1229,8 @@ function codegenTerminal(
             `Unhandled lvalue kind: ${iterableItem.value.lvalue.kind}`,
           );
       }
-      return t.forOfStatement(
+      return createForOfStatement(
+        terminal.loc,
         /*
          * Special handling here since we only want the VariableDeclarators without any inits
          * This needs to be updated when we handle non-trivial ForOf inits
@@ -1247,7 +1252,7 @@ function codegenTerminal(
           alternate = block;
         }
       }
-      return t.ifStatement(test, consequent, alternate);
+      return createIfStatement(terminal.loc, test, consequent, alternate);
     }
     case 'return': {
       const value = codegenPlaceToExpression(cx, terminal.value);
@@ -1258,7 +1263,8 @@ function codegenTerminal(
       return t.returnStatement(value);
     }
     case 'switch': {
-      return t.switchStatement(
+      return createSwitchStatement(
+        terminal.loc,
         codegenPlaceToExpression(cx, terminal.test),
         terminal.cases.map(case_ => {
           const test =
@@ -1271,15 +1277,26 @@ function codegenTerminal(
       );
     }
     case 'throw': {
-      return t.throwStatement(codegenPlaceToExpression(cx, terminal.value));
+      return createThrowStatement(
+        terminal.loc,
+        codegenPlaceToExpression(cx, terminal.value),
+      );
     }
     case 'do-while': {
       const test = codegenInstructionValueToExpression(cx, terminal.test);
-      return t.doWhileStatement(test, codegenBlock(cx, terminal.loop));
+      return createDoWhileStatement(
+        terminal.loc,
+        test,
+        codegenBlock(cx, terminal.loop),
+      );
     }
     case 'while': {
       const test = codegenInstructionValueToExpression(cx, terminal.test);
-      return t.whileStatement(test, codegenBlock(cx, terminal.loop));
+      return createWhileStatement(
+        terminal.loc,
+        test,
+        codegenBlock(cx, terminal.loop),
+      );
     }
     case 'label': {
       return codegenBlock(cx, terminal.block);
@@ -1290,7 +1307,8 @@ function codegenTerminal(
         catchParam = convertIdentifier(terminal.handlerBinding.identifier);
         cx.temp.set(terminal.handlerBinding.identifier.declarationId, null);
       }
-      return t.tryStatement(
+      return createTryStatement(
+        terminal.loc,
         codegenBlock(cx, terminal.block),
         t.catchClause(catchParam, codegenBlock(cx, terminal.handler)),
       );
@@ -1695,7 +1713,13 @@ const createExpressionStatement = withLoc(t.expressionStatement);
 const _createLabelledStatement = withLoc(t.labeledStatement);
 const createVariableDeclaration = withLoc(t.variableDeclaration);
 const createFunctionDeclaration = withLoc(t.functionDeclaration);
-const _createWhileStatement = withLoc(t.whileStatement);
+const createWhileStatement = withLoc(t.whileStatement);
+const createDoWhileStatement = withLoc(t.doWhileStatement);
+const createSwitchStatement = withLoc(t.switchStatement);
+const createIfStatement = withLoc(t.ifStatement);
+const createForStatement = withLoc(t.forStatement);
+const createForOfStatement = withLoc(t.forOfStatement);
+const createForInStatement = withLoc(t.forInStatement);
 const createTaggedTemplateExpression = withLoc(t.taggedTemplateExpression);
 const createLogicalExpression = withLoc(t.logicalExpression);
 const createSequenceExpression = withLoc(t.sequenceExpression);
@@ -1710,6 +1734,10 @@ const createJsxText = withLoc(t.jsxText);
 const createJsxClosingElement = withLoc(t.jsxClosingElement);
 const createJsxOpeningElement = withLoc(t.jsxOpeningElement);
 const createStringLiteral = withLoc(t.stringLiteral);
+const createThrowStatement = withLoc(t.throwStatement);
+const createTryStatement = withLoc(t.tryStatement);
+const createBreakStatement = withLoc(t.breakStatement);
+const createContinueStatement = withLoc(t.continueStatement);
 
 function createHookGuard(
   guard: ExternalFunction,
@@ -2524,6 +2552,9 @@ function codegenInstructionValue(
       );
     }
   }
+  if (instrValue.loc != null && instrValue.loc != GeneratedSource) {
+    value.loc = instrValue.loc;
+  }
   return value;
 }
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/align-scopes-nested-block-structure.expect.md

@@ -129,6 +129,7 @@ function useFoo(t0) {
         t1 = null;
         break bb0;
       }
+
       if (cond2) {
         mutate(s);
       }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/align-scopes-reactive-scope-overlaps-label.expect.md

@@ -43,6 +43,7 @@ function useFoo(t0) {
   if ($[0] !== cond || $[1] !== value) {
     bb0: {
       items = [];
+
       if (cond) {
         break bb0;
       }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/align-scopes-trycatch-nested-overlapping-range.expect.md

@@ -39,6 +39,7 @@ function Foo() {
       if (cond) {
         thing = makeObject_Primitives();
       }
+
       if (CONST_TRUE) {
         mutate(thing);
       }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/assignment-in-nested-if.expect.md

@@ -23,6 +23,7 @@ import { c as _c } from "react/compiler-runtime";
 function useBar(props) {
   const $ = _c(1);
   let z;
+
   if (props.a) {
     if (props.b) {
       let t0;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/capture-param-mutate.expect.md

@@ -67,6 +67,7 @@ function getNativeLogFunction(level) {
       ) {
         logLevel = LOG_LEVELS.warn;
       }
+
       if (global.__inspectorLog) {
         global.__inspectorLog(
           INSPECTOR_LEVELS[logLevel],
@@ -75,6 +76,7 @@ function getNativeLogFunction(level) {
           INSPECTOR_FRAMES_TO_SKIP,
         );
       }
+
       if (groupStack.length) {
         str = groupFormat("", str);
       }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/capture-ref-for-later-mutation.expect.md

@@ -47,9 +47,7 @@ function useKeyCommand() {
     };
 
     const moveLeft = { handler: handleKey("left") };
-
     const moveRight = { handler: handleKey("right") };
-
     t0 = [moveLeft, moveRight];
     $[0] = t0;
   } else {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/chained-assignment-context-variable.expect.md

@@ -33,7 +33,6 @@ function Component() {
   let y;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
     y = x = {};
-
     const foo = () => {
       x = makeArray();
     };

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/conditional-on-mutable.expect.md

@@ -44,6 +44,7 @@ function ComponentA(props) {
     if (b) {
       a.push(props.p0);
     }
+
     if (props.p1) {
       b.push(props.p2);
     }
@@ -68,6 +69,7 @@ function ComponentB(props) {
     if (mayMutate(b)) {
       a.push(props.p0);
     }
+
     if (props.p1) {
       b.push(props.p2);
     }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/conditional-set-state-in-render.expect.md

@@ -34,6 +34,7 @@ function Component(props) {
   const foo = () => {
     setX(1);
   };
+
   if (props.cond) {
     setX(2);
     foo();