Fix element visitor for yielding

Author: kitten

src/index.js

@@ -56,16 +56,16 @@ const updateWithFrame = (
  // Update the component after we've suspended to rerender it,
  // at which point we'll actually get its children
  if (frame.kind === 'frame.class') {
- children = getChildrenArray(updateClassComponent(queue, frame))
+ children = updateClassComponent(queue, frame)
  } else if (frame.kind === 'frame.hooks') {
- children = getChildrenArray(updateFunctionComponent(queue, frame))
+ children = updateFunctionComponent(queue, frame)
  } else if (frame.kind === 'frame.lazy') {
- children = getChildrenArray(updateLazyComponent(queue, frame))
+ children = updateLazyComponent(queue, frame)
  }
 
  // Now continue walking the previously suspended component's
  // children (which might also suspend)
- visitChildren(children, queue, visitor)
+ visitChildren(getChildrenArray(children), queue, visitor)
  ReactCurrentDispatcher.current = prevDispatcher
  })
 }
@@ -80,11 +80,11 @@ const flushFrames = (queue: Frame[], visitor: Visitor): Promise<void> => {
  )
 }
 
-const defaultVisitor = () => {}
+const defaultVisitor = () => undefined
 
 const renderPrepass = (element: Node, visitor?: Visitor): Promise<void> => {
  const queue: Frame[] = []
- let fn = visitor !== undefined ? visitor : defaultVisitor
+ const fn = visitor !== undefined ? visitor : defaultVisitor
 
  // Context state is kept globally and is modified in-place.
  // Before we start walking the element tree we need to reset

src/render/classComponent.js

@@ -151,17 +151,15 @@ export const mount = (
  props: DefaultProps,
  queue: Frame[],
  visitor: Visitor,
- element?: UserElement
+ element: UserElement
 ) => {
  setCurrentIdentity(null)
- const instance = createInstance(type, props)
 
- if (element !== undefined) {
- const p = visitor(element, instance)
- if (typeof p === 'object' && p !== null && typeof p.then === 'function') {
- queue.push(makeFrame(type, instance, p))
- return null
- }
+ const instance = createInstance(type, props)
+ const promise = visitor(element, instance)
+ if (promise) {
+ queue.push(makeFrame(type, instance, promise))
+ return null
  }
 
  return render(type, instance, queue)

src/render/functionComponent.js

@@ -61,8 +61,6 @@ const render = (
 
  queue.push(makeFrame(type, props, error))
  return null
- } finally {
- setCurrentIdentity(null)
  }
 }
 
@@ -72,17 +70,15 @@ export const mount = (
  props: DefaultProps,
  queue: Frame[],
  visitor: Visitor,
- element?: UserElement
+ element: UserElement
 ): Node => {
  setFirstHook(null)
  setCurrentIdentity(makeIdentity())
 
- if (element !== undefined) {
- const p = visitor(element)
- if (typeof p === 'object' && p !== null && typeof p.then === 'function') {
- queue.push(makeFrame(type, props, p))
- return null
- }
+ const promise = visitor(element)
+ if (promise) {
+ queue.push(makeFrame(type, props, promise))
+ return null
  }
 
  return render(type, props, queue)

src/render/styledComponent.js

@@ -7,6 +7,7 @@ import { readContextValue } from '../internals'
 import { mount as mountFunctionComponent } from './functionComponent'
 
 import type {
+ AbstractElement,
  DefaultProps,
  ForwardRefElement,
  Frame,
@@ -89,7 +90,7 @@ export const isStyledElement = (element: ForwardRefElement): boolean %checks =>
 /** This is an optimised faux mounting strategy for StyledComponents.
  It is only enabled when styled-components is installed and the component
  can safely be skipped */
-export const mount = (element: ForwardRefElement): Node => {
+export const mount = (element: ForwardRefElement): AbstractElement[] => {
  // Imitate styled-components' attrs props without computing styles
  const type = ((element.type: any): StyledComponentStatics)
  const attrs: Attr[] = Array.isArray(type.attrs) ? type.attrs : [type.attrs]
@@ -101,9 +102,9 @@ export const mount = (element: ForwardRefElement): Node => {
 
  // StyledComponents rendering DOM elements can safely be skipped like normal DOM elements
  if (typeof as === 'string') {
- return children
+ return getChildrenArray(children)
  } else {
  delete props.as
- return createElement((as: any), props, children)
+ return [(createElement((as: any), props, children): any)]
  }
 }

src/types/frames.js

@@ -39,7 +39,6 @@ export type HooksFrame = BaseFrame & {
 export type YieldFrame = BaseFrame & {
  kind: 'frame.yield',
  children: AbstractElement[][],
- index: number[],
  map: Array<void | ContextMap>,
  store: Array<void | ContextEntry>
 }

src/visitor.js

@@ -1,6 +1,6 @@
 // @flow
 
-import type { Node, ComponentType } from 'react'
+import { type Node, type ComponentType, createElement } from 'react'
 
 import {
  typeOf,
@@ -75,7 +75,7 @@ const render = (
  props: DefaultProps,
  queue: Frame[],
  visitor: Visitor,
- element?: UserElement
+ element: UserElement
 ) => {
  return shouldConstruct(type)
  ? mountClassComponent(type, props, queue, visitor, element)
@@ -134,7 +134,8 @@ export const visitElement = (
  case REACT_MEMO_TYPE: {
  const memoElement = ((element: any): MemoElement)
  const type = memoElement.type.type
- const child = render(type, memoElement.props, queue, visitor)
+ const fauxElement = (createElement((type: any), memoElement.props): any)
+ const child = render(type, memoElement.props, queue, visitor, fauxElement)
  return getChildrenArray(child)
  }
 
@@ -143,15 +144,14 @@ export const visitElement = (
 
  // If we find a StyledComponent, we trigger a specific optimisation
  // that allows quick rendering of them without computing styles
- let child = null
  if (isStyledElement(refElement)) {
- child = mountStyledComponent(refElement)
- } else {
- const { render, defaultProps } = refElement.type
- const props = computeProps(refElement.props, defaultProps)
- child = mountFunctionComponent(render, props, queue, visitor)
+ return mountStyledComponent(refElement)
  }
 
+ const { render: type, defaultProps } = refElement.type
+ const props = computeProps(refElement.props, defaultProps)
+ const fauxElement = (createElement((render: any), props): any)
+ const child = render(type, props, queue, visitor, fauxElement)
  return getChildrenArray(child)
  }
 
@@ -177,43 +177,29 @@ export const visitElement = (
 
 const visitLoop = (
  traversalChildren: AbstractElement[][],
- traversalIndex: number[],
  traversalMap: Array<void | ContextMap>,
  traversalStore: Array<void | ContextEntry>,
  queue: Frame[],
  visitor: Visitor
-) => {
+): boolean => {
  const start = Date.now()
 
  while (traversalChildren.length > 0) {
- const currChildren = traversalChildren[traversalChildren.length - 1]
- const currIndex = traversalIndex[traversalIndex.length - 1]++
-
- if (currIndex < currChildren.length) {
- const element = currChildren[currIndex]
+ const element = traversalChildren[traversalChildren.length - 1].shift()
+ if (element !== undefined) {
  const children = visitElement(element, queue, visitor)
-
- if (children.length > 0) {
- traversalChildren.push(children)
- traversalIndex.push(0)
- traversalMap.push(flushPrevContextMap())
- traversalStore.push(flushPrevContextStore())
- } else {
- restoreContextMap(flushPrevContextMap())
- restoreContextStore(flushPrevContextStore())
- }
+ traversalChildren.push(children)
+ traversalMap.push(flushPrevContextMap())
+ traversalStore.push(flushPrevContextStore())
  } else {
  traversalChildren.pop()
- traversalIndex.pop()
  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
@@ -225,12 +211,10 @@ export const visitChildren = (
  visitor: Visitor
 ) => {
  const traversalChildren: AbstractElement[][] = [init]
- const traversalIndex: number[] = [0]
  const traversalMap: Array<void | ContextMap> = [flushPrevContextMap()]
  const traversalStore: Array<void | ContextEntry> = [flushPrevContextStore()]
  const hasYielded = visitLoop(
  traversalChildren,
- traversalIndex,
  traversalMap,
  traversalStore,
  queue,
@@ -244,7 +228,6 @@ export const visitChildren = (
  thenable: Promise.resolve(),
  kind: 'frame.yield',
  children: traversalChildren,
- index: traversalIndex,
  map: traversalMap,
  store: traversalStore
  })
@@ -260,5 +243,5 @@ export const resumeVisitChildren = (
  setCurrentContextMap(frame.contextMap)
  setCurrentContextStore(frame.contextStore)
 
- visitLoop(frame.children, frame.index, frame.map, frame.store, queue, visitor)
+ visitLoop(frame.children, frame.map, frame.store, queue, visitor)
 }