Implement frame.yield to yield to split up synchronous work

The problem with the current implementation is that like with
renderToString we never yield to the event loop. When a large
app is rendered a lot of the work happens synchronously.

This reimplements the visitor to not be recursive and yield
when it's taking longer than 10ms. It does this by putting
a YieldFrame onto the suspense queue that preserves the
visitor's state and returns to it like with other suspense
frames.

Author: kitten

src/index.js

@@ -2,7 +2,7 @@
 
 import React, { type Node, type Element } from 'react'
 import type { Visitor, Frame, AbstractElement } from './types'
-import { visitChildren } from './visitor'
+import { visitChildren, resumeVisitChildren } from './visitor'
 import { getChildrenArray } from './element'
 
 import {
@@ -42,20 +42,31 @@ const flushFrames = (queue: Frame[], visitor: Visitor): Promise<void> => {
     let children = []
 
     // Update the component after we've suspended to rerender it,
-    // at which point we'll actually get its children
+    // at which point we'll actually get its children and continue
+    // walking the tree
     if (frame.kind === 'frame.class') {
-      children = getChildrenArray(updateClassComponent(queue, frame))
+      visitChildren(
+        getChildrenArray(updateClassComponent(queue, frame)),
+        queue,
+        visitor
+      )
     } else if (frame.kind === 'frame.hooks') {
-      children = getChildrenArray(updateFunctionComponent(queue, frame))
+      visitChildren(
+        getChildrenArray(updateFunctionComponent(queue, frame)),
+        queue,
+        visitor
+      )
     } else if (frame.kind === 'frame.lazy') {
-      children = getChildrenArray(updateLazyComponent(queue, frame))
+      visitChildren(
+        getChildrenArray(updateLazyComponent(queue, frame)),
+        queue,
+        visitor
+      )
+    } else if (frame.kind === 'frame.yield') {
+      resumeVisitChildren(frame, queue, visitor)
     }
 
-    // Now continue walking the previously suspended component's
-    // children (which might also suspend)
-    visitChildren(children, queue, visitor)
     ReactCurrentDispatcher.current = prevDispatcher
-
     return flushFrames(queue, visitor)
   })
 }

src/internals/context.js

@@ -4,7 +4,8 @@ import type {
   AbstractContext,
   UserElement,
   ContextMap,
-  ContextStore
+  ContextStore,
+  ContextEntry
 } from '../types'
 
 /** The context is kept as a Map from a Context value to the current
@@ -17,8 +18,6 @@ import type {
    After walking the children they can be restored.
    This way the context recursively restores itself on the way up. */
 
-type ContextEntry = [AbstractContext, mixed]
-
 let currentContextStore: ContextStore = new Map()
 let currentContextMap: ContextMap = {}
 
@@ -42,12 +41,16 @@ export const flushPrevContextStore = (): void | ContextEntry => {
   return prev
 }
 
