react-side-effect using createContext

react-side-effect.jsx

import React, { Component, PureComponent, createContext } from "react";
import flatten from "lodash/flatten";

function flattenPropsTree(arr) {
  return flatten(arr.map(a => [a.props, ...flattenPropsTree(a.children)]));
}

function getPropsTree(arr) {
  return arr.map(a => ({
    props: a.props,
    children: getPropsTree(a.state.instances)
  }));
}

export default function withSideEffect(
  reducePropsToState,
  handleStateChangeOnClient,
  mapStateOnServer
) {
  if (typeof reducePropsToState !== "function") {
    throw new Error("Expected reducePropsToState to be a function.");
  }
  if (typeof handleStateChangeOnClient !== "function") {
    throw new Error("Expected handleStateChangeOnClient to be a function.");
  }
  if (
    typeof mapStateOnServer !== "undefined" &&
    typeof mapStateOnServer !== "function"
  ) {
    throw new Error(
      "Expected mapStateOnServer to either be undefined or a function."
    );
  }

  function getDisplayName(WrappedComponent) {
    return WrappedComponent.displayName || WrappedComponent.name || "Component";
  }

  return function wrap(WrappedComponent) {
    if (typeof WrappedComponent !== "function") {
      throw new Error("Expected WrappedComponent to be a React component.");
    }

    const SideEffectContext = createContext();

    class SideEffectProvider extends Component {
      state = {
        instances: [],
        count: 0
      };

      register = instance => {
        this.setState(prevState => ({
          instances: [...prevState.instances, instance]
        }));

        return {
          unregister: () => {
            this.setState(prevState => ({
              instances: prevState.instances.filter(inst => inst !== instance)
            }));
          },
          onUpdate: () => {
            this.setState(prevState => ({ count: prevState.count + 1 }));
          }
        };
      };

      render() {
        return (
          <SideEffectContext.Provider value={{ register: this.register }}>
            {this.props.children}
          </SideEffectContext.Provider>
        );
      }

      componentDidUpdate() {
        const props = flattenPropsTree(getPropsTree(this.state.instances));
        let state = reducePropsToState(props);

        handleStateChangeOnClient(state);
      }
    }

    class SideEffect extends PureComponent {
      static Provider = SideEffectProvider;

      state = {
        instances: [],
        count: 0
      };

      // Try to use displayName of wrapped component
      static displayName = `SideEffect(${getDisplayName(WrappedComponent)})`;

      componentDidUpdate(prevProps, prevState) {
        this.onUpdate();
      }

      componentDidMount() {
        const { unregister, onUpdate } = this.registerSelf(this);

        this.unregister = unregister;
        this.onUpdate = onUpdate;
      }

      componentWillUnmount() {
        this.unregister();
      }

      register = instance => {
        this.setState(prevState => ({
          instances: [...prevState.instances, instance]
        }));

        return {
          unregister: () => {
            this.setState(prevState => ({
              instances: prevState.instances.filter(inst => inst !== instance)
            }));
          },
          onUpdate: () => {
            this.setState(prevState => ({ count: prevState.count + 1 }));
          }
        };
      };

      render() {
        return (
          <SideEffectContext.Consumer>
            {value => {
              this.registerSelf = value.register;

              return (
                <SideEffectContext.Provider value={{ register: this.register }}>
                  <WrappedComponent {...this.props} />
                </SideEffectContext.Provider>
              );
            }}
          </SideEffectContext.Consumer>
        );
      }
    }

    return SideEffect;
  };
}