index.ts

import { TApp, Configuration, EventType, Overmind } from 'overmind'
import {
  // @ts-ignore
  __SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED,
  ClassicComponentClass,
  Component,
  ComponentClass,
  createElement,
  StatelessComponent,
  // @ts-ignore
  useEffect,
  // @ts-ignore
  useState,
  useLayoutEffect,
} from 'react'
import { IMutation } from 'proxy-state-tree'

const IS_PRODUCTION = process.env.NODE_ENV === 'production'

export type IReactComponent<P = any> =
  | StatelessComponent<P>
  | ComponentClass<P>
  | ClassicComponentClass<P>

// Diff / Omit taken from https://github.com/Microsoft/TypeScript/issues/12215#issuecomment-311923766
type Omit<T, K extends keyof T> = Pick<
  T,
  ({ [P in keyof T]: P } &
    { [P in K]: never } & { [x: string]: never; [x: number]: never })[keyof T]
>

export type TConnect<Config extends Configuration> = {
  overmind: {
    state: TApp<Config>['state']
    actions: TApp<Config>['actions']
    addMutationListener: (cb: (mutation: IMutation) => void) => () => void
  }
}

let nextComponentId = 0

export const createHook = <A extends Overmind<Configuration>>(
  overmind: A
): (() => {
  state: A['state']
  actions: A['actions']
  addMutationListener: (cb: (mutation: IMutation) => void) => () => void
}) => {
  let currentComponentInstanceId = 0
  const {
    ReactCurrentOwner,
  } = __SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
  const useCurrentComponent = () => {
    return ReactCurrentOwner &&
      ReactCurrentOwner.current &&
      ReactCurrentOwner.current.elementType
      ? ReactCurrentOwner.current.elementType
      : {}
  }

  return () => {
    const component = useCurrentComponent()
    const name = component.name
    component.__componentId =
      typeof component.__componentId === 'undefined'
        ? nextComponentId++
        : component.__componentId
    const [tree, updateComponent] = useState<any>(() =>
      (overmind as any).proxyStateTree.getTrackStateTree()
    )

    if (IS_PRODUCTION) {
      tree.track(() => {
        updateComponent((state) => state)
      })

      useEffect(
        () => () => {
          ;(overmind as any).proxyStateTree.disposeTree(tree)
        },
        []
      )
    } else {
      const [debugging] = useState<any>(() => ({
        isFirstRender: true,
        currentFlushId: 0,
        componentInstanceId: currentComponentInstanceId++,
      }))

      tree.track((mutations, paths, flushId) => {
        debugging.currentFlushId = flushId
        updateComponent((state) => state)
      })

      useLayoutEffect(() => {
        if (debugging.isFirstRender) {
          overmind.eventHub.emitAsync(EventType.COMPONENT_ADD, {
            componentId: component.__componentId,
            componentInstanceId: debugging.componentInstanceId,
            name,
            paths: Array.from(tree.pathDependencies) as any,
          })
          debugging.isFirstRender = false
        } else {
          overmind.eventHub.emitAsync(EventType.COMPONENT_UPDATE, {
            componentId: component.__componentId,
            componentInstanceId: debugging.componentInstanceId,
            name,
            flushId: debugging.currentFlushId,
            paths: Array.from(tree.pathDependencies as Set<string>),
          })
        }
      })

      useEffect(
        () => () => {
          ;(overmind as any).proxyStateTree.disposeTree(tree)
          overmind.eventHub.emitAsync(EventType.COMPONENT_REMOVE, {
            componentId: component.__componentId,
            componentInstanceId: debugging.componentInstanceId,
            name,
          })
        },
        []
      )
    }

    return {
      state: tree.state,
      actions: overmind.actions,
      addMutationListener: overmind.addMutationListener,
    }
  }
}

