Add store option to unique()

src/unique.ts

@@ -1,6 +1,11 @@
 import { isAsyncIterable } from './predicate/isAsyncIterable.js';
 import { isIterable } from './predicate/isIterable.js';
 
+export type MembershipStore<T> = {
+  has(value: T): boolean;
+  add(value: T): void;
+};
+
 export type UniqueOptions<Type> = {
   /**
    * Return a custom ID to uniquely identify an item.
@@ -9,9 +14,19 @@ export type UniqueOptions<Type> = {
   /**
    * Return `true` to yield the item.
    *
-   * Use https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Iterator/filter when available.
+   * @todo Refactor to use `filter` on the (async)iterator when available.
+   * @see https://github.com/tc39/proposal-async-iterator-helpers
+   * @see https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Iterator/filter
    */
   filter?: (item: Type) => boolean;
+  /**
+   * A store to keep track of which items have been seen.
+   *
+   * If you have a large dataset, you may want to use a more memory-efficient store than a `Set`, such as a Bloom filter.
+   *
+   * @default `Set<unknown>`.
+   */
+  store?: MembershipStore<unknown>;
 };
 
 export function unique<Type>(items: AsyncIterable<Type>, options?: UniqueOptions<Type>): AsyncIterable<Type>;
@@ -27,8 +42,7 @@ export function unique<Type>(
   items: AsyncIterable<Type> | Iterable<Type>,
   options: UniqueOptions<Type> = {},
 ): AsyncIterable<Type> | Iterable<Type> {
-  const ids = new Set<unknown>();
-  const { filter, identity } = options;
+  const { filter, identity, store = new Set<unknown>() } = options;
 
   if (typeof items === 'string' || isIterable(items)) {
     return {
@@ -37,8 +51,8 @@ export function unique<Type>(
           if (!filter || filter(item) === true) {
             const id = identity?.(item) ?? item;
 
-            if (!ids.has(id)) {
-              ids.add(id);
+            if (!store.has(id)) {
+              store.add(id);
 
               yield item;
             }
@@ -55,8 +69,8 @@ export function unique<Type>(
           if (!filter || filter(item) === true) {
             const id = identity?.(item) ?? item;
 
-            if (!ids.has(id)) {
-              ids.add(id);
+            if (!store.has(id)) {
+              store.add(id);
 
               yield item;
             }

test/unique.test.ts

@@ -1,6 +1,6 @@
-import { describe, expect, it } from 'vitest';
+import { describe, expect, it, vi } from 'vitest';
 
-import { unique } from '../src/unique.js';
+import { unique, type MembershipStore } from '../src/unique.js';
 
 describe('unique()', () => {
   it('Yields unique items from an Iterable', () => {
@@ -70,4 +70,38 @@ describe('unique()', () => {
       unique(false as unknown as Iterable<boolean>);
     }).toThrow();
   });
+
+  it('Can use a custom MembershipStore via the `store` option', () => {
+    const has = vi.fn((value: unknown) => value === 2);
+    const add = vi.fn();
+
+    const store: MembershipStore<unknown> = { has, add };
+
+    expect([...unique([1, 2, 3], { store })]).toEqual([1, 3]);
+
+    expect(has).toHaveBeenCalledWith(1);
+    expect(has).toHaveBeenCalledWith(2);
+    expect(has).toHaveBeenCalledWith(3);
+    expect(add).toHaveBeenCalledWith(1);
+    expect(add).not.toHaveBeenCalledWith(2);
+    expect(add).toHaveBeenCalledWith(3);
+  });
+
+  it('Shares uniqueness state across calls when the same `store` is reused', () => {
+    const store: MembershipStore<unknown> = new Set();
+
+    expect([...unique([1, 2, 3], { store })]).toEqual([1, 2, 3]);
+    expect([...unique([2, 3, 4], { store })]).toEqual([4]);
+  });
+
+  it('Uses the `store` option with an AsyncIterable', async () => {
+    const demo = async function* (): AsyncGenerator<number> {
+      yield 1;
+      yield 2;
+      yield 2;
+      yield 3;
+    };
+
+    await expect(Array.fromAsync(unique(demo(), { store: new Set([2]) }))).resolves.toEqual([1, 3]);
+  });
 });