[Blocks] Make it possible to have lazy initialized and lazy loaded Blocks

packages/react-reconciler/src/ReactChildFiber.js

@@ -9,6 +9,7 @@
 
 import type {ReactElement} from 'shared/ReactElementType';
 import type {ReactPortal} from 'shared/ReactTypes';
+import type {BlockComponent} from 'react/src/block';
 import type {Fiber} from './ReactFiber';
 import type {ExpirationTime} from './ReactFiberExpirationTime';
 
@@ -47,6 +48,7 @@ import {
 } from './ReactCurrentFiber';
 import {isCompatibleFamilyForHotReloading} from './ReactFiberHotReloading';
 import {StrictMode} from './ReactTypeOfMode';
+import {initializeBlockComponentType} from 'shared/ReactLazyComponent';
 
 let didWarnAboutMaps;
 let didWarnAboutGenerators;
@@ -420,18 +422,25 @@ function ChildReconciler(shouldTrackSideEffects) {
       } else if (
         enableBlocksAPI &&
         current.tag === Block &&
-        element.type.$$typeof === REACT_BLOCK_TYPE &&
-        element.type.render === current.type.render
+        element.type.$$typeof === REACT_BLOCK_TYPE
       ) {
-        // Same as above but also update the .type field.
-        const existing = useFiber(current, element.props);
-        existing.return = returnFiber;
-        existing.type = element.type;
-        if (__DEV__) {
-          existing._debugSource = element._source;
-          existing._debugOwner = element._owner;
+        // 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 ===
+          (current.type: BlockComponent<any, any, any>)._fn
+        ) {
+          // Same as above but also update the .type field.
+          const existing = useFiber(current, element.props);
+          existing.return = returnFiber;
+          existing.type = element.type;
+          if (__DEV__) {
+            existing._debugSource = element._source;
+            existing._debugOwner = element._owner;
+          }
+          return existing;
         }
-        return existing;
       }
     }
     // Insert
