Add proper styled-components optimisation (#3)

Add proper styled-components optimisation

Author: kitten

.flowconfig

@@ -4,6 +4,7 @@
 
 [untyped]
 .*/node_modules/react-is/.*
+.*/node_modules/styled-components/.*
 
 [libs]
 types/

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 />
       }
 
@@ -37,6 +37,78 @@ describe('renderPrepass', () => {
       })
     })
 
+    it('preserves the correct legacy context values across yields', () => {
+      let called = false
+      const Inner = (_, context) => {
+        expect(context.test).toBe(123)
+        called = true
+        return null
+      }
+
+      const Wait = props => {
+        const start = Date.now()
+        while (Date.now() - start < 21) {}
+        return props.children
+      }
+
+      class Outer extends Component {
+        getChildContext() {
+          return { test: 123 }
+        }
+
+        render() {
+          return (
+            <Wait>
+              <Wait>
+                <Inner />
+              </Wait>
+            </Wait>
+          )
+        }
+      }
+
+      Inner.contextTypes = { test: null }
+      Outer.childContextTypes = { test: null }
+
+      const render$ = renderPrepass(<Outer />)
+      expect(called).toBe(false)
+      return render$.then(() => {
+        expect(called).toBe(true)
+      })
+    })
+
+    it('preserves the correct context values across yields', () => {
+      const Context = createContext(null)
+
+      let called = false
+      const Inner = () => {
+        const value = useContext(Context)
+        expect(value).toBe(123)
+        called = true
+        return null
+      }
+
+      const Wait = () => {
+        const start = Date.now()
+        while (Date.now() - start < 21) {}
+        return <Inner />
+      }
+
+      const Outer = () => {
+        return (
+          <Context.Provider value={123}>
+            <Wait />
+          </Context.Provider>
+        )
+      }
+
+      const render$ = renderPrepass(<Outer />)
+      expect(called).toBe(false)
+      return render$.then(() => {
+        expect(called).toBe(true)
+      })
+    })
+
     it('does not yields when work is below the threshold', () => {
       const Inner = jest.fn(() => null)
       const Outer = () => <Inner />

src/__tests__/visitor.test.js

@@ -31,6 +31,10 @@ const {
   ReactCurrentDispatcher
 } = (React: any).__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
 
+const {
+  StyleSheet
+} = require('styled-components').__DO_NOT_USE_OR_YOU_WILL_BE_HAUNTED_BY_SPOOKY_GHOSTS
+
 let prevDispatcher = null
 
 beforeEach(() => {
@@ -48,6 +52,10 @@ afterEach(() => {
 const Noop = () => null
 
 describe('visitElement', () => {
+  beforeEach(() => {
+    StyleSheet.reset(true)
+  })
+
   it('walks Fragments', () => {
     const element = (
       <Fragment>
@@ -115,14 +123,32 @@ describe('visitElement', () => {
       [],
       () => {}
     )
+
     expect(children.length).toBe(1)
     expect(children[0].type).toBe(Noop)
+    expect(StyleSheet.master.tags.length).toBe(1)
+
+    const tag = StyleSheet.master.tags[0]
+    expect(tag.css().trim()).toBe('')
+
+    expect(Object.keys(StyleSheet.master.deferred)).toEqual([
+      expect.any(String)
+    ])
   })
 
   it('walks StyledComponent wrapper elements', () => {
     const Comp = styled(Noop)``
     const children = visitElement(<Comp />, [], () => {})
+
     expect(children.length).toBe(1)
+    expect(StyleSheet.master.tags.length).toBe(1)
+
+    const tag = StyleSheet.master.tags[0]
+    expect(tag.css().trim()).toBe('')
+
+    expect(Object.keys(StyleSheet.master.deferred)).toEqual([
+      expect.any(String)
+    ])
   })
 
   it('walks Providers and Consumers', () => {

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/index.js

@@ -22,3 +22,8 @@ export {
   mount as mountClassComponent,
   update as updateClassComponent
 } from './classComponent'
+
+export {
+  isStyledElement,
+  mount as mountStyledComponent
+} from './styledComponent'

src/render/styledComponent.js

@@ -0,0 +1,110 @@
+// @flow
+
+import { createElement, type ComponentType, type Node } from 'react'
+
+import { getChildrenArray, computeProps } from '../element'
+import { readContextValue } from '../internals'
+import { mount as mountFunctionComponent } from './functionComponent'
+
+import type {
+  AbstractElement,
+  DefaultProps,
+  ForwardRefElement,
+  Frame,
+  ComponentStatics
+} from '../types'
+
+let styledComponents: any
+try {
+  styledComponents = require('styled-components')
+
+  if (
+    styledComponents.__DO_NOT_USE_OR_YOU_WILL_BE_HAUNTED_BY_SPOOKY_GHOSTS ===
+      undefined ||
+    styledComponents.ThemeContext === undefined
+  ) {
+    styledComponents = undefined
+  }
+} catch (_error) {}
+
+type AttrsFn = (context: mixed) => DefaultProps
+type Attr = void | AttrsFn | { [propName: string]: ?AttrsFn }
+
+type StyledComponentStatics = {
+  styledComponentId: string,
+  attrs: Attr | Attr[],
+  target: ComponentType<DefaultProps> & ComponentStatics,
+  defaultProps?: Object
+}
+
+/** Determines a StyledComponent's theme taking defaults into account */
+const computeTheme = (props: Object, defaultProps: Object): Object => {
+  const defaultTheme = defaultProps ? defaultProps.theme : undefined
+  const isDefaultTheme = defaultTheme ? props.theme === defaultTheme : false
+
+  if (props.theme && !isDefaultTheme) {
+    return props.theme
+  } else {
+    const contextTheme = readContextValue(styledComponents.ThemeContext)
+    return contextTheme || defaultTheme
+  }
+}
+
+/** Computes a StyledComponent's props with attributes */
+const computeAttrsProps = (input: Attr[], props: any, theme: any): any => {
+  const executionContext = { ...props, theme }
+
+  const attrs = input.reduce((acc, attr) => {
+    if (typeof attr === 'function') {
+      return Object.assign(acc, attr(executionContext))
+    } else if (typeof attr !== 'object' || attr === null) {
+      return acc
+    }
+
+    for (const key in attr) {
+      const attrProp = attr[key]
+      if (typeof attrProp === 'function') {
+        acc[key] = attrProp(executionContext)
+      } else if (attr.hasOwnProperty(key)) {
+        acc[key] = attrProp
+      }
+    }
+
+    return acc
+  }, {})
+
+  const newProps = (Object.assign(attrs, props): any)
+  newProps.className = props.className || ''
+  newProps.style = props.style
+    ? Object.assign({}, attrs.style, props.style)
+    : attrs.style
+  return newProps
+}
+
+/** Checks whether a ForwardRefElement is a StyledComponent element */
+export const isStyledElement = (element: ForwardRefElement): boolean %checks =>
+  styledComponents !== undefined &&
+  typeof element.type.target !== undefined &&
+  typeof element.type.styledComponentId === 'string'
+
+/** 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): 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]
+  const computedProps = computeProps(element.props, type.defaultProps)
+  const theme = computeTheme(element.props, type)
+  const props = computeAttrsProps(attrs, computedProps, theme)
+  const as = props.as || type.target
+  const children = computedProps.children || null
+
+  // StyledComponents rendering DOM elements can safely be skipped like normal DOM elements
+  if (typeof as === 'string') {
+    return getChildrenArray(children)
+  } else {
+    delete props.as
+    return [(createElement((as: any), props, children): any)]
+  }
+}

src/types/element.js

@@ -104,6 +104,7 @@ export type ForwardRefElement = {
   type: {
     render: ComponentType<DefaultProps> & ComponentStatics,
     $$typeof: typeof REACT_FORWARD_REF_TYPE,
+    defaultProps?: Object,
     // styled-components specific properties
     styledComponentId?: string,
     target?: ComponentType<mixed> | string

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>
 }