Merge pull request #2 from FormidableLabs/event-loop-yielding

Yield to event loop when synchronous work takes too long

Author: kitten

package.json

@@ -35,6 +35,11 @@
       "@babel/preset-react"
     ]
   },
+  "jest": {
+    "globals": {
+      "__DEV__": true
+    }
+  },
   "lint-staged": {
     "**/*.js": [
       "flow focus-check",
@@ -70,7 +75,8 @@
     "prettier": "^1.16.4",
     "react": "^16.8.4",
     "react-dom": "^16.8.4",
-    "react-is": "^16.8.4"
+    "react-is": "^16.8.4",
+    "styled-components": "^4.2.0"
   },
   "dependencies": {
     "object-is": "^1.0.1"

src/__tests__/suspense.test.js

@@ -12,6 +12,40 @@ import React, {
 import renderPrepass from '..'
 
 describe('renderPrepass', () => {
+  describe('event loop yielding', () => {
+    it('yields to the event loop when work is taking too long', () => {
+      const Inner = jest.fn(() => null)
+
+      const Outer = () => {
+        const start = Date.now()
+        while (Date.now() - start < 21) {}
+        return <Inner />
+      }
+
+      const render$ = renderPrepass(<Outer />)
+
+      expect(Inner).toHaveBeenCalledTimes(0)
+
+      setImmediate(() => {
+        setImmediate(() => {
+          expect(Inner).toHaveBeenCalledTimes(1)
+        })
+      })
+
+      return render$.then(() => {
+        expect(Inner).toHaveBeenCalledTimes(1)
+      })
+    })
+
+    it('does not yields when work is below the threshold', () => {
+      const Inner = jest.fn(() => null)
+      const Outer = () => <Inner />
+      const render$ = renderPrepass(<Outer />)
+
+      expect(Inner).toHaveBeenCalledTimes(1)
+    })
+  })
+
   describe('function components', () => {
     it('supports suspending subtrees', () => {
       const value = {}

src/__tests__/visitor.test.js

@@ -12,6 +12,7 @@ import React, {
 } from 'react'
 
 import { createPortal } from 'react-dom'
+import styled from 'styled-components'
 
 import {
   Dispatcher,
@@ -105,6 +106,25 @@ describe('visitElement', () => {
     expect(children[1].type).toBe(Noop)
   })
 
+  it('walks StyledComponent DOM elements', () => {
+    const Comp = styled.div``
+    const children = visitElement(
+      <Comp>
+        <Noop />
+      </Comp>,
+      [],
+      () => {}
+    )
+    expect(children.length).toBe(1)
+    expect(children[0].type).toBe(Noop)
+  })
+
+  it('walks StyledComponent wrapper elements', () => {
+    const Comp = styled(Noop)``
+    const children = visitElement(<Comp />, [], () => {})
+    expect(children.length).toBe(1)
+  })
+
   it('walks Providers and Consumers', () => {
     const Context = createContext('default')
     const leaf = jest.fn().mockReturnValue(null)

src/index.js

@@ -1,8 +1,8 @@
 // @flow
 
 import React, { type Node, type Element } from 'react'
-import type { Visitor, Frame, AbstractElement } from './types'
-import { visitChildren } from './visitor'
+import type { Visitor, YieldFrame, Frame, AbstractElement } from './types'
+import { visitChildren, resumeVisitChildren } from './visitor'
 import { getChildrenArray } from './element'
 
 import {
@@ -28,13 +28,25 @@ 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
@@ -55,11 +67,19 @@ const flushFrames = (queue: Frame[], visitor: Visitor): Promise<void> => {
     // 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/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/element.js

@@ -103,7 +103,10 @@ export type MemoElement = {
 export type ForwardRefElement = {
   type: {
     render: ComponentType<DefaultProps> & ComponentStatics,
-    $$typeof: typeof REACT_FORWARD_REF_TYPE
+    $$typeof: typeof REACT_FORWARD_REF_TYPE,
+    // styled-components specific properties
+    styledComponentId?: string,
+    target?: ComponentType<mixed> | string
   },
   props: DefaultProps,
   $$typeof: typeof REACT_ELEMENT_TYPE

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