Switch Blocks to using a Lazy component wrapper

sebmarkbage · sebmarkbage · commit d26721f21ae6 · 2020-03-21T11:46:40.000-07:00
Then resolve to a true Block inside.
diff --git a/packages/react-reconciler/src/ReactChildFiber.js b/packages/react-reconciler/src/ReactChildFiber.js
@@ -10,6 +10,7 @@
 import type {ReactElement} from 'shared/ReactElementType';
 import type {ReactPortal} from 'shared/ReactTypes';
 import type {BlockComponent} from 'react/src/block';
+import type {LazyComponent} from 'react/src/ReactLazy';
 import type {Fiber} from './ReactFiber';
 import type {ExpirationTime} from './ReactFiberExpirationTime';
 
@@ -20,6 +21,7 @@ import {
   REACT_ELEMENT_TYPE,
   REACT_FRAGMENT_TYPE,
   REACT_PORTAL_TYPE,
+  REACT_LAZY_TYPE,
   REACT_BLOCK_TYPE,
 } from 'shared/ReactSymbols';
 import {
@@ -48,7 +50,6 @@ import {
 } from './ReactCurrentFiber';
 import {isCompatibleFamilyForHotReloading} from './ReactFiberHotReloading';
 import {StrictMode} from './ReactTypeOfMode';
-import {initializeBlockComponentType} from 'shared/ReactLazyComponent';
 
 let didWarnAboutMaps;
 let didWarnAboutGenerators;
@@ -263,6 +264,22 @@ function warnOnFunctionType() {
   }
 }
 
+// We avoid inlining this to avoid potential deopts from using try/catch.
+/** @noinline */
+function resolveLazyType<T, P>(
+  lazyComponent: LazyComponent<T, P>,
+): LazyComponent<T, P> | T {
+  try {
+    // If we can, let's peak at the resulting type.
+    let payload = lazyComponent._payload;
+    let init = lazyComponent._init;
+    return init(payload);
+  } catch (x) {
+    // Leave it in place and let it throw again in the begin phase.
+    return lazyComponent;
+  }
+}
+
 // This wrapper function exists because I expect to clone the code in each path
 // to be able to optimize each path individually by branching early. This needs
 // a compiler or we can do it manually. Helpers that don't need this branching
@@ -419,22 +436,22 @@ function ChildReconciler(shouldTrackSideEffects) {
           existing._debugOwner = element._owner;
         }
         return existing;
