Repository: facebook/react
Pull request: [Flight] Source Map Actions in Reference Node Loader Transforms

Follow up to #30741.

This is just for the reference Webpack implementation.

If there is a source map associated with a Node ESM loader, we generate new source map entries for every `registerServerReference` call.

To avoid messing too much with it, this doesn't rewrite the original mappings. It just reads them while finding each of the exports in the original mappings. We need to read all since whatever we append at the end is relative. Then we just generate new appended entries at the end.

For the location I picked the location of the local name identifier. Since that's the name of the function and that gives us a source map name index. It means it jumps to the name rather than the beginning of the function declaration. It could be made more clever like finding a local function definition if it is reexported. We could also point to the line/column of the function declaration rather than the identifier but point to the name index of the identifier name.

Now jumping to definition works in the fixture.

<img width="574" alt="Screenshot 2024-08-20 at 2 49 07 PM" src="https://github.com/user-attachments/assets/7710f0e6-2cee-4aad-8d4c-ae985f8289eb">

Unfortunately this technique doesn't seem to work in Firefox nor Safari. They don't apply the source map for jumping to the definition.

Changed files:
filename: fixtures/flight/package.json
patch: @@ -49,6 +49,7 @@
     "react-dev-utils": "^12.0.1",
     "react-dom": "experimental",
     "react-refresh": "^0.11.0",
+    "react-server-dom-webpack": "experimental",
     "resolve": "^1.20.0",
     "resolve-url-loader": "^4.0.0",
     "sass-loader": "^12.3.0",

filename: packages/react-server-dom-webpack/package.json
patch: @@ -106,6 +106,7 @@
   },
   "dependencies": {
     "acorn-loose": "^8.3.0",
-    "neo-async": "^2.6.1"
+    "neo-async": "^2.6.1",
+    "webpack-sources": "^3.2.0"
   }
 }

filename: packages/react-server-dom-webpack/src/ReactFlightWebpackNodeLoader.js
patch: @@ -9,6 +9,9 @@
 
 import * as acorn from 'acorn-loose';
 
