[Flight] Encode enclosing line/column numbers and use it to align the fake function

packages/react-client/src/ReactFlightClient.js

@@ -2226,6 +2226,8 @@ function createFakeFunction<T>(
   sourceMap: null | string,
   line: number,
   col: number,
+  enclosingLine: number,
+  enclosingCol: number,
   environmentName: string,
 ): FakeFunction<T> {
   // This creates a fake copy of a Server Module. It represents a module that has already
@@ -2243,26 +2245,104 @@ function createFakeFunction<T>(
   // This allows us to use the original source map as the source map of this fake file to
   // point to the original source.
   let code;
-  if (line <= 1) {
-    const minSize = encodedName.length + 7;
+  // Normalize line/col to zero based.
+  if (enclosingLine < 1) {
+    enclosingLine = 0;
+  } else {
+    enclosingLine--;
+  }
+  if (enclosingCol < 1) {
+    enclosingCol = 0;
+  } else {
+    enclosingCol--;
+  }
+  if (line < 1) {
+    line = 0;
+  } else {
+    line--;
+  }
+  if (col < 1) {
+    col = 0;
+  } else {
+    col--;
+  }
+  if (line < enclosingLine || (line === enclosingLine && col < enclosingCol)) {
+    // Protection against invalid enclosing information. Should not happen.
+    enclosingLine = 0;
+    enclosingCol = 0;
+  }
+  if (line < 1) {
+    // Fit everything on the first line.
+    const minCol = encodedName.length + 3;
+    let enclosingColDistance = enclosingCol - minCol;
+    if (enclosingColDistance < 0) {
+      enclosingColDistance = 0;
+    }
+    let colDistance = col - enclosingColDistance - minCol - 3;
+    if (colDistance < 0) {
+      colDistance = 0;
+    }
     code =
       '({' +
       encodedName +
-      ':_=>' +
-      ' '.repeat(col < minSize ? 0 : col - minSize) +
-      '_()})\n' +
-      comment;
+      ':' +
+      ' '.repeat(enclosingColDistance) +
+      '_=>' +
+      ' '.repeat(colDistance) +
+      '_()})';
+  } else if (enclosingLine < 1) {
+    // Fit just the enclosing function on the first line.
+    const minCol = encodedName.length + 3;
+    let enclosingColDistance = enclosingCol - minCol;
+    if (enclosingColDistance < 0) {
+      enclosingColDistance = 0;
+    }
+    code =
+      '({' +
+      encodedName +
+      ':' +
+      ' '.repeat(enclosingColDistance) +
+      '_=>' +
+      '\n'.repeat(line - enclosingLine) +
+      ' '.repeat(col) +
+      '_()})';
+  } else if (enclosingLine === line) {
+    // Fit the enclosing function and callsite on same line.
+    let colDistance = col - enclosingCol - 3;
+    if (colDistance < 0) {
+      colDistance = 0;
+    }
+    code =
+      '\n'.repeat(enclosingLine - 1) +
+      '({' +
+      encodedName +
+      ':\n' +
+      ' '.repeat(enclosingCol) +
+      '_=>' +
+      ' '.repeat(colDistance) +
+      '_()})';
   } else {
+    // This is the ideal because we can always encode any position.
     code =
-      comment +
-      '\n'.repeat(line - 2) +
+      '\n'.repeat(enclosingLine - 1) +
       '({' +
       encodedName +
-      ':_=>\n' +
-      ' '.repeat(col < 1 ? 0 : col - 1) +
+      ':\n' +
+      ' '.repeat(enclosingCol) +
+      '_=>' +
+      '\n'.repeat(line - enclosingLine) +
+      ' '.repeat(col) +
       '_()})';
   }
 
+  if (enclosingLine < 1) {
+    // If the function starts at the first line, we append the comment after.
+    code = code + '\n' + comment;
+  } else {
+    // Otherwise we prepend the comment on the first line.
+    code = comment + code;
+  }
+
   if (filename.startsWith('/')) {
     // If the filename starts with `/` we assume that it is a file system file
     // rather than relative to the current host. Since on the server fully qualified
@@ -2320,7 +2400,7 @@ function buildFakeCallStack<T>(
     const frameKey = frame.join('-') + '-' + environmentName;
     let fn = fakeFunctionCache.get(frameKey);
     if (fn === undefined) {
-      const [name, filename, line, col] = frame;
+      const [name, filename, line, col, enclosingLine, enclosingCol] = frame;
       const findSourceMapURL = response._debugFindSourceMapURL;
       const sourceMap = findSourceMapURL
         ? findSourceMapURL(filename, environmentName)
@@ -2331,6 +2411,8 @@ function buildFakeCallStack<T>(
         sourceMap,
         line,
         col,
+        enclosingLine,
+        enclosingCol,
         environmentName,
       );
       // TODO: This cache should technically live on the response since the _debugFindSourceMapURL

packages/react-client/src/ReactFlightReplyClient.js

@@ -1350,10 +1350,11 @@ function parseStackLocation(error: Error): null | ReactCallSite {
   if (filename === '<anonymous>') {
     filename = '';
   }
+  // This is really the enclosingLine/Column.
   const line = +(parsed[3] || parsed[6]);
   const col = +(parsed[4] || parsed[7]);
 
-  return [name, filename, line, col];
+  return [name, filename, line, col, line, col];
 }
 
 export function createServerReference<A: Iterable<any>, T>(

packages/react-server/src/ReactFlightStackConfigV8.js

@@ -63,7 +63,7 @@ function collectStackTrace(
       // Skip everything after the bottom frame since it'll be internals.
       break;
     } else if (callSite.isNative()) {
-      result.push([name, '', 0, 0]);
+      result.push([name, '', 0, 0, 0, 0]);
     } else {
       // We encode complex function calls as if they're part of the function
       // name since we cannot simulate the complex ones and they look the same
@@ -88,7 +88,17 @@ function collectStackTrace(
       }
       const line = callSite.getLineNumber() || 0;
       const col = callSite.getColumnNumber() || 0;
-      result.push([name, filename, line, col]);
+      const enclosingLine: number =
+        // $FlowFixMe[prop-missing]
+        typeof callSite.getEnclosingLineNumber === 'function'
+          ? (callSite: any).getEnclosingLineNumber() || 0
+          : 0;
+      const enclosingCol: number =
+        // $FlowFixMe[prop-missing]
+        typeof callSite.getEnclosingColumnNumber === 'function'
+          ? (callSite: any).getEnclosingColumnNumber() || 0
+          : 0;
+      result.push([name, filename, line, col, enclosingLine, enclosingCol]);
     }
   }
   // At the same time we generate a string stack trace just in case someone
@@ -179,7 +189,7 @@ export function parseStackTrace(
     }
     const line = +(parsed[3] || parsed[6]);
     const col = +(parsed[4] || parsed[7]);
-    parsedFrames.push([name, filename, line, col]);
+    parsedFrames.push([name, filename, line, col, 0, 0]);
   }
   return parsedFrames;
 }

packages/shared/ReactTypes.js

@@ -186,6 +186,8 @@ export type ReactCallSite = [
   string, // file name TODO: model nested eval locations as nested arrays
   number, // line number
   number, // column number
+  number, // enclosing line number
+  number, // enclosing column number
 ];
 
 export type ReactStackTrace = Array<ReactCallSite>;