Fix theme and attrs compute

Author: kitten

src/render/styledComponent.js

@@ -10,7 +10,6 @@ import type {
   DefaultProps,
   ForwardRefElement,
   Frame,
-  Visitor,
   ComponentStatics
 } from '../types'
 
@@ -33,15 +32,25 @@ type Attr = void | AttrsFn | { [propName: string]: ?AttrsFn }
 type StyledComponentStatics = {
   styledComponentId: string,
   attrs: Attr | Attr[],
-  target: ComponentType<DefaultProps> & ComponentStatics
+  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: DefaultProps,
-  theme: mixed
-): any => {
+const computeAttrsProps = (input: Attr[], props: any, theme: any): any => {
   const executionContext = { ...props, theme }
 
   const attrs = input.reduce((acc, attr) => {
@@ -63,40 +72,37 @@ const computeAttrsProps = (
     return acc
   }, {})
 
-  return Object.assign(attrs, props)
+  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.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,
-  queue: Frame[],
-  visitor: Visitor
-): Node => {
-  if (styledComponents === undefined) {
-    // styled-components is not installed or incompatible, so the component will have to be
-    // mounted normally
-    const { render } = element.type
-    return mountFunctionComponent(render, element.props, queue, visitor)
-  }
-
+export const mount = (element: ForwardRefElement): Node => {
   // Imitate styled-components' attrs props without computing styles
   const type = ((element.type: any): StyledComponentStatics)
-  const theme = readContextValue(styledComponents.ThemeContext) || {}
   const attrs: Attr[] = Array.isArray(type.attrs) ? type.attrs : [type.attrs]
-  const computedProps = computeProps(element.props, (type: any).defaultProps)
+  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
 
-  if (typeof as !== 'function') {
-    // StyledComponents rendering DOM elements can safely be skipped like normal DOM elements
-    return element.props.children || null
+  // StyledComponents rendering DOM elements can safely be skipped like normal DOM elements
+  if (typeof as === 'string') {
+    return children
   } else {
-    return createElement((as: any), props)
+    delete props.as
+    return createElement((as: any), props, children)
   }
 }

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

@@ -1,7 +1,13 @@
 // @flow
 
 import type { Node, ComponentType } from 'react'
-import { typeOf, shouldConstruct, getChildrenArray } from './element'
+
+import {
+  typeOf,
+  shouldConstruct,
+  getChildrenArray,
+  computeProps
+} from './element'
 
 import {
   mountFunctionComponent,
@@ -64,9 +70,6 @@ import {
 // the event loop is not interrupted for too long
 const YIELD_AFTER_MS = process.env.NODE_ENV !== 'production' ? 20 : 5
 
-// A no-op function for styled-components' ComponentStyle class
-const NOOP_GEN_CLASSNAME = () => ''
-
 const render = (
   type: ComponentType<DefaultProps> & ComponentStatics,
   props: DefaultProps,
@@ -142,10 +145,10 @@ export const visitElement = (
       // that allows quick rendering of them without computing styles
       let child = null
       if (isStyledElement(refElement)) {
-        child = mountStyledComponent(refElement, queue, visitor)
+        child = mountStyledComponent(refElement)
       } else {
-        const { render } = refElement.type
-        const props = refElement.props
+        const { render, defaultProps } = refElement.type
+        const props = computeProps(refElement.props, defaultProps)
         child = mountFunctionComponent(render, props, queue, visitor)
       }
 
@@ -182,7 +185,7 @@ const visitLoop = (
 ) => {
   const start = Date.now()
 
-  while (traversalChildren.length > 0 && Date.now() - start <= YIELD_AFTER_MS) {
+  while (traversalChildren.length > 0) {
     const currChildren = traversalChildren[traversalChildren.length - 1]
     const currIndex = traversalIndex[traversalIndex.length - 1]++
 
@@ -205,7 +208,15 @@ const visitLoop = (
       restoreContextMap(traversalMap.pop())
       restoreContextStore(traversalStore.pop())
     }
+
+    /*
+    if (Date.now() - start > YIELD_AFTER_MS) {
+      return true
+    }
+    */
   }
+
+  return false
 }
 
 export const visitChildren = (
@@ -217,8 +228,7 @@ export const visitChildren = (
   const traversalIndex: number[] = [0]
   const traversalMap: Array<void | ContextMap> = [flushPrevContextMap()]
   const traversalStore: Array<void | ContextEntry> = [flushPrevContextStore()]
-
-  visitLoop(
+  const hasYielded = visitLoop(
     traversalChildren,
     traversalIndex,
     traversalMap,
@@ -227,7 +237,7 @@ export const visitChildren = (
     visitor
   )
 
-  if (traversalChildren.length > 0) {
+  if (hasYielded) {
     queue.unshift({
       contextMap: getCurrentContextMap(),
       contextStore: getCurrentContextStore(),