Fix two event loop yields in a row

Author: kitten

src/__tests__/suspense.test.js

@@ -18,7 +18,7 @@ describe('renderPrepass', () => {
 
       const Outer = () => {
         const start = Date.now()
-        while (Date.now() - start < 21) {}
+        while (Date.now() - start < 40) {}
         return <Inner />
       }
 
@@ -45,10 +45,10 @@ describe('renderPrepass', () => {
         return null
       }
 
-      const Wait = () => {
+      const Wait = props => {
         const start = Date.now()
         while (Date.now() - start < 21) {}
-        return <Inner />
+        return props.children
       }
 
       class Outer extends Component {
@@ -57,7 +57,13 @@ describe('renderPrepass', () => {
         }
 
         render() {
-          return <Wait />
+          return (
+            <Wait>
+              <Wait>
+                <Inner />
+              </Wait>
+            </Wait>
+          )
         }
       }
 

src/visitor.js

@@ -196,15 +196,29 @@ const visitLoop = (
       restoreContextMap(traversalMap.pop())
       restoreContextStore(traversalStore.pop())
     }
-  }
 
-  if (Date.now() - start > YIELD_AFTER_MS) {
-    return true
+    if (Date.now() - start > YIELD_AFTER_MS) {
+      return true
+    }
   }
 
   return false
 }
 
+const makeYieldFrame = (
+  traversalChildren: AbstractElement[][],
+  traversalMap: Array<void | ContextMap>,
+  traversalStore: Array<void | ContextEntry>
+): Frame => ({
+  contextMap: getCurrentContextMap(),
+  contextStore: getCurrentContextStore(),
+  thenable: Promise.resolve(),
+  kind: 'frame.yield',
+  children: traversalChildren,
+  map: traversalMap,
+  store: traversalStore
+})
+
 export const visitChildren = (
   init: AbstractElement[],
   queue: Frame[],
@@ -213,6 +227,7 @@ export const visitChildren = (
   const traversalChildren: AbstractElement[][] = [init]
   const traversalMap: Array<void | ContextMap> = [flushPrevContextMap()]
   const traversalStore: Array<void | ContextEntry> = [flushPrevContextStore()]
+
   const hasYielded = visitLoop(
     traversalChildren,
     traversalMap,
@@ -222,15 +237,9 @@ export const visitChildren = (
   )
 
   if (hasYielded) {
-    queue.unshift({
-      contextMap: getCurrentContextMap(),
-      contextStore: getCurrentContextStore(),
-      thenable: Promise.resolve(),
-      kind: 'frame.yield',
-      children: traversalChildren,
-      map: traversalMap,
-      store: traversalStore
-    })
+    queue.unshift(
+      makeYieldFrame(traversalChildren, traversalMap, traversalStore)
+    )
   }
 }
 
@@ -243,5 +252,15 @@ export const resumeVisitChildren = (
   setCurrentContextMap(frame.contextMap)
   setCurrentContextStore(frame.contextStore)
 
-  visitLoop(frame.children, frame.map, frame.store, queue, visitor)
+  const hasYielded = visitLoop(
+    frame.children,
+    frame.map,
+    frame.store,
+    queue,
+    visitor
+  )
+
+  if (hasYielded) {
+    queue.unshift(makeYieldFrame(frame.children, frame.map, frame.store))
+  }
 }