Angular 2 ngrx/store HMR

app.module.ts

import 'rxjs/add/operator/take';

import {
  Action,
  ActionReducer,
  combineReducers,
  Store,
  StoreModule
} from '@ngrx/store';
import { compose } from '@ngrx/core/compose';

// AppState is an interface custom to your application
type StoreType = {
  state: AppState,
  restoreInputValues: () => void,
  disposeOldHosts: () => void
};

// Reducer for setting the root STATE for HMR
const stateSetter: ActionReducer<any> = (reducer: ActionReducer<any>) => {
  return (state, action ) => {
    if (action.type === 'SET_ROOT_STATE') {
      return action.payload;
    }
    return reducer(state, action);
  };
};

const rootReducer = compose(stateSetter, combineReducers)({
  // Add your reducers here
});

// ...

export class AppModule {
  constructor(public appRef: ApplicationRef, private _store: Store<any>) {}

  hmrOnInit(store: StoreType) {
    if (!store || !store.state) return;
    console.log('HMR store', JSON.stringify(store, null, 2));

    // set state
    // restore state by dispatch a SET_ROOT_STATE action
    if (store.state) {
      this._store.dispatch(<Action>{ type: 'SET_ROOT_STATE', payload: store.state });
    }

    // set input values
    if ('restoreInputValues' in store) {
      let restoreInputValues = store.restoreInputValues;
      setTimeout(restoreInputValues);
    }

    this.appRef.tick();
    delete store.state;
    delete store.restoreInputValues;
  }

  hmrOnDestroy(store: StoreType) {
    const cmpLocation = this.appRef.components.map(cmp => cmp.location.nativeElement);
    // save state
    this._store.take(1).subscribe(s => store.state = s);

    // recreate root elements
    store.disposeOldHosts = createNewHosts(cmpLocation);

    // save input values
    store.restoreInputValues  = createInputTransfer();

    // remove styles
    removeNgStyles();
  }

  hmrAfterDestroy(store: StoreType) {
    // display new elements
    store.disposeOldHosts();
    delete store.disposeOldHosts;
  }

}

Yes, I removed Angular and AngularClass/hmr imports for brevity.