From 340820784903ae033b116e564a1eda90a1f3d745 Mon Sep 17 00:00:00 2001
From: Josh Story <story@hey.com>
Date: Thu, 15 Feb 2024 19:23:50 -0800
Subject: [PATCH] Retain component stacks when rethrowing an error that already
 generated a component stack

---
 .../src/ReactCapturedValue.js                 |  28 ++++-
 .../src/ReactFiberBeginWork.js                |   6 +-
 .../src/__tests__/ReactErrorStacks-test.js    | 103 ++++++++++++++++++
 3 files changed, 129 insertions(+), 8 deletions(-)
 create mode 100644 packages/react-reconciler/src/__tests__/ReactErrorStacks-test.js

diff --git a/packages/react-reconciler/src/ReactCapturedValue.js b/packages/react-reconciler/src/ReactCapturedValue.js
index feab8b590a7f5..b4c1b88dccf0e 100644
--- a/packages/react-reconciler/src/ReactCapturedValue.js
+++ b/packages/react-reconciler/src/ReactCapturedValue.js
@@ -11,8 +11,10 @@ import type {Fiber} from './ReactInternalTypes';
 
 import {getStackByFiberInDevAndProd} from './ReactFiberComponentStack';
 
+const CapturedStacks: WeakMap<any, string> = new WeakMap();
+
 export type CapturedValue<T> = {
-  value: T,
+  +value: T,
   source: Fiber | null,
   stack: string | null,
   digest: string | null,
@@ -24,19 +26,35 @@ export function createCapturedValueAtFiber<T>(
 ): CapturedValue<T> {
   // If the value is an error, call this function immediately after it is thrown
   // so the stack is accurate.
+  let stack;
+  if (typeof value === 'object' && value !== null) {
+    const capturedStack = CapturedStacks.get(value);
+    if (typeof capturedStack === 'string') {
+      stack = capturedStack;
+    } else {
+      stack = getStackByFiberInDevAndProd(source);
+      CapturedStacks.set(value, stack);
+    }
+  } else {
+    stack = getStackByFiberInDevAndProd(source);
+  }
+
   return {
     value,
     source,
-    stack: getStackByFiberInDevAndProd(source),
+    stack,
     digest: null,
   };
 }
 
-export function createCapturedValue<T>(
-  value: T,
+export function createCapturedValueFromError(
+  value: Error,
   digest: ?string,
   stack: ?string,
-): CapturedValue<T> {
+): CapturedValue<Error> {
+  if (typeof stack === 'string') {
+    CapturedStacks.set(value, stack);
+  }
   return {
     value,
     source: null,
diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.js b/packages/react-reconciler/src/ReactFiberBeginWork.js
index 8a7ed49f322d7..a272da942e665 100644
--- a/packages/react-reconciler/src/ReactFiberBeginWork.js
+++ b/packages/react-reconciler/src/ReactFiberBeginWork.js
@@ -264,7 +264,7 @@ import {
 import {enqueueConcurrentRenderForLane} from './ReactFiberConcurrentUpdates';
 import {pushCacheProvider, CacheContext} from './ReactFiberCacheComponent';
 import {
-  createCapturedValue,
+  createCapturedValueFromError,
   createCapturedValueAtFiber,
   type CapturedValue,
 } from './ReactCapturedValue';
@@ -2804,7 +2804,7 @@ function updateDehydratedSuspenseComponent(
           );
         }
         (error: any).digest = digest;
-        capturedValue = createCapturedValue<mixed>(error, digest, stack);
+        capturedValue = createCapturedValueFromError(error, digest, stack);
       }
       return retrySuspenseComponentWithoutHydrating(
         current,
@@ -2941,7 +2941,7 @@ function updateDehydratedSuspenseComponent(
       pushPrimaryTreeSuspenseHandler(workInProgress);
 
       workInProgress.flags &= ~ForceClientRender;
-      const capturedValue = createCapturedValue<mixed>(
+      const capturedValue = createCapturedValueFromError(
         new Error(
           'There was an error while hydrating this Suspense boundary. ' +
             'Switched to client rendering.',
diff --git a/packages/react-reconciler/src/__tests__/ReactErrorStacks-test.js b/packages/react-reconciler/src/__tests__/ReactErrorStacks-test.js
new file mode 100644
index 0000000000000..6e6a9d471c4ef
--- /dev/null
+++ b/packages/react-reconciler/src/__tests__/ReactErrorStacks-test.js
@@ -0,0 +1,103 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @emails react-core
+ * @jest-environment node
+ */
+'use strict';
+
+let React;
+let ReactNoop;
+let waitForAll;
+
+describe('ReactFragment', () => {
+  beforeEach(function () {
+    jest.resetModules();
+
+    React = require('react');
+    ReactNoop = require('react-noop-renderer');
+    const InternalTestUtils = require('internal-test-utils');
+    waitForAll = InternalTestUtils.waitForAll;
+  });
+
+  function componentStack(components) {
+    return components
+      .map(component => `\n    in ${component} (at **)`)
+      .join('');
+  }
+
+  function normalizeCodeLocInfo(str) {
+    return (
+      str &&
+      str.replace(/\n +(?:at|in) ([\S]+)[^\n]*/g, function (m, name) {
+        return '\n    in ' + name + ' (at **)';
+      })
+    );
+  }
+
+  it('retains component stacks when rethrowing an error', async () => {
+    function Foo() {
+      return (
+        <RethrowingBoundary>
+          <Bar />
+        </RethrowingBoundary>
+      );
+    }
+    function Bar() {
+      return <SomethingThatErrors />;
+    }
+    function SomethingThatErrors() {
+      throw new Error('uh oh');
+    }
+
+    class RethrowingBoundary extends React.Component {
+      static getDerivedStateFromError(error) {
+        throw error;
+      }
+
+      render() {
+        return this.props.children;
+      }
+    }
+
+    const errors = [];
+    class CatchingBoundary extends React.Component {
+      constructor() {
+        super();
+        this.state = {};
+      }
+      static getDerivedStateFromError(error) {
+        return {errored: true};
+      }
+      componentDidCatch(err, errInfo) {
+        errors.push(err.message, normalizeCodeLocInfo(errInfo.componentStack));
+      }
+      render() {
+        if (this.state.errored) {
+          return null;
+        }
+        return this.props.children;
+      }
+    }
+
+    ReactNoop.render(
+      <CatchingBoundary>
+        <Foo />
+      </CatchingBoundary>,
+    );
+    await waitForAll([]);
+    expect(errors).toEqual([
+      'uh oh',
+      componentStack([
+        'SomethingThatErrors',
+        'Bar',
+        'RethrowingBoundary',
+        'Foo',
+        'CatchingBoundary',
+      ]),
+    ]);
+  });
+});