-      } else if (
-        enableBlocksAPI &&
-        current.tag === Block &&
-        element.type.$$typeof === REACT_BLOCK_TYPE
-      ) {
+      } else if (enableBlocksAPI && current.tag === Block) {
         // The new Block might not be initialized yet. We need to initialize
         // it in case initializing it turns out it would match.
-        initializeBlockComponentType(element.type);
+        let type = element.type;
+        if (type.$$typeof === REACT_LAZY_TYPE) {
+          type = resolveLazyType(type);
+        }
         if (
-          (element.type: BlockComponent<any, any, any>)._fn ===
-          (current.type: BlockComponent<any, any, any>)._fn
+          type.$$typeof === REACT_BLOCK_TYPE &&
+          ((type: any): BlockComponent<any, any>)._render ===
+            (current.type: BlockComponent<any, any>)._render
         ) {
           // Same as above but also update the .type field.
           const existing = useFiber(current, element.props);
           existing.return = returnFiber;
-          existing.type = element.type;
+          existing.type = type;
           if (__DEV__) {
             existing._debugSource = element._source;
             existing._debugOwner = element._owner;
@@ -1188,17 +1205,20 @@ function ChildReconciler(shouldTrackSideEffects) {
           }
           case Block:
             if (enableBlocksAPI) {
-              if (element.type.$$typeof === REACT_BLOCK_TYPE) {
+              let type = element.type;
+              if (type.$$typeof === REACT_LAZY_TYPE) {
+                type = resolveLazyType(type);
+              }
+              if (type.$$typeof === REACT_BLOCK_TYPE) {
                 // The new Block might not be initialized yet. We need to initialize
                 // it in case initializing it turns out it would match.
-                initializeBlockComponentType(element.type);
                 if (
-                  (element.type: BlockComponent<any, any, any>)._fn ===
-                  (child.type: BlockComponent<any, any, any>)._fn
+                  ((type: any): BlockComponent<any, any>)._render ===
+                  (child.type: BlockComponent<any, any>)._render
                 ) {
                   deleteRemainingChildren(returnFiber, child.sibling);
                   const existing = useFiber(child, element.props);
-                  existing.type = element.type;
+                  existing.type = type;
                   existing.return = returnFiber;
                   if (__DEV__) {
                     existing._debugSource = element._source;
diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.js b/packages/react-reconciler/src/ReactFiberBeginWork.js
@@ -165,7 +165,6 @@ import {
   updateClassInstance,
 } from './ReactFiberClassComponent';
 import {resolveDefaultProps} from './ReactFiberLazyComponent';
-import {initializeBlockComponentType} from 'shared/ReactLazyComponent';
 import {
   resolveLazyComponentTag,
   createFiberFromTypeAndProps,
@@ -181,7 +180,6 @@ import {
   renderDidSuspendDelayIfPossible,
   markUnprocessedUpdateTime,
 } from './ReactFiberWorkLoop';
-import {Resolved} from 'shared/ReactLazyStatusTags';
 
 const ReactCurrentOwner = ReactSharedInternals.ReactCurrentOwner;
 
@@ -707,23 +705,18 @@ function updateFunctionComponent(
   return workInProgress.child;
 }
 
-function updateBlock<Props, Payload, Data>(
+function updateBlock<Props, Data>(
   current: Fiber | null,
   workInProgress: Fiber,
-  block: BlockComponent<Props, Payload, Data>,
+  block: BlockComponent<Props, Data>,
   nextProps: any,
   renderExpirationTime: ExpirationTime,
 ) {
   // TODO: current can be non-null here even if the component
   // hasn't yet mounted. This happens after the first render suspends.
   // We'll need to figure out if this is fine or can cause issues.
 
-  initializeBlockComponentType(block);
-  if (block._status !== Resolved) {
-    throw block._data;
-  }
-
-  const render = block._fn;
+  const render = block._render;
   const data = block._data;
 
   // The rest is a fork of updateFunctionComponent
diff --git a/packages/react-reconciler/src/__tests__/ReactBlocks-test.js b/packages/react-reconciler/src/__tests__/ReactBlocks-test.js
@@ -14,11 +14,13 @@ let useState;
 let Suspense;
 let block;
 let readString;
+let Scheduler;
 
 describe('ReactBlocks', () => {
   beforeEach(() => {
     jest.resetModules();
 
+    Scheduler = require('scheduler');
     React = require('react');
     ReactNoop = require('react-noop-renderer');
 
@@ -86,64 +88,64 @@ describe('ReactBlocks', () => {
     expect(ReactNoop).toMatchRenderedOutput(<span>Name: Sebastian</span>);
   });
 
-  it.experimental('supports a lazy wrapper around a chunk', async () => {
-    function Query(id) {
-      return {
-        id: id,
-        name: readString('Sebastian'),
-      };
-    }
+  it.experimental(
+    'does not support a lazy wrapper around a chunk',
+    async () => {
+      function Query(id) {
+        return {
+          id: id,
+          name: readString('Sebastian'),
+        };
+      }
 
-    function Render(props, data) {
-      return (
-        <span>
-          {props.title}: {data.name}
-        </span>
-      );
-    }
+      function Render(props, data) {
+        return (
+          <span>
+            {props.title}: {data.name}
+          </span>
+        );
+      }
 
-    let loadUser = block(Query, Render);
+      let loadUser = block(Query, Render);
 
-    function App({User}) {
-      return (
-        <Suspense fallback={'Loading...'}>
-          <User title="Name" />
-        </Suspense>
+      function App({User}) {
+        return (
+          <Suspense fallback={'Loading...'}>
+            <User title="Name" />
+          </Suspense>
+        );
+      }
+
+      let resolveLazy;
+      let LazyUser = React.lazy(
+        () =>
+          new Promise(resolve => {
+            resolveLazy = function() {
+              resolve({
+                default: loadUser(123),
+              });
+            };
+          }),
       );
-    }
 
-    let resolveLazy;
-    let LazyUser = React.lazy(
-      () =>
-        new Promise(resolve => {
-          resolveLazy = function() {
-            resolve({
-              default: loadUser(123),
-            });
-          };
-        }),
-    );
+      await ReactNoop.act(async () => {
+        ReactNoop.render(<App User={LazyUser} />);
+      });
 
-    await ReactNoop.act(async () => {
-      ReactNoop.render(<App User={LazyUser} />);
-    });
+      expect(ReactNoop).toMatchRenderedOutput('Loading...');
 
-    expect(ReactNoop).toMatchRenderedOutput('Loading...');
-
-    // Resolve the component.
-    await ReactNoop.act(async () => {
+      // Resolve the component.
       await resolveLazy();
-    });
-
-    // We're still waiting on the data.
-    expect(ReactNoop).toMatchRenderedOutput('Loading...');
-
-    await ReactNoop.act(async () => {
-      jest.advanceTimersByTime(1000);
-    });
 
-    expect(ReactNoop).toMatchRenderedOutput(<span>Name: Sebastian</span>);
-  });
+      expect(Scheduler).toFlushAndThrow(
+        'Element type is invalid. Received a promise that resolves to: [object Object]. ' +
+          'Lazy element type must resolve to a class or function.' +
+          (__DEV__
+            ? ' Did you wrap a component in React.lazy() more than once?'
+            : ''),
+      );
+    },
+  );
 
   it.experimental(
     'can receive updated data for the same component',
diff --git a/packages/react/src/block.js b/packages/react/src/block.js
@@ -7,7 +7,10 @@
  * @flow
  */
 
+import type {LazyComponent} from './ReactLazy';
+
 import {
+  REACT_LAZY_TYPE,
   REACT_BLOCK_TYPE,
   REACT_MEMO_TYPE,
   REACT_FORWARD_REF_TYPE,
@@ -19,55 +22,33 @@ type BlockRenderFunction<Props, Data> = (
   data: Data,
 ) => React$Node;
 
-type Thenable<T, R> = {
-  then(resolve: (T) => mixed, reject: (mixed) => mixed): R,
-};
-
-type Initializer<Props, Payload, Data> = (
-  payload: Payload,
-) =>
-  | [Data, BlockRenderFunction<Props, Data>]
-  | Thenable<[Data, BlockRenderFunction<Props, Data>], mixed>;
-
-export type UninitializedBlockComponent<Props, Payload, Data> = {
-  $$typeof: Symbol | number,
-  _status: -1,
-  _data: Payload,
-  _fn: Initializer<Props, Payload, Data>,
-};
-
-export type PendingBlockComponent<Props, Data> = {
-  $$typeof: Symbol | number,
-  _status: 0,
-  _data: Thenable<[Data, BlockRenderFunction<Props, Data>], mixed>,
-  _fn: null,
+type Payload<Props, Args: Iterable<any>, Data> = {
+  query: BlockQueryFunction<Args, Data>,
+  args: Args,
+  render: BlockRenderFunction<Props, Data>,
 };
 
-export type ResolvedBlockComponent<Props, Data> = {
+export type BlockComponent<Props, Data> = {
   $$typeof: Symbol | number,
-  _status: 1,
   _data: Data,
-  _fn: BlockRenderFunction<Props, Data>,
+  _render: BlockRenderFunction<Props, Data>,
 };
 
-export type RejectedBlockComponent = {
-  $$typeof: Symbol | number,
-  _status: 2,
-  _data: mixed,
-  _fn: null,
-};
-
-export type BlockComponent<Props, Payload, Data> =
-  | UninitializedBlockComponent<Props, Payload, Data>
-  | PendingBlockComponent<Props, Data>
-  | ResolvedBlockComponent<Props, Data>
-  | RejectedBlockComponent;
-
 opaque type Block<Props>: React$AbstractComponent<
   Props,
   null,
 > = React$AbstractComponent<Props, null>;
 
+function lazyInitializer<Props, Args: Iterable<any>, Data>(
+  payload: Payload<Props, Args, Data>,
+): BlockComponent<Props, Data> {
+  return {
+    $$typeof: REACT_BLOCK_TYPE,
+    _data: payload.query.apply(null, payload.args),
+    _render: payload.render,
+  };
+}
+
 export default function block<Args: Iterable<any>, Props, Data>(
   query: BlockQueryFunction<Args, Data>,
   render: BlockRenderFunction<Props, Data>,
@@ -115,19 +96,26 @@ export default function block<Args: Iterable<any>, Props, Data>(
       );
     }
   }
-  function initializer(args) {
-    let data = query.apply(null, args);
-    return [data, render];
-  }
+
   return function(): Block<Props> {
     let args: Args = arguments;
-    let blockComponent: UninitializedBlockComponent<Props, Args, Data> = {
-      $$typeof: REACT_BLOCK_TYPE,
-      _status: -1,
-      _data: args,
-      _fn: initializer,
+
+    let payload: Payload<Props, Args, Data> = {
+      query: query,
+      args: args,
+      render: render,
     };
+
+    let lazyType: LazyComponent<
+      BlockComponent<Props, Data>,
+      Payload<Props, Args, Data>,
+    > = {
+      $$typeof: REACT_LAZY_TYPE,
+      _payload: payload,
+      _init: lazyInitializer,
+    };
+
     // $FlowFixMe
-    return blockComponent;
+    return lazyType;
   };
 }
diff --git a/packages/shared/ReactLazyComponent.js b/packages/shared/ReactLazyComponent.js
diff --git a/packages/shared/ReactLazyStatusTags.js b/packages/shared/ReactLazyStatusTags.js
