Add applyPatches function to BaseControllerV2 (#980)

Author: hmalik88

packages/base-controller/src/BaseControllerV2.test.ts

@@ -73,7 +73,12 @@ class CountController extends BaseController<
       state: Draft<CountControllerState>,
     ) => void | CountControllerState,
   ) {
-    super.update(callback);
+    const res = super.update(callback);
+    return res;
+  }
+
+  applyPatches(patches: Patch[]) {
+    super.applyPatches(patches);
   }
 
   destroy() {
@@ -170,6 +175,29 @@ describe('BaseController', () => {
     expect(controller.state).toStrictEqual({ count: 1 });
   });
 
+  it('should return next state, patches and inverse patches after an update', () => {
+    const controller = new CountController({
+      messenger: getCountMessenger(),
+      name: 'CountController',
+      state: { count: 0 },
+      metadata: countControllerStateMetadata,
+    });
+
+    const returnObj = controller.update((draft) => {
+      draft.count += 1;
+    });
+
+    expect(returnObj).not.toBeUndefined();
+    expect(returnObj.nextState).toStrictEqual({ count: 1 });
+    expect(returnObj.patches).toStrictEqual([
+      { op: 'replace', path: ['count'], value: 1 },
+    ]);
+
+    expect(returnObj.inversePatches).toStrictEqual([
+      { op: 'replace', path: ['count'], value: 0 },
+    ]);
+  });
+
   it('should throw an error if update callback modifies draft and returns value', () => {
     const controller = new CountController({
       messenger: getCountMessenger(),
@@ -188,6 +216,55 @@ describe('BaseController', () => {
     );
   });
 
+  it('should allow for applying immer patches to state', () => {
+    const controller = new CountController({
+      messenger: getCountMessenger(),
+      name: 'CountController',
+      state: { count: 0 },
+      metadata: countControllerStateMetadata,
+    });
+
+    const returnObj = controller.update((draft) => {
+      draft.count += 1;
+    });
+
+    controller.applyPatches(returnObj.inversePatches);
+
+    expect(controller.state).toStrictEqual({ count: 0 });
+  });
+
+  it('should inform subscribers of state changes as a result of applying patches', () => {
+    const controllerMessenger = new ControllerMessenger<
+      never,
+      CountControllerEvent
+    >();
+    const controller = new CountController({
+      messenger: getCountMessenger(controllerMessenger),
+      name: 'CountController',
+      state: { count: 0 },
+      metadata: countControllerStateMetadata,
+    });
+    const listener1 = sinon.stub();
+
+    controllerMessenger.subscribe('CountController:stateChange', listener1);
+    const { inversePatches } = controller.update(() => {
+      return { count: 1 };
+    });
+
+    controller.applyPatches(inversePatches);
+
+    expect(listener1.callCount).toStrictEqual(2);
+    expect(listener1.firstCall.args).toStrictEqual([
+      { count: 1 },
+      [{ op: 'replace', path: [], value: { count: 1 } }],
+    ]);
+
+    expect(listener1.secondCall.args).toStrictEqual([
+      { count: 0 },
+      [{ op: 'replace', path: [], value: { count: 0 } }],
+    ]);
+  });
+
   it('should inform subscribers of state changes', () => {
     const controllerMessenger = new ControllerMessenger<
       never,

packages/base-controller/src/BaseControllerV2.ts

@@ -1,4 +1,4 @@
-import { enablePatches, produceWithPatches } from 'immer';
+import { enablePatches, produceWithPatches, applyPatches } from 'immer';
 
 // Imported separately because only the type is used
 // eslint-disable-next-line no-duplicate-imports
@@ -146,12 +146,17 @@ export class BaseController<
    *
    * @param callback - Callback for updating state, passed a draft state
    * object. Return a new state object or mutate the draft to update state.
+   * @returns An object that has the next state, patches applied in the update and inverse patches to
+   * rollback the update.
    */
-  protected update(callback: (state: Draft<S>) => void | S) {
+  protected update(callback: (state: Draft<S>) => void | S): {
+    nextState: S;
+    patches: Patch[];
+    inversePatches: Patch[];
+  } {
     // We run into ts2589, "infinite type depth", if we don't cast
     // produceWithPatches here.
-    // The final, omitted member of the returned tuple are the inverse patches.
-    const [nextState, patches] = (
+    const [nextState, patches, inversePatches] = (
       produceWithPatches as unknown as (
         state: S,
         cb: typeof callback,
@@ -164,6 +169,25 @@ export class BaseController<
       nextState,
       patches,
     );
+
+    return { nextState, patches, inversePatches };
+  }
+
+  /**
+   * Applies immer patches to the current state. The patches come from the
+   * update function itself and can either be normal or inverse patches.
+   *
+   * @param patches - An array of immer patches that are to be applied to make
+   * or undo changes.
+   */
+  protected applyPatches(patches: Patch[]) {
+    const nextState = applyPatches(this.internalState, patches);
+    this.internalState = nextState;
+    this.messagingSystem.publish(
+      `${this.name}:stateChange` as Namespaced<N, any>,
+      nextState,
+      patches,
+    );
   }
 
   /**