index.js

var React = require('react');
var ReactRedux = require('react-redux');

var connect = ReactRedux.connect;
var Provider = ReactRedux.Provider;

var memoizedStore;
var _Promise;
var _debug = false;
var skipMerge = ['initialState', 'initialProps', 'isServer', 'store'];

function initStore(makeStore, req, initialState) {

    // Always make a new store if server
    if (!!req && typeof window === 'undefined') {
        if (!req._store) {
            req._store = makeStore(initialState);
        }
        return req._store;
    }

    // Memoize store if client
    if (!memoizedStore) {
        memoizedStore = makeStore(initialState);
    }

    return memoizedStore;

}

module.exports = function(createStore) {

    var connectArgs = [].slice.call(arguments).slice(1);

    return function(Cmp) {

        // Since provide should always be after connect we connect here
        var ConnectedCmp = (connect.apply(null, connectArgs))(Cmp);

        function WrappedCmp(props) {

            props = props || {};

            var initialState = props.initialState;
            var initialProps = props.initialProps;
            var hasStore = props.store && props.store.dispatch && props.store.getState;
            var store = hasStore
                ? props.store
                : initStore(createStore, {}, initialState); // client case, no store but has initialState

            if (_debug) console.log(Cmp.name, '- 4. WrappedCmp.render', (hasStore ? 'picked up existing one,' : 'created new store with'), 'initialState', initialState);

            // Fix for _document
            var mergedProps = {};
            Object.keys(props).forEach(function(p) { if (!~skipMerge.indexOf(p)) mergedProps[p] = props[p]; });
            Object.keys(initialProps).forEach(function(p) { mergedProps[p] = initialProps[p]; });

            return React.createElement( //FIXME This will create double Provider for _document case
                Provider,
                {store: store},
                React.createElement(ConnectedCmp, mergedProps)
            );

        }

        WrappedCmp.getInitialProps = function(ctx) {

            return new _Promise(function(res) {

                ctx = ctx || {};

                if (_debug) console.log(Cmp.name, '- 1. WrappedCmp.getInitialProps wrapper', (ctx.req && ctx.req._store ? 'takes the req store' : 'creates the store'));

                ctx.isServer = !!ctx.req;
                ctx.store = initStore(createStore, ctx.req);

                res(_Promise.all([
                    ctx.isServer,
                    ctx.store,
                    ctx.req,
                    Cmp.getInitialProps ? Cmp.getInitialProps.call(Cmp, ctx) : {}
                ]));

            }).then(function(arr) {

                if (_debug) console.log(Cmp.name, '- 3. WrappedCmp.getInitialProps has store state', arr[1].getState());

                return {
                    isServer: arr[0],
                    store: arr[1],
                    initialState: arr[1].getState(),
                    initialProps: arr[3]
                };

            });

        };

        return WrappedCmp;

    };

};

module.exports.setPromise = function(Promise) {
    _Promise = Promise;
};

module.exports.setDebug = function(debug) {
    _debug = debug;
};

module.exports.setPromise(Promise);