// @flow
import React, { PureComponent } from 'react' // eslint-disable-line import/no-unresolved
import type { Node } from 'react' // eslint-disable-line import/no-unresolved
import type { Persistor } from '../types'

type Props = {
  onBeforeLift?: Function,
  children?: Node,
  loading?: Node,
  persistor: Persistor,
}

type State = {
  bootstrapped: boolean,
}

export class PersistGate extends PureComponent<Props, State> {
  static defaultProps = {
    loading: null,
  }

  state = {
    bootstrapped: false,
  }
  _unsubscribe: ?Function

  componentDidMount() {
    this._unsubscribe = this.props.persistor.subscribe(
      this.handlePersistorState
    )
    this.handlePersistorState()
  }

  handlePersistorState = () => {
    const { persistor } = this.props
    let { bootstrapped } = persistor.getState()
    if (bootstrapped) {
      if (this.props.onBeforeLift) {
        Promise.resolve(this.props.onBeforeLift())
          .then(() => this.setState({ bootstrapped: true }))
          .catch(() => this.setState({ bootstrapped: true }))
      } else {
        this.setState({ bootstrapped: true })
      }
      this._unsubscribe && this._unsubscribe()
    }
  }

  componentWillUnmount() {
    this._unsubscribe && this._unsubscribe()
  }

  render() {
    return this.state.bootstrapped ? this.props.children : this.props.loading
  }
}

@@ -267,24 +267,35 @@
 ldren?: Node
+ %7C Function
 ,%0A  loading?
@@ -1274,24 +1274,142 @@
  render() %7B%0A
+    if (typeof this.props.children === 'function') %7B%0A      return this.props.children(this.state.bootstrapped)%0A    %7D%0A%0A
     return t

make PersistGate take function as a child component (#718)