[Flight] Support more element types and Hooks for Server and Hybrid Components (#19711)

* Shim support for more element types

* Shim commonly used Hooks that are safe

* Flow

* Oopsie

packages/react-server/src/ReactFlightServer.js

@@ -7,6 +7,7 @@
  * @flow
  */
 
+import type {Dispatcher as DispatcherType} from 'react-reconciler/src/ReactInternalTypes';
 import type {
   Destination,
   Chunk,
@@ -29,12 +30,24 @@ import {
 
 import {
   REACT_BLOCK_TYPE,
-  REACT_SERVER_BLOCK_TYPE,
   REACT_ELEMENT_TYPE,
+  REACT_DEBUG_TRACING_MODE_TYPE,
+  REACT_FORWARD_REF_TYPE,
   REACT_FRAGMENT_TYPE,
   REACT_LAZY_TYPE,
+  REACT_LEGACY_HIDDEN_TYPE,
+  REACT_MEMO_TYPE,
+  REACT_OFFSCREEN_TYPE,
+  REACT_PROFILER_TYPE,
+  REACT_SCOPE_TYPE,
+  REACT_SERVER_BLOCK_TYPE,
+  REACT_STRICT_MODE_TYPE,
+  REACT_SUSPENSE_TYPE,
+  REACT_SUSPENSE_LIST_TYPE,
 } from 'shared/ReactSymbols';
 
+import * as React from 'react';
+import ReactSharedInternals from 'shared/ReactSharedInternals';
 import invariant from 'shared/invariant';
 
 type ReactJSONValue =
@@ -74,6 +87,8 @@ export type Request = {
   toJSON: (key: string, value: ReactModel) => ReactJSONValue,
 };
 
+const ReactCurrentDispatcher = ReactSharedInternals.ReactCurrentDispatcher;
+
 export function createRequest(
   model: ReactModel,
   destination: Destination,
@@ -110,11 +125,33 @@ function attemptResolveElement(element: React$Element<any>): ReactModel {
     return [REACT_ELEMENT_TYPE, type, element.key, element.props];
   } else if (type[0] === REACT_SERVER_BLOCK_TYPE) {
     return [REACT_ELEMENT_TYPE, type, element.key, element.props];
-  } else if (type === REACT_FRAGMENT_TYPE) {
+  } else if (
+    type === REACT_FRAGMENT_TYPE ||
+    type === REACT_STRICT_MODE_TYPE ||
+    type === REACT_PROFILER_TYPE ||
+    type === REACT_SCOPE_TYPE ||
+    type === REACT_DEBUG_TRACING_MODE_TYPE ||
+    type === REACT_LEGACY_HIDDEN_TYPE ||
+    type === REACT_OFFSCREEN_TYPE ||
+    // TODO: These are temporary shims
+    // and we'll want a different behavior.
+    type === REACT_SUSPENSE_TYPE ||
+    type === REACT_SUSPENSE_LIST_TYPE
+  ) {
     return element.props.children;
-  } else {
-    invariant(false, 'Unsupported type.');
+  } else if (type != null && typeof type === 'object') {
+    switch (type.$$typeof) {
+      case REACT_FORWARD_REF_TYPE: {
+        const render = type.render;
+        return render(props, undefined);
+      }
+      case REACT_MEMO_TYPE: {
+        const nextChildren = React.createElement(type.type, element.props);
+        return attemptResolveElement(nextChildren);
+      }
+    }
   }
+  invariant(false, 'Unsupported type.');
 }
 
 function pingSegment(request: Request, segment: Segment): void {
@@ -236,9 +273,11 @@ export function resolveModelToJSON(
     value !== null &&
     value.$$typeof === REACT_ELEMENT_TYPE
   ) {
+    const prevDispatcher = ReactCurrentDispatcher.current;
     // TODO: Concatenate keys of parents onto children.
     const element: React$Element<any> = (value: any);
     try {
+      ReactCurrentDispatcher.current = Dispatcher;
       // Attempt to render the server component.
       value = attemptResolveElement(element);
     } catch (x) {
@@ -253,6 +292,8 @@ export function resolveModelToJSON(
         // Something errored. Don't bother encoding anything up to here.
         throw x;
       }
+    } finally {
+      ReactCurrentDispatcher.current = prevDispatcher;
     }
   }
 
@@ -378,3 +419,33 @@ export function startFlowing(request: Request): void {
   request.flowing = true;
   flushCompletedChunks(request);
 }
+
+function unsupportedHook(): void {
+  invariant(false, 'This Hook is not supported in Server Components.');
+}
+
+const Dispatcher: DispatcherType = {
+  useMemo<T>(nextCreate: () => T): T {
+    return nextCreate();
+  },
+  useCallback<T>(callback: T): T {
+    return callback;
+  },
+  useDebugValue(): void {},
+  useDeferredValue<T>(value: T): T {
+    return value;
+  },
+  useTransition(): [(callback: () => void) => void, boolean] {
+    return [() => {}, false];
+  },
+  readContext: (unsupportedHook: any),
+  useContext: (unsupportedHook: any),
+  useReducer: (unsupportedHook: any),
+  useRef: (unsupportedHook: any),
+  useState: (unsupportedHook: any),
+  useLayoutEffect: (unsupportedHook: any),
+  useImperativeHandle: (unsupportedHook: any),
+  useEffect: (unsupportedHook: any),
+  useOpaqueIdentifier: (unsupportedHook: any),
+  useMutableSource: (unsupportedHook: any),
+};

packages/react-transport-dom-relay/src/__tests__/ReactFlightDOMRelay-test.internal.js

@@ -124,4 +124,93 @@ describe('ReactFlightDOMRelay', () => {
 
     expect(container.innerHTML).toEqual('<span>Hello, Seb Smith</span>');
   });
+
+  // @gate experimental
+  it('can reasonably handle different element types', () => {
+    const {
+      forwardRef,
+      memo,
+      Fragment,
+      StrictMode,
+      Profiler,
+      Suspense,
+      SuspenseList,
+    } = React;
+
+    const Inner = memo(
+      forwardRef((props, ref) => {
+        return <div ref={ref}>{'Hello ' + props.name}</div>;
+      }),
+    );
+
+    function Foo() {
+      return {
+        bar: (
+          <div>
+            <Fragment>Fragment child</Fragment>
+            <Profiler>Profiler child</Profiler>
+            <StrictMode>StrictMode child</StrictMode>
+            <Suspense fallback="Loading...">Suspense child</Suspense>
+            <SuspenseList fallback="Loading...">
+              {'SuspenseList row 1'}
+              {'SuspenseList row 2'}
+            </SuspenseList>
+            <Inner name="world" />
+          </div>
+        ),
+      };
+    }
+    const transport = [];
+    ReactDOMFlightRelayServer.render(
+      {
+        foo: <Foo />,
+      },
+      transport,
+    );
+
+    const model = readThrough(transport);
+    expect(model).toEqual({
+      foo: {
+        bar: (
+          <div>
+            {'Fragment child'}
+            {'Profiler child'}
+            {'StrictMode child'}
+            {'Suspense child'}
+            {['SuspenseList row 1', 'SuspenseList row 2']}
+            <div>Hello world</div>
+          </div>
+        ),
+      },
+    });
+  });
+
+  it('can handle a subset of Hooks', () => {
+    const {useMemo, useCallback} = React;
+    function Inner({x}) {
+      const foo = useMemo(() => x + x, [x]);
+      const bar = useCallback(() => 10 + foo, [foo]);
+      return bar();
+    }
+
+    function Foo() {
+      return {
+        bar: <Inner x={2} />,
+      };
+    }
+    const transport = [];
+    ReactDOMFlightRelayServer.render(
+      {
+        foo: <Foo />,
+      },
+      transport,
+    );
+
+    const model = readThrough(transport);
+    expect(model).toEqual({
+      foo: {
+        bar: 14,
+      },
+    });
+  });
 });

scripts/error-codes/codes.json

@@ -360,5 +360,6 @@
   "369": "ReactDOM.createEventHandle: setter called on an invalid target. Provide a valid EventTarget or an element managed by React.",
   "370": "ReactDOM.createEventHandle: setter called with an invalid callback. The callback must be a function.",
   "371": "Text string must be rendered within a <Text> component.\n\nText: %s",
-  "372": "Cannot call unstable_createEventHandle with \"%s\", as it is not an event known to React."
+  "372": "Cannot call unstable_createEventHandle with \"%s\", as it is not an event known to React.",
+  "373": "This Hook is not supported in Server Components."
 }