@@ -1179,19 +1188,24 @@ function ChildReconciler(shouldTrackSideEffects) {
           }
           case Block:
             if (enableBlocksAPI) {
-              if (
-                element.type.$$typeof === REACT_BLOCK_TYPE &&
-                element.type.render === child.type.render
-              ) {
-                deleteRemainingChildren(returnFiber, child.sibling);
-                const existing = useFiber(child, element.props);
-                existing.type = element.type;
-                existing.return = returnFiber;
-                if (__DEV__) {
-                  existing._debugSource = element._source;
-                  existing._debugOwner = element._owner;
+              if (element.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
+                ) {
+                  deleteRemainingChildren(returnFiber, child.sibling);
+                  const existing = useFiber(child, element.props);
+                  existing.type = element.type;
+                  existing.return = returnFiber;
+                  if (__DEV__) {
+                    existing._debugSource = element._source;
+                    existing._debugOwner = element._owner;
+                  }
+                  return existing;
                 }
-                return existing;
               }
             }
           // We intentionally fallthrough here if enableBlocksAPI is not on.

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -8,6 +8,7 @@
  */
 
 import type {ReactProviderType, ReactContext} from 'shared/ReactTypes';
+import type {BlockComponent} from 'react/src/block';
 import type {Fiber} from './ReactFiber';
 import type {FiberRoot} from './ReactFiberRoot';
 import type {ExpirationTime} from './ReactFiberExpirationTime';
@@ -17,6 +18,7 @@ import type {
   SuspenseListTailMode,
 } from './ReactFiberSuspenseComponent';
 import type {SuspenseContext} from './ReactFiberSuspenseContext';
+import type {} from 'react/src/block';
 
 import checkPropTypes from 'shared/checkPropTypes';
 
@@ -166,6 +168,7 @@ import {
   readLazyComponentType,
   resolveDefaultProps,
 } from './ReactFiberLazyComponent';
+import {initializeBlockComponentType} from 'shared/ReactLazyComponent';
 import {
   resolveLazyComponentTag,
   createFiberFromTypeAndProps,
@@ -181,6 +184,7 @@ import {
   renderDidSuspendDelayIfPossible,
   markUnprocessedUpdateTime,
 } from './ReactFiberWorkLoop';
+import {Resolved} from 'shared/ReactLazyStatusTags';
 
 const ReactCurrentOwner = ReactSharedInternals.ReactCurrentOwner;
 
@@ -693,19 +697,24 @@ function updateFunctionComponent(
   return workInProgress.child;
 }
 
-function updateBlock(
+function updateBlock<Props, Payload, Data>(
   current: Fiber | null,
   workInProgress: Fiber,
-  block: any,
+  block: BlockComponent<Props, Payload, 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.
 
-  const render = block.render;
-  const data = block.query();
+  initializeBlockComponentType(block);
+  if (block._status !== Resolved) {
+    throw block._data;
+  }
+
+  const render = block._fn;
+  const data = block._data;
 
   // The rest is a fork of updateFunctionComponent
   let nextChildren;

packages/react-reconciler/src/ReactFiberHooks.js

@@ -343,12 +343,12 @@ function areHookInputsEqual(
   return true;
 }
 
-export function renderWithHooks(
+export function renderWithHooks<Props, SecondArg>(
   current: Fiber | null,
   workInProgress: Fiber,
-  Component: any,
-  props: any,
-  secondArg: any,
+  Component: (p: Props, arg: SecondArg) => any,
+  props: Props,
+  secondArg: SecondArg,
   nextRenderExpirationTime: ExpirationTime,
 ): any {
   renderExpirationTime = nextRenderExpirationTime;

packages/react/src/block.js

@@ -3,6 +3,8 @@
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
  */
 
 import {
@@ -11,14 +13,64 @@ import {
   REACT_FORWARD_REF_TYPE,
 } from 'shared/ReactSymbols';
 
+type BlockQueryFunction<Args: Iterable<any>, Data> = (...args: Args) => Data;
+type BlockRenderFunction<Props, Data> = (
+  props: Props,
+  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,
+};
+
+export type ResolvedBlockComponent<Props, Data> = {
+  $$typeof: Symbol | number,
+  _status: 1,
+  _data: Data,
+  _fn: 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>;
 
-export default function block<Args, Props, Data>(
-  query: (...args: Args) => Data,
-  render: (props: Props, data: Data) => React$Node,
+export default function block<Args: Iterable<any>, Props, Data>(
+  query: BlockQueryFunction<Args, Data>,
+  render: BlockRenderFunction<Props, Data>,
 ): (...args: Args) => Block<Props> {
   if (__DEV__) {
     if (typeof query !== 'function') {
@@ -63,14 +115,19 @@ export default function block<Args, Props, Data>(
       );
     }
   }
+  function initializer(args) {
+    let data = query.apply(null, args);
+    return [data, render];
+  }
   return function(): Block<Props> {
-    let args = arguments;
-    return {
+    let args: Args = arguments;
+    let blockComponent: UninitializedBlockComponent<Props, Args, Data> = {
       $$typeof: REACT_BLOCK_TYPE,
-      query: function() {
-        return query.apply(null, args);
-      },
-      render: render,
+      _status: -1,
+      _data: args,
+      _fn: initializer,
     };
+    // $FlowFixMe
+    return blockComponent;
   };
 }

packages/shared/ReactLazyComponent.js

@@ -14,6 +14,13 @@ import type {
   LazyComponent,
 } from 'react/src/ReactLazy';
 
+import type {
+  PendingBlockComponent,
+  ResolvedBlockComponent,
+  RejectedBlockComponent,
+  BlockComponent,
+} from 'react/src/block';
+
 import {
   Uninitialized,
   Pending,
@@ -72,3 +79,50 @@ export function initializeLazyComponentType(
     );
   }
 }
+
+export function initializeBlockComponentType<Props, Payload, Data>(
+  blockComponent: BlockComponent<Props, Payload, Data>,
+): void {
+  if (blockComponent._status === Uninitialized) {
+    const thenableOrTuple = blockComponent._fn(blockComponent._data);
+    if (typeof thenableOrTuple.then !== 'function') {
+      let tuple: [any, any] = (thenableOrTuple: any);
+      const resolved: ResolvedBlockComponent<
+        Props,
+        Data,
+      > = (blockComponent: any);
+      resolved._status = Resolved;
+      resolved._data = tuple[0];
+      resolved._fn = tuple[1];
+      return;
+    }
+    const thenable = (thenableOrTuple: any);
+    // Transition to the next state.
+    const pending: PendingBlockComponent<Props, Data> = (blockComponent: any);
+    pending._status = Pending;
+    pending._data = thenable;
+    pending._fn = null;
+    thenable.then(
+      (tuple: [any, any]) => {
+        if (blockComponent._status === Pending) {
+          // Transition to the next state.
+          const resolved: ResolvedBlockComponent<
+            Props,
+            Data,
+          > = (blockComponent: any);
+          resolved._status = Resolved;
+          resolved._data = tuple[0];
+          resolved._fn = tuple[1];
+        }
+      },
+      error => {
+        if (blockComponent._status === Pending) {
+          // Transition to the next state.
+          const rejected: RejectedBlockComponent = (blockComponent: any);
+          rejected._status = Rejected;
+          rejected._data = error;
+        }
+      },
+    );
+  }
+}