export const createConnect = <A extends Overmind<Configuration>>(
  overmind: A
) => {
  return <Props>(
    component: IReactComponent<
      Props & { overmind: { state: A['state']; actions: A['actions'] } }
    >
  ): IReactComponent<Omit<Props & TConnect<A>, keyof TConnect<A>>> => {
    let componentInstanceId = 0
    const name = component.name
    const populatedComponent = component as any
    populatedComponent.__componentId =
      typeof populatedComponent.__componentId === 'undefined'
        ? nextComponentId++
        : populatedComponent.__componentId
    const isClassComponent =
      component.prototype && typeof component.prototype.render === 'function'

    if (isClassComponent) {
      const originalRender = component.prototype.render
      component.prototype.render = function() {
        this.props.overmind &&
          this.props.overmind.tree.track(this.props.overmind.onUpdate)
        return originalRender.call(this)
      }
    }

    if (IS_PRODUCTION) {
      class HOC extends Component {
        tree = (overmind as any).proxyStateTree.getTrackStateTree()
        state: {
          overmind: any
        }
        constructor(props) {
          super(props)
          this.state = {
            overmind: {
              state: this.tree.state,
              actions: overmind.actions,
              addMutationListener: overmind.addMutationListener,
              onUpdate: this.onUpdate,
              tree: this.tree,
            },
          }
        }
        componentWillUnmount() {
          ;(overmind as any).proxyStateTree.disposeTree(this.tree)
        }
        onUpdate = () => {
          this.setState({
            overmind: {
              state: this.tree.state,
              actions: overmind.actions,
              addMutationListener: overmind.addMutationListener,
              onUpdate: this.onUpdate,
              tree: this.tree,
            },
          })
        }
        render() {
          if (isClassComponent) {
            return createElement(component, {
              ...this.props,
              overmind: this.state.overmind,
            } as any)
          }

          this.tree.track(this.onUpdate)

          return createElement(component, {
            ...this.props,
            overmind: this.state.overmind,
          } as any)
        }
      }

      return HOC as any
    } else {
      class HOC extends Component {
        tree = (overmind as any).proxyStateTree.getTrackStateTree()
        componentInstanceId = componentInstanceId++
        currentFlushId = 0
        state: {
          overmind: any
        }
        constructor(props) {
          super(props)
          this.state = {
            overmind: {
              state: this.tree.state,
              actions: overmind.actions,
              addMutationListener: overmind.addMutationListener,
              onUpdate: this.onUpdate,
              tree: this.tree,
            },
          }
        }
        componentDidMount() {
          overmind.eventHub.emitAsync(EventType.COMPONENT_ADD, {
            componentId: populatedComponent.__componentId,
            componentInstanceId: this.componentInstanceId,
            name,
            paths: Array.from(this.tree.pathDependencies) as any,
          })
        }
        componentDidUpdate() {
          overmind.eventHub.emitAsync(EventType.COMPONENT_UPDATE, {
            componentId: populatedComponent.__componentId,
            componentInstanceId: this.componentInstanceId,
            name,
            flushId: this.currentFlushId,
            paths: Array.from(this.tree.pathDependencies as Set<string>),
          })
        }
        componentWillUnmount() {
          ;(overmind as any).proxyStateTree.disposeTree(this.tree)
          overmind.eventHub.emitAsync(EventType.COMPONENT_REMOVE, {
            componentId: populatedComponent.__componentId,
            componentInstanceId: this.componentInstanceId,
            name,
          })
        }
        onUpdate = (mutatons, paths, flushId) => {
          this.currentFlushId = flushId
          this.setState({
            overmind: {
              state: this.tree.state,
              actions: overmind.actions,
              addMutationListener: overmind.addMutationListener,
              onUpdate: this.onUpdate,
              tree: this.tree,
            },
          })
        }
        render() {
          if (isClassComponent) {
            return createElement(component, {
              ...this.props,
              overmind: this.state.overmind,
            } as any)
          }

          this.tree.track(this.onUpdate)

          return createElement(component, {
            ...this.props,
            overmind: this.state.overmind,
          } as any)
        }
      }

      Object.defineProperties(HOC, {
        name: {
          value: 'Connect' + (component.displayName || component.name || ''),
        },
      })

      return HOC as any
    }
  }
}