+import readMappings from 'webpack-sources/lib/helpers/readMappings.js';
+import createMappingsSerializer from 'webpack-sources/lib/helpers/createMappingsSerializer.js';
+
 type ResolveContext = {
   conditions: Array<string>,
   parentURL: string | void,
@@ -95,45 +98,102 @@ export async function getSource(
   return defaultGetSource(url, context, defaultGetSource);
 }
 
-function addLocalExportedNames(names: Map<string, string>, node: any) {
+type ExportedEntry = {
+  localName: string,
+  exportedName: string,
+  type: null | string,
+  loc: {
+    start: {line: number, column: number},
+    end: {line: number, column: number},
+  },
+  originalLine: number,
+  originalColumn: number,
+  originalSource: number,
+  nameIndex: number,
+};
+
+function addExportedEntry(
+  exportedEntries: Array<ExportedEntry>,
+  localNames: Set<string>,
+  localName: string,
+  exportedName: string,
+  type: null | 'function',
+  loc: {
+    start: {line: number, column: number},
+    end: {line: number, column: number},
+  },
+) {
+  if (localNames.has(localName)) {
+    // If the same local name is exported more than once, we only need one of the names.
+    return;
+  }
+  exportedEntries.push({
+    localName,
+    exportedName,
+    type,
+    loc,
+    originalLine: -1,
+    originalColumn: -1,
+    originalSource: -1,
+    nameIndex: -1,
+  });
+}
+
+function addLocalExportedNames(
+  exportedEntries: Array<ExportedEntry>,
+  localNames: Set<string>,
+  node: any,
+) {
   switch (node.type) {
     case 'Identifier':
-      names.set(node.name, node.name);
+      addExportedEntry(
+        exportedEntries,
+        localNames,
+        node.name,
+        node.name,
+        null,
+        node.loc,
+      );
       return;
     case 'ObjectPattern':
       for (let i = 0; i < node.properties.length; i++)
-        addLocalExportedNames(names, node.properties[i]);
+        addLocalExportedNames(exportedEntries, localNames, node.properties[i]);
       return;
     case 'ArrayPattern':
       for (let i = 0; i < node.elements.length; i++) {
         const element = node.elements[i];
-        if (element) addLocalExportedNames(names, element);
+        if (element)
+          addLocalExportedNames(exportedEntries, localNames, element);
       }
       return;
     case 'Property':
-      addLocalExportedNames(names, node.value);
+      addLocalExportedNames(exportedEntries, localNames, node.value);
       return;
     case 'AssignmentPattern':
-      addLocalExportedNames(names, node.left);
+      addLocalExportedNames(exportedEntries, localNames, node.left);
       return;
     case 'RestElement':
-      addLocalExportedNames(names, node.argument);
+      addLocalExportedNames(exportedEntries, localNames, node.argument);
       return;
     case 'ParenthesizedExpression':
-      addLocalExportedNames(names, node.expression);
+      addLocalExportedNames(exportedEntries, localNames, node.expression);
       return;
   }
 }
 
 function transformServerModule(
   source: string,
-  body: any,
+  program: any,
   url: string,
+  sourceMap: any,
   loader: LoadFunction,
 ): string {
-  // If the same local name is exported more than once, we only need one of the names.
-  const localNames: Map<string, string> = new Map();
-  const localTypes: Map<string, string> = new Map();
+  const body = program.body;
+
+  // This entry list needs to be in source location order.
+  const exportedEntries: Array<ExportedEntry> = [];
+  // Dedupe set.
+  const localNames: Set<string> = new Set();
 
   for (let i = 0; i < body.length; i++) {
     const node = body[i];
@@ -143,11 +203,24 @@ function transformServerModule(
         break;
       case 'ExportDefaultDeclaration':
         if (node.declaration.type === 'Identifier') {
-          localNames.set(node.declaration.name, 'default');
+          addExportedEntry(
+            exportedEntries,
+            localNames,
+            node.declaration.name,
+            'default',
+            null,
+            node.declaration.loc,
+          );
         } else if (node.declaration.type === 'FunctionDeclaration') {
           if (node.declaration.id) {
-            localNames.set(node.declaration.id.name, 'default');
-            localTypes.set(node.declaration.id.name, 'function');
+            addExportedEntry(
+              exportedEntries,
+              localNames,
+              node.declaration.id.name,
+              'default',
+              'function',
+              node.declaration.id.loc,
+            );
           } else {
             // TODO: This needs to be rewritten inline because it doesn't have a local name.
           }
@@ -158,41 +231,230 @@ function transformServerModule(
           if (node.declaration.type === 'VariableDeclaration') {
             const declarations = node.declaration.declarations;
             for (let j = 0; j < declarations.length; j++) {
-              addLocalExportedNames(localNames, declarations[j].id);
+              addLocalExportedNames(
+                exportedEntries,
+                localNames,
+                declarations[j].id,
+              );
             }
           } else {
             const name = node.declaration.id.name;
-            localNames.set(name, name);
-            if (node.declaration.type === 'FunctionDeclaration') {
-              localTypes.set(name, 'function');
-            }
+            addExportedEntry(
+              exportedEntries,
+              localNames,
+              name,
+              name,
+
+              node.declaration.type === 'FunctionDeclaration'
+                ? 'function'
+                : null,
+              node.declaration.id.loc,
+            );
           }
         }
         if (node.specifiers) {
           const specifiers = node.specifiers;
           for (let j = 0; j < specifiers.length; j++) {
             const specifier = specifiers[j];
-            localNames.set(specifier.local.name, specifier.exported.name);
+            addExportedEntry(
+              exportedEntries,
+              localNames,
+              specifier.local.name,
+              specifier.exported.name,
+              null,
+              specifier.local.loc,
+            );
           }
         }
         continue;
     }
   }
-  if (localNames.size === 0) {
-    return source;
-  }
-  let newSrc = source + '\n\n;';
-  newSrc +=
-    'import {registerServerReference} from "react-server-dom-webpack/server";\n';
-  localNames.forEach(function (exported, local) {
-    if (localTypes.get(local) !== 'function') {
-      // We first check if the export is a function and if so annotate it.
-      newSrc += 'if (typeof ' + local + ' === "function") ';
+
+  let mappings =
+    sourceMap && typeof sourceMap.mappings === 'string'
+      ? sourceMap.mappings
+      : '';
+  let newSrc = source;
+
+  if (exportedEntries.length > 0) {
+    let lastSourceIndex = 0;
+    let lastOriginalLine = 0;
+    let lastOriginalColumn = 0;
+    let lastNameIndex = 0;
+    let sourceLineCount = 0;
+    let lastMappedLine = 0;
+
+    if (sourceMap) {
+      // We iterate source mapping entries and our matched exports in parallel to source map
+      // them to their original location.
+      let nextEntryIdx = 0;
+      let nextEntryLine = exportedEntries[nextEntryIdx].loc.start.line;
+      let nextEntryColumn = exportedEntries[nextEntryIdx].loc.start.column;
+      readMappings(
+        mappings,
+        (
+          generatedLine: number,
+          generatedColumn: number,
+          sourceIndex: number,
+          originalLine: number,
+          originalColumn: number,
+          nameIndex: number,
+        ) => {
+          if (
+            generatedLine > nextEntryLine ||
+            (generatedLine === nextEntryLine &&
+              generatedColumn > nextEntryColumn)
+          ) {
+            // We're past the entry which means that the best match we have is the previous entry.
+            if (lastMappedLine === nextEntryLine) {
+              // Match
+              exportedEntries[nextEntryIdx].originalLine = lastOriginalLine;
+              exportedEntries[nextEntryIdx].originalColumn = lastOriginalColumn;
+              exportedEntries[nextEntryIdx].originalSource = lastSourceIndex;
+              exportedEntries[nextEntryIdx].nameIndex = lastNameIndex;
+            } else {
+              // Skip if we didn't have any mappings on the exported line.
+            }
+            nextEntryIdx++;
+            if (nextEntryIdx < exportedEntries.length) {
+              nextEntryLine = exportedEntries[nextEntryIdx].loc.start.line;
+              nextEntryColumn = exportedEntries[nextEntryIdx].loc.start.column;
+            } else {
+              nextEntryLine = -1;
+              nextEntryColumn = -1;
+            }
+          }
+          lastMappedLine = generatedLine;
+          if (sourceIndex > -1) {
+            lastSourceIndex = sourceIndex;
+          }
+          if (originalLine > -1) {
+            lastOriginalLine = originalLine;
+          }
+          if (originalColumn > -1) {
+            lastOriginalColumn = originalColumn;
+          }
+          if (nameIndex > -1) {
+            lastNameIndex = nameIndex;
+          }
+        },
+      );
+      if (nextEntryIdx < exportedEntries.length) {
+        if (lastMappedLine === nextEntryLine) {
+          // Match
+          exportedEntries[nextEntryIdx].originalLine = lastOriginalLine;
+          exportedEntries[nextEntryIdx].originalColumn = lastOriginalColumn;
+          exportedEntries[nextEntryIdx].originalSource = lastSourceIndex;
+          exportedEntries[nextEntryIdx].nameIndex = lastNameIndex;
+        }
+      }
+
+      for (
+        let lastIdx = mappings.length - 1;
+        lastIdx >= 0 && mappings[lastIdx] === ';';
+        lastIdx--
+      ) {
+        // If the last mapped lines don't contain any segments, we don't get a callback from readMappings
+        // so we need to pad the number of mapped lines, with one for each empty line.
+        lastMappedLine++;
+      }
+
+      sourceLineCount = program.loc.end.line;
+      if (sourceLineCount < lastMappedLine) {
+        throw new Error(
+          'The source map has more mappings than there are lines.',
+        );
+      }
+      // If the original source string had more lines than there are mappings in the source map.
+      // Add some extra padding of unmapped lines so that any lines that we add line up.
+      for (
+        let extraLines = sourceLineCount - lastMappedLine;
+        extraLines > 0;
+        extraLines--
+      ) {
+        mappings += ';';
+      }
+    } else {
+      // If a file doesn't have a source map then we generate a blank source map that just
+      // contains the original content and segments pointing to the original lines.
+      sourceLineCount = 1;
+      let idx = -1;
+      while ((idx = source.indexOf('\n', idx + 1)) !== -1) {
+        sourceLineCount++;
+      }
+      mappings = 'AAAA' + ';AACA'.repeat(sourceLineCount - 1);
+      sourceMap = {
+        version: 3,
+        sources: [url],
+        sourcesContent: [source],
+        mappings: mappings,
+        sourceRoot: '',
+      };
+      lastSourceIndex = 0;
+      lastOriginalLine = sourceLineCount;
+      lastOriginalColumn = 0;
+      lastNameIndex = -1;
+      lastMappedLine = sourceLineCount;
+
+      for (let i = 0; i < exportedEntries.length; i++) {
+        // Point each entry to original location.
+        const entry = exportedEntries[i];
+        entry.originalSource = 0;
+        entry.originalLine = entry.loc.start.line;
+        // We use column zero since we do the short-hand line-only source maps above.
+        entry.originalColumn = 0; // entry.loc.start.column;
+      }
     }
-    newSrc += 'registerServerReference(' + local + ',';
-    newSrc += JSON.stringify(url) + ',';
-    newSrc += JSON.stringify(exported) + ');\n';
-  });
+
+    newSrc += '\n\n;';
+    newSrc +=
+      'import {registerServerReference} from "react-server-dom-webpack/server";\n';
+    if (mappings) {
+      mappings += ';;';
+    }
+
+    const createMapping = createMappingsSerializer();
+
+    // Create an empty mapping pointing to where we last left off to reset the counters.
+    let generatedLine = 1;
+    createMapping(
+      generatedLine,
+      0,
+      lastSourceIndex,
+      lastOriginalLine,
+      lastOriginalColumn,
+      lastNameIndex,
+    );
+    for (let i = 0; i < exportedEntries.length; i++) {
+      const entry = exportedEntries[i];
+      generatedLine++;
+      if (entry.type !== 'function') {
+        // We first check if the export is a function and if so annotate it.
+        newSrc += 'if (typeof ' + entry.localName + ' === "function") ';
+      }
+      newSrc += 'registerServerReference(' + entry.localName + ',';
+      newSrc += JSON.stringify(url) + ',';
+      newSrc += JSON.stringify(entry.exportedName) + ');\n';
+
+      mappings += createMapping(
+        generatedLine,
+        0,
+        entry.originalSource,
+        entry.originalLine,
+        entry.originalColumn,
+        entry.nameIndex,
+      );
+    }
+  }
+
+  if (sourceMap) {
+    // Override with an new mappings and serialize an inline source map.
+    sourceMap.mappings = mappings;
+    newSrc +=
+      '//# sourceMappingURL=data:application/json;charset=utf-8;base64,' +
+      Buffer.from(JSON.stringify(sourceMap)).toString('base64');
+  }
+
   return newSrc;
 }
 
@@ -307,10 +569,13 @@ async function parseExportNamesInto(
 }
 
 async function transformClientModule(
-  body: any,
+  program: any,
   url: string,
+  sourceMap: any,
   loader: LoadFunction,
 ): Promise<string> {
+  const body = program.body;
+
   const names: Array<string> = [];
 
   await parseExportNamesInto(body, names, url, loader);
@@ -351,6 +616,9 @@ async function transformClientModule(
     newSrc += JSON.stringify(url) + ',';
     newSrc += JSON.stringify(name) + ');\n';
   }
+
+  // TODO: Generate source maps for Client Reference functions so they can point to their
+  // original locations.
   return newSrc;
 }
 
@@ -391,12 +659,36 @@ async function transformModuleIfNeeded(
     return source;
   }
 
-  let body;
+  let sourceMappingURL = null;
+  let sourceMappingStart = 0;
+  let sourceMappingEnd = 0;
+  let sourceMappingLines = 0;
+
+  let program;
   try {
-    body = acorn.parse(source, {
+    program = acorn.parse(source, {
       ecmaVersion: '2024',
       sourceType: 'module',
-    }).body;
+      locations: true,
+      onComment(
+        block: boolean,
+        text: string,
+        start: number,
+        end: number,
+        startLoc: {line: number, column: number},
+        endLoc: {line: number, column: number},
+      ) {
+        if (
+          text.startsWith('# sourceMappingURL=') ||
+          text.startsWith('@ sourceMappingURL=')
+        ) {
+          sourceMappingURL = text.slice(19);
+          sourceMappingStart = start;
+          sourceMappingEnd = end;
+          sourceMappingLines = endLoc.line - startLoc.line;
+        }
+      },
+    });
   } catch (x) {
     // eslint-disable-next-line react-internal/no-production-logging
     console.error('Error parsing %s %s', url, x.message);
@@ -405,6 +697,8 @@ async function transformModuleIfNeeded(
 
   let useClient = false;
   let useServer = false;
+
+  const body = program.body;
   for (let i = 0; i < body.length; i++) {
     const node = body[i];
     if (node.type !== 'ExpressionStatement' || !node.directive) {
@@ -428,11 +722,38 @@ async function transformModuleIfNeeded(
     );
   }
 
+  let sourceMap = null;
+  if (sourceMappingURL) {
+    const sourceMapResult = await loader(
+      sourceMappingURL,
+      // $FlowFixMe
+      {
+        format: 'json',
+        conditions: [],
+        importAssertions: {type: 'json'},
+        importAttributes: {type: 'json'},
+      },
+      loader,
+    );
+    const sourceMapString =
+      typeof sourceMapResult.source === 'string'
+        ? sourceMapResult.source
+        : // $FlowFixMe
+          sourceMapResult.source.toString('utf8');
+    sourceMap = JSON.parse(sourceMapString);
+
+    // Strip the source mapping comment. We'll re-add it below if needed.
+    source =
+      source.slice(0, sourceMappingStart) +
+      '\n'.repeat(sourceMappingLines) +
+      source.slice(sourceMappingEnd);
+  }
+
   if (useClient) {
-    return transformClientModule(body, url, loader);
+    return transformClientModule(program, url, sourceMap, loader);
   }
 
-  return transformServerModule(source, body, url, loader);
+  return transformServerModule(source, program, url, sourceMap, loader);
 }
 
 export async function transformSource(

filename: scripts/rollup/modules.js
patch: @@ -22,6 +22,9 @@ const importSideEffects = Object.freeze({
   'react-dom': HAS_NO_SIDE_EFFECTS_ON_IMPORT,
   url: HAS_NO_SIDE_EFFECTS_ON_IMPORT,
   ReactNativeInternalFeatureFlags: HAS_NO_SIDE_EFFECTS_ON_IMPORT,
+  'webpack-sources/lib/helpers/createMappingsSerializer.js':
+    HAS_NO_SIDE_EFFECTS_ON_IMPORT,
+  'webpack-sources/lib/helpers/readMappings.js': HAS_NO_SIDE_EFFECTS_ON_IMPORT,
 });
 
 // Bundles exporting globals that other modules rely on.

filename: yarn.lock
patch: @@ -11347,11 +11347,6 @@ lodash.omitby@4.6.0:
   resolved "https://registry.yarnpkg.com/lodash.omitby/-/lodash.omitby-4.6.0.tgz#5c15ff4754ad555016b53c041311e8f079204791"
   integrity sha1-XBX/R1StVVAWtTwEExHo8HkgR5E=
 
-lodash.throttle@^4.1.1:
-  version "4.1.1"
-  resolved "https://registry.yarnpkg.com/lodash.throttle/-/lodash.throttle-4.1.1.tgz#c23e91b710242ac70c37f1e1cda9274cc39bf2f4"
-  integrity sha1-wj6RtxAkKscMN/HhzaknTMOb8vQ=
-
 lodash.truncate@^4.4.2:
   version "4.4.2"
   resolved "https://registry.yarnpkg.com/lodash.truncate/-/lodash.truncate-4.4.2.tgz#5a350da0b1113b837ecfffd5812cbe58d6eae193"
@@ -11603,11 +11598,6 @@ memfs@^3.4.3:
   resolved "https://registry.yarnpkg.com/memoize-one/-/memoize-one-5.2.1.tgz#8337aa3c4335581839ec01c3d594090cebe8f00e"
   integrity sha512-zYiwtZUcYyXKo/np96AGZAckk+FWWsUdJ3cHGGmld7+AhvcWmQyGCYUh1hc4Q/pkOhb65dQR/pqCyK0cOaHz4Q==
 
-memoize-one@^3.1.1:
-  version "3.1.1"
-  resolved "https://registry.yarnpkg.com/memoize-one/-/memoize-one-3.1.1.tgz#ef609811e3bc28970eac2884eece64d167830d17"
-  integrity sha512-YqVh744GsMlZu6xkhGslPSqSurOv6P+kLN2J3ysBZfagLcL5FdRK/0UpgLoL8hwjjEvvAVkjJZyFP+1T6p1vgA==
-
 memoize-one@^5.1.1:
   version "5.1.1"
   resolved "https://registry.yarnpkg.com/memoize-one/-/memoize-one-5.1.1.tgz#047b6e3199b508eaec03504de71229b8eb1d75c0"
@@ -16499,7 +16489,7 @@ webpack-merge@^5.7.3:
     clone-deep "^4.0.1"
     wildcard "^2.0.0"
 
-webpack-sources@^3.2.3:
+webpack-sources@^3.2.0, webpack-sources@^3.2.3:
   version "3.2.3"
   resolved "https://registry.yarnpkg.com/webpack-sources/-/webpack-sources-3.2.3.tgz#2d4daab8451fd4b240cc27055ff6a0c2ccea0cde"
   integrity sha512-/DyMEOrDgLKKIG0fmvtz+4dUX/3Ghozwgm6iPp8KRhvn+eQf9+Q7GWxVNMk3+uCPWfdXYC4ExGBckIXdFEfH1w==