chore: adding batch function (#22)

Author: divdavem

src/index.spec.ts

@@ -1,4 +1,13 @@
-import { derived, DerivedStore, Store, SubscribableStore, writable, get, readable } from './index';
+import {
+  derived,
+  DerivedStore,
+  Store,
+  SubscribableStore,
+  writable,
+  get,
+  readable,
+  batch,
+} from './index';
 import { from } from 'rxjs';
 import { Component, Injectable } from '@angular/core';
 import { TestBed } from '@angular/core/testing';
@@ -221,6 +230,20 @@ describe('stores', () => {
       expect(values).toEqual([0, 1]); // unsubscribe worked, value 2 was not received
     });
 
+    it('should work when changing a store in the listener of the store', () => {
+      const store = writable(0);
+      const calls: number[] = [];
+      const unsubscribe = store.subscribe((n) => {
+        calls.push(n);
+        if (n < 10) {
+          store.set(n + 1);
+        }
+      });
+      expect(calls).toEqual([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+      expect(get(store)).toEqual(10);
+      unsubscribe();
+    });
+
     it('should be compatible with rxjs "from" (writable)', () => {
       const store = writable(0);
       const observable = from(store);
@@ -718,4 +741,114 @@ describe('stores', () => {
       expect(cleanUpFn).toHaveBeenCalledWith(3);
     });
   });
+
+  describe('batch', () => {
+    it('should work with two writables and a derived', () => {
+      const firstName = writable('Arsène');
+      const lastName = writable('Lupin');
+      const fullName = derived([firstName, lastName], ([a, b]) => `${a} ${b}`);
+      const values: string[] = [];
+      const unsubscribe = fullName.subscribe((value) => values.push(value));
+      const expectedRes = {};
+      const res = batch(() => {
+        firstName.set('Sherlock');
+        lastName.set('Holmes');
+        return expectedRes;
+      });
+      expect(values).toEqual(['Arsène Lupin', 'Sherlock Holmes']);
+      expect(res).toBe(expectedRes);
+      unsubscribe();
+    });
+
+    it('should not call listeners multiple times', () => {
+      const store = writable(0);
+      const calls: number[] = [];
+      const unsubscribe = store.subscribe((n) => calls.push(n));
+      expect(calls).toEqual([0]);
+      batch(() => {
+        store.set(1);
+        store.set(2);
+        expect(calls).toEqual([0]);
+      });
+      expect(calls).toEqual([0, 2]);
+      unsubscribe();
+    });
+
+    it('should allow nested calls', () => {
+      const store = writable(0);
+      const calls: number[] = [];
+      const unsubscribe = store.subscribe((n) => calls.push(n));
+      expect(calls).toEqual([0]);
+      batch(() => {
+        store.set(1);
+        store.set(2);
+        batch(() => {
+          store.set(3);
+          store.set(4);
+        });
+        store.set(5);
+        batch(() => {
+          store.set(6);
+          store.set(7);
+        });
+        expect(calls).toEqual([0]);
+      });
+      expect(calls).toEqual([0, 7]);
+      unsubscribe();
+    });
+
+    it('should still call multiple times the listeners registered multiple times', () => {
+      const store = writable(0);
+      const calls: number[] = [];
+      const fn = (n: number) => calls.push(n);
+      const unsubscribe1 = store.subscribe(fn);
+      const unsubscribe2 = store.subscribe(fn);
+      expect(calls).toEqual([0, 0]);
+      batch(() => {
+        store.set(1);
+        store.set(2);
+        expect(calls).toEqual([0, 0]);
+      });
+      expect(calls).toEqual([0, 0, 2, 2]);
+      unsubscribe1();
+      unsubscribe2();
+    });
+
+    it('should not call invalidate multiple times', () => {
+      const store = writable(0);
+      let invalidateCalls = 0;
+      const unsubscribe = store.subscribe({
+        next: () => {},
+        invalidate: () => {
+          invalidateCalls++;
+        },
+      });
+      expect(invalidateCalls).toEqual(0);
+      batch(() => {
+        expect(invalidateCalls).toEqual(0);
+        store.set(1);
+        expect(invalidateCalls).toEqual(1);
+        store.set(2);
+        expect(invalidateCalls).toEqual(1);
+      });
+      expect(invalidateCalls).toEqual(1);
+      unsubscribe();
+    });
+
+    it('should not call an unregistered listener', () => {
+      const store = writable(0);
+      const calls: number[] = [];
+      const fn = (n: number) => calls.push(n);
+      const unsubscribe = store.subscribe(fn);
+      expect(calls).toEqual([0]);
+      batch(() => {
+        store.set(1);
+        store.set(2);
+        expect(calls).toEqual([0]);
+        unsubscribe();
+        store.set(3);
+      });
+      expect(calls).toEqual([0]);
+    });
+  });
 });

src/index.ts

@@ -113,14 +113,8 @@ const asReadable = <T>(store: Store<T>): Readable<T> => ({
   [symbolObservable]: returnThis,
 });
 
-const queue: [SubscriberFunction<any>, any][] = [];
-
-function processQueue() {
-  for (const [subscriberFn, value] of queue) {
-    subscriberFn(value);
-  }
-  queue.length = 0;
-}
+let willProcessQueue = false;
+const queue = new Map<SubscriberObject<any>, any>();
 
 const callUnsubscribe = (unsubscribe: Unsubscriber) =>
   typeof unsubscribe === 'function' ? unsubscribe() : unsubscribe.unsubscribe();
@@ -133,6 +127,50 @@ function notEqual(a: any, b: any): boolean {
   return true;
 }
 
+/**
+ * Batches multiple changes to stores while calling the provided function,
+ * preventing derived stores from updating until the function returns,
+ * to avoid unnecessary recomputations.
+ * If a store is updated multiple times in the provided function, listeners
+ * of that store will only be called once when the provided function returns.
+ * It is possible to have nested calls of batch, in which case only the first
+ * (outer) call has an effect, inner calls only call the provided function.
+ *
+ * @param fn - a function that can update stores. Its return value is
+ * returned by the batch function.
+ *
+ * @example
+ * Using batch in the following example prevents logging the intermediate "Sherlock Lupin" value.
+ *
+ * ```typescript
+ * const firstName = writable('Arsène');
+ * const lastName = writable('Lupin');
+ * const fullName = derived([firstName, lastName], ([a, b]) => `${a} ${b}`);
+ * fullName.subscribe((name) => console.log(name)); // logs any change to fullName
+ * batch(() => {
+ *     firstName.set('Sherlock');
+ *     lastName.set('Holmes');
+ * });
+ * ```
+ */
+export const batch = <T>(fn: () => T): T => {
+  const needsProcessQueue = !willProcessQueue;
+  if (needsProcessQueue) {
+    willProcessQueue = true;
+  }
+  try {
+    return fn();
+  } finally {
+    if (needsProcessQueue) {
+      for (const [subscriberObject, value] of queue) {
+        queue.delete(subscriberObject);
+        subscriberObject.next(value);
+      }
+      willProcessQueue = false;
+    }
+  }
+};
+
 /**
  * A utility function to get the current value from a given store.
  * It works by subscribing to a store, capturing the value (synchronously) and unsubscribing just after.
@@ -217,14 +255,15 @@ export abstract class Store<T> implements Readable<T> {
         // subscriber not yet initialized
         return;
       }
-      const needsProcessQueue = queue.length == 0;
-      for (const subscriber of this._subscribers) {
-        subscriber.invalidate();
-        queue.push([subscriber.next, value]);
-      }
-      if (needsProcessQueue) {
-        processQueue();
-      }
+      batch(() => {
+        for (const subscriber of this._subscribers) {
+          const needInvalidate = !queue.has(subscriber);
+          queue.set(subscriber, value);
+          if (needInvalidate) {
+            subscriber.invalidate();
+          }
+        }
+      });
     }
   }
 
@@ -278,6 +317,8 @@ export abstract class Store<T> implements Readable<T> {
 
     const unsubscribe = () => {
       const removed = this._subscribers.delete(subscriberObject);
+      subscriberObject.next = noop;
+      subscriberObject.invalidate = noop;
       if (removed && this._subscribers.size === 0) {
         this._stop();
       }