Clean up and fix incorrect context setter

Author: kitten

src/__tests__/suspense.test.js

@@ -18,7 +18,7 @@ describe('renderPrepass', () => {
 
       const Outer = () => {
         const start = Date.now()
-        while (Date.now() - start < 11) {}
+        while (Date.now() - start < 6) {}
         return <Inner />
       }
 

src/index.js

@@ -1,7 +1,7 @@
 // @flow
 
 import React, { type Node, type Element } from 'react'
-import type { Visitor, Frame, AbstractElement } from './types'
+import type { Visitor, YieldFrame, Frame, AbstractElement } from './types'
 import { visitChildren, resumeVisitChildren } from './visitor'
 import { getChildrenArray } from './element'
 
@@ -28,49 +28,58 @@ let prevDispatcher = ReactCurrentDispatcher.current
     the queue. Hence we recursively look at suspended components in
     this queue, wait for their promises to resolve, and continue
     calling visitChildren on their children. */
-const flushFrames = (queue: Frame[], visitor: Visitor): Promise<void> => {
-  if (queue.length === 0) {
-    return Promise.resolve()
+const updateWithFrame = (
+  frame: Frame,
+  queue: Frame[],
+  visitor: Visitor
+): Promise<void> => {
+  if (frame.kind === 'frame.yield') {
+    const yieldFrame: YieldFrame = frame
+
+    return new Promise(resolve => {
+      setImmediate(() => {
+        prevDispatcher = ReactCurrentDispatcher.current
+        ReactCurrentDispatcher.current = Dispatcher
+        resumeVisitChildren(yieldFrame, queue, visitor)
+        ReactCurrentDispatcher.current = prevDispatcher
+        resolve()
+      })
+    })
   }
 
-  const frame = queue.shift()
-
   return frame.thenable.then(() => {
     prevDispatcher = ReactCurrentDispatcher.current
     ReactCurrentDispatcher.current = Dispatcher
 
     let children = []
 
     // Update the component after we've suspended to rerender it,
-    // at which point we'll actually get its children and continue
-    // walking the tree
+    // at which point we'll actually get its children
     if (frame.kind === 'frame.class') {
-      visitChildren(
-        getChildrenArray(updateClassComponent(queue, frame)),
-        queue,
-        visitor
-      )
+      children = getChildrenArray(updateClassComponent(queue, frame))
     } else if (frame.kind === 'frame.hooks') {
-      visitChildren(
-        getChildrenArray(updateFunctionComponent(queue, frame)),
-        queue,
-        visitor
-      )
+      children = getChildrenArray(updateFunctionComponent(queue, frame))
     } else if (frame.kind === 'frame.lazy') {
-      visitChildren(
-        getChildrenArray(updateLazyComponent(queue, frame)),
-        queue,
-        visitor
-      )
-    } else if (frame.kind === 'frame.yield') {
-      resumeVisitChildren(frame, queue, visitor)
+      children = getChildrenArray(updateLazyComponent(queue, frame))
     }
 
+    // Now continue walking the previously suspended component's
+    // children (which might also suspend)
+    visitChildren(children, queue, visitor)
     ReactCurrentDispatcher.current = prevDispatcher
-    return flushFrames(queue, visitor)
   })
 }
 
+const flushFrames = (queue: Frame[], visitor: Visitor): Promise<void> => {
+  if (queue.length === 0) {
+    return Promise.resolve()
+  }
+
+  return updateWithFrame(queue.shift(), queue, visitor).then(() =>
+    flushFrames(queue, visitor)
+  )
+}
+
 const defaultVisitor = () => {}
 
 const renderPrepass = (element: Node, visitor?: Visitor): Promise<void> => {

src/visitor.js

@@ -58,12 +58,6 @@ import {
   REACT_LAZY_TYPE
 } from './symbols'
 
-const makeImmediatePromise = () => {
-  return new Promise(resolve => {
-    setImmediate(resolve)
-  })
-}
-
 const render = (
   type: ComponentType<DefaultProps> & ComponentStatics,
   props: DefaultProps,
@@ -171,7 +165,7 @@ const visitLoop = (
 ) => {
   const start = Date.now()
 
-  while (traversalChildren.length > 0 && Date.now() - start <= 10) {
+  while (traversalChildren.length > 0 && Date.now() - start <= 5) {
     const currChildren = traversalChildren[traversalChildren.length - 1]
     const currIndex = traversalIndex[traversalIndex.length - 1]++
 
@@ -199,8 +193,8 @@ export const visitChildren = (
 ) => {
   const traversalChildren: AbstractElement[][] = [init]
   const traversalIndex: number[] = [0]
-  const traversalMap: Array<void | ContextMap> = [undefined]
-  const traversalStore: Array<void | ContextEntry> = [undefined]
+  const traversalMap: Array<void | ContextMap> = [flushPrevContextMap()]
+  const traversalStore: Array<void | ContextEntry> = [flushPrevContextStore()]
 
   visitLoop(
     traversalChildren,
@@ -212,11 +206,11 @@ export const visitChildren = (
   )
 
   if (traversalChildren.length > 0) {
-    queue.push({
+    queue.unshift({
       contextMap: getCurrentContextMap(),
       contextStore: getCurrentContextStore(),
+      thenable: Promise.resolve(),
       kind: 'frame.yield',
-      thenable: makeImmediatePromise(),
       children: traversalChildren,
       index: traversalIndex,
       map: traversalMap,
@@ -232,7 +226,7 @@ export const resumeVisitChildren = (
 ) => {
   setCurrentIdentity(null)
   setCurrentContextMap(frame.contextMap)
-  setCurrentContextMap(frame.contextStore)
+  setCurrentContextStore(frame.contextStore)
 
   visitLoop(frame.children, frame.index, frame.map, frame.store, queue, visitor)
 }