-export const restoreContextMap = (prev: ContextMap) => {
-  Object.assign(currentContextMap, prev)
+export const restoreContextMap = (prev: void | ContextMap) => {
+  if (prev !== undefined) {
+    Object.assign(currentContextMap, prev)
+  }
 }
 
-export const restoreContextStore = (prev: ContextEntry) => {
-  currentContextStore.set(prev[0], prev[1])
+export const restoreContextStore = (prev: void | ContextEntry) => {
+  if (prev !== undefined) {
+    currentContextStore.set(prev[0], prev[1])
+  }
 }
 
 export const setCurrentContextMap = (map: ContextMap) => {

src/types/frames.js

@@ -3,8 +3,8 @@
 import type { ComponentType } from 'react'
 import type { Identity } from '../internals'
 import type { LazyComponent } from '../types'
-import type { DefaultProps, ComponentStatics } from './element'
-import type { ContextMap, ContextStore, Hook } from './state'
+import type { ContextMap, ContextStore, ContextEntry, Hook } from './state'
+import type { AbstractElement, DefaultProps, ComponentStatics } from './element'
 
 export type BaseFrame = {
   contextMap: ContextMap,
@@ -35,4 +35,13 @@ export type HooksFrame = BaseFrame & {
   hook: Hook | null
 }
 
-export type Frame = ClassFrame | HooksFrame | LazyFrame
+/** Description of a pause to yield to the event loop */
+export type YieldFrame = BaseFrame & {
+  kind: 'frame.yield',
+  children: AbstractElement[][],
+  index: number[],
+  map: Array<void | ContextMap>,
+  store: Array<void | ContextEntry>
+}
+
+export type Frame = ClassFrame | HooksFrame | LazyFrame | YieldFrame

src/types/state.js

@@ -4,6 +4,7 @@ import type { AbstractContext } from './element'
 
 export type ContextStore = Map<AbstractContext, mixed>
 export type ContextMap = { [name: string]: mixed }
+export type ContextEntry = [AbstractContext, mixed]
 
 export type Dispatch<A> = A => void
 

src/visitor.js

@@ -11,8 +11,10 @@ import {
 
 import type {
   Visitor,
+  YieldFrame,
   Frame,
   ContextMap,
+  ContextEntry,
   DefaultProps,
   ComponentStatics,
   LazyElement,
@@ -28,12 +30,17 @@ import type {
 } from './types'
 
 import {
+  getCurrentContextMap,
+  getCurrentContextStore,
+  setCurrentContextMap,
+  setCurrentContextStore,
   flushPrevContextMap,
   flushPrevContextStore,
   restoreContextMap,
   restoreContextStore,
   readContextValue,
-  setContextValue
+  setContextValue,
+  setCurrentIdentity
 } from './internals'
 
 import {
@@ -51,6 +58,12 @@ import {
   REACT_LAZY_TYPE
 } from './symbols'
 
+const makeImmediatePromise = () => {
+  return new Promise(resolve => {
+    setImmediate(resolve)
+  })
+}
+
 const render = (
   type: ComponentType<DefaultProps> & ComponentStatics,
   props: DefaultProps,
@@ -148,39 +161,78 @@ export const visitElement = (
   }
 }
 
-/** The context (legacy and createContext are separate) are kept
-    as global state. As we're walking the tree depth-first, we
-    simply overwrite it with new values. This is recursive and
-    as we walk back upwards (after visitChildren) we restore
-    the value that we have previously overwritten. */
-const visitChild = (
-  child: AbstractElement,
+const visitLoop = (
+  traversalChildren: AbstractElement[][],
+  traversalIndex: number[],
+  traversalMap: Array<void | ContextMap>,
+  traversalStore: Array<void | ContextEntry>,
   queue: Frame[],
   visitor: Visitor
 ) => {
-  const children = visitElement(child, queue, visitor)
-  // Flush changes (if any) that have been made to the context
-  const prevMap = flushPrevContextMap()
-  const prevStore = flushPrevContextStore()
-
-  visitChildren(children, queue, visitor)
-
-  // Restore context changes after children have been walked
-  if (prevMap !== undefined) {
-    restoreContextMap(prevMap)
-  }
-
-  if (prevStore !== undefined) {
-    restoreContextStore(prevStore)
+  const start = Date.now()
+
+  while (traversalChildren.length > 0 && Date.now() - start <= 10) {
+    const currChildren = traversalChildren[traversalChildren.length - 1]
+    const currIndex = traversalIndex[traversalIndex.length - 1]++
+
+    if (currIndex < currChildren.length) {
+      const element = currChildren[currIndex]
+      const children = visitElement(element, queue, visitor)
+
+      traversalChildren.push(children)
+      traversalIndex.push(0)
+      traversalMap.push(flushPrevContextMap())
+      traversalStore.push(flushPrevContextStore())
+    } else {
+      traversalChildren.pop()
+      traversalIndex.pop()
+      restoreContextMap(traversalMap.pop())
+      restoreContextStore(traversalStore.pop())
+    }
   }
 }
 
 export const visitChildren = (
-  children: AbstractElement[],
+  init: AbstractElement[],
   queue: Frame[],
   visitor: Visitor
 ) => {
-  for (let i = 0, l = children.length; i < l; i++) {
-    visitChild(children[i], queue, visitor)
+  const traversalChildren: AbstractElement[][] = [init]
+  const traversalIndex: number[] = [0]
+  const traversalMap: Array<void | ContextMap> = [undefined]
+  const traversalStore: Array<void | ContextEntry> = [undefined]
+
+  visitLoop(
+    traversalChildren,
+    traversalIndex,
+    traversalMap,
+    traversalStore,
+    queue,
+    visitor
+  )
+
+  if (traversalChildren.length > 0) {
+    queue.push({
+      contextMap: getCurrentContextMap(),
+      contextStore: getCurrentContextStore(),
+      kind: 'frame.yield',
+      thenable: makeImmediatePromise(),
+      children: traversalChildren,
+      index: traversalIndex,
+      map: traversalMap,
+      store: traversalStore
+    })
   }
 }
+
+export const resumeVisitChildren = (
+  frame: YieldFrame,
+  queue: Frame[],
+  visitor: Visitor
+) => {
+  setCurrentIdentity(null)
+  setCurrentContextMap(frame.contextMap)
+  setCurrentContextMap(frame.contextStore)
+
+  visitLoop(frame.children, frame.index, frame.map, frame.store, queue, visitor)
+}