Fixing partial state rehydration with lazy loaded modules

spec/index_spec.ts

@@ -1,6 +1,6 @@
 declare var beforeEachProviders, it, describe, expect, inject;
 require('es6-shim');
-import { syncStateUpdate, rehydrateApplicationState, dateReviver } from '../src/index';
+import { syncStateUpdate, rehydrateApplicationState, dateReviver, mergePartialStates } from '../src/index';
 import *  as CryptoJS from 'crypto-js';
 
 // Very simple classes to test serialization options.  They cover string, number, date, and nested classes
@@ -374,4 +374,21 @@ describe('ngrxLocalStorage', () => {
         expect(t1 instanceof TypeA).toBeTruthy();
         expect(finalState.simple instanceof TypeA).toBeFalsy();
     });
-});
+
+    it('merges partial states', () => {
+        const keyDef = ['foo', { bar: ['baz'] }];
+        expect(
+          mergePartialStates(
+            keyDef,
+            { bar: {baz: 123, nitch: 456} },
+            { foo: 'hello', bar: {baz: 1} },
+          )
+        ).toEqual({
+          foo: 'hello',
+          bar: {
+            baz: 1,
+            nitch: 456
+          }
+        });
+    });
+});

src/index.ts

@@ -232,7 +232,7 @@ export const localStorageSync = (config: LocalStorageConfig) => (
       (action.type === INIT_ACTION || action.type === UPDATE_ACTION) &&
       rehydratedState
     ) {
-      state = Object.assign({}, state, rehydratedState);
+      state = mergePartialStates(stateKeys, state, rehydratedState);
     }
     const nextState = reducer(state, action);
     syncStateUpdate(
@@ -246,6 +246,37 @@ export const localStorageSync = (config: LocalStorageConfig) => (
   };
 };
 
+export const mergePartialStates = (stateKeys: (string | { [idx: string]: null | undefined | string[] })[], currentState: {}, rehydratedState: {}) => {
+  let finalState = { ...currentState };
+
+  stateKeys.map((key) => {
+    if ( typeof key === 'string' ) {
+      finalState = {
+        ...finalState,
+        [key]: rehydratedState[key],
+      };
+    } else {
+      const name = Object.keys(key)[0];
+      const newValues = {};
+
+      const keys = key[name]
+      if (keys) {
+        keys.map((key) => newValues[key] = rehydratedState[name][key])
+      }
+
+      finalState = {
+        ...finalState,
+        [name]: {
+          ...finalState[name],
+          ...newValues,
+        },
+      };
+    }
+  });
+
+  return finalState;
+}
+
 /*
     @deprecated: Use localStorageSync(LocalStorageConfig)