stash the component stack on the thrown value and reuse (facebook#25790)

ErrorBoundaries are currently not fully composable. The reason is if you
decide your boundary cannot handle a particular error and rethrow it to
higher boundary the React runtime does not understand that this throw is
a forward and it recreates the component stack from the Boundary
position. This loses fidelity and is especially bad if the boundary is
limited it what it handles and high up in the component tree.

This implementation uses a WeakMap to store component stacks for values
that are objects. If an error is rethrown from an ErrorBoundary the
stack will be pulled from the map if it exists. This doesn't work for
thrown primitives but this is uncommon and stashing the stack on the
primitive also wouldn't work

Author: gnoff

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,

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';
@@ -2805,7 +2805,7 @@ function updateDehydratedSuspenseComponent(
           );
         }
         (error: any).digest = digest;
-        capturedValue = createCapturedValue<mixed>(error, digest, stack);
+        capturedValue = createCapturedValueFromError(error, digest, stack);
       }
       return retrySuspenseComponentWithoutHydrating(
         current,
@@ -2942,7 +2942,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.',

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',
+      ]),
+    ]);
+  });
+});