Modularised application state

app.state.ts

import {Set, Map} from 'immutable';

import {Action, Reducer} from "@ngrx/store";
import {Type} from "@angular/core";

type TypeToken = Type<any>;

export interface StatefulModule<T> {
  ngModule: Type<T>;
  reducer<TState>(state: TState, action: Action<any>): TState;

  /**
   * The initial module state.
   */
  initState?: any;

  /**
   * By default, a module reducer is only applied to actions from it's own module,
   * or actions without an ngModule.
   *
   * Actions defined in the following set will also trigger the module's reducer.
   */
  watch?: Set<TypeToken>;
}

export class AppState {
  constructor(public statefulModules: Map<TypeToken, StatefulModule<any>>,
              public moduleStates: Map<TypeToken,any>) {
  }
}


export class RegisterStatefulModule<T,S> implements Action<StatefulModule<T,S>> {
    static ngModule = AppModule;
    constructor(public payload: StatefulModule<T,S>) {}
}

export function appStateReducer(state: AppState, action: Action<any>) {
  let actionType = actionType(action);

  if (type === RegisterStatefulModule) {
    let module: RegisterStatefulModule<any,any> = action.payload;

    return new AppState(
      ngModules.set(module.ngModule, module),
      moduleState.set(module.ngModule, module.initState)
    );
  }

  let actionModule = actionType(action).ngModule;

  // By default (action type has no ngModule), apply to every state.
  let applyModules = state.statefulModules;

  if (actionModule) {
    // If this action has a module, only apply it to:
    //  - the module it was defined in
    //  - any additional modules that `watch` the action's module.
    applyModules = applyModules
        .filter(module => module.ngModule === actionType || module.watch.contains(actionType));
  }

  let updatedStates = applyModules.map(module => {
    let state = this.moduleStates.get(module.ngModule);
    return module.reducer(state, action);
  });

  return new AppState(this.statefulModules, this.moduleStates.merge(updatedStates));
}

my.module.ts

import {NgModule} from '@angular/core';
import {RegisterStatefulModule} from './app-state';

@NgModule({/* etc. */})
export class MyModule {
  /**
   * This may not work, I didn't actually check angular to determine whether `NgModule` were
   * actually every constructed (or dependency injected). The issue I wrote this for was closed
   * before I got the chance.
   *
   * If not, the `StatefulModule` could be added as a service in `Module.dependencies`.
   */
  constructor(store: Store<AppState>) {
     store.dispatch(new RegisterStatefulModule({
       ngModule: MyModule,
       reducer: myModuleReducer,
       initState: new MyModuleState(),
     });
  }
}