Maintain unique indexes incrementally for O(n) bulk writes

src/collection.ts

@@ -18,17 +18,45 @@ import type {
   UpdateChanges,
 } from './types.js';
 
-interface IndexDefinition {
-  readonly fields: string[];
-  readonly spec: IndexSpec;
-  readonly unique: boolean;
-}
-
 /** Deep copy used at every boundary so callers can never mutate stored state. */
 function clone<T>(value: T): T {
   return structuredClone(value);
 }
 
+/**
+ * A collection index.
+ *
+ * A **unique** index keeps a live `entries` map from each document's canonical
+ * key (the combination of its indexed field values) to the id that holds it.
+ * Maintaining this map incrementally on every write makes uniqueness checks
+ * O(1) per document, so inserting `n` documents is O(n) rather than O(n²).
+ *
+ * A **non-unique** index has no `entries` map: its direction is recorded but
+ * nothing is enforced and queries still scan.
+ */
+class Index {
+  /** Canonical key -> id of the holder, for unique indexes; `null` otherwise. */
+  readonly entries: Map<string, DocumentId> | null;
+
+  constructor(
+    readonly fields: string[],
+    readonly spec: IndexSpec,
+    readonly unique: boolean,
+  ) {
+    this.entries = unique ? new Map<string, DocumentId>() : null;
+  }
+
+  /** The indexed field values of a document (a missing field reads as `null`). */
+  valuesOf(doc: Document): unknown[] {
+    return this.fields.map((field) => getPath(doc, field) ?? null);
+  }
+
+  /** The canonical lookup key of a document on this index. */
+  keyOf(doc: Document): string {
+    return canonicalKey(this.valuesOf(doc));
+  }
+}
+
 /**
  * A MongoDB-like cursor over query results, returned by {@link Collection.list}.
  *
@@ -93,7 +121,7 @@ export class DocumentCursor implements AsyncIterableIterator<Document> {
  */
 export class Collection {
   private readonly documents = new Map<DocumentId, Document>();
-  private readonly indexes: IndexDefinition[] = [];
+  private readonly indexes: Index[] = [];
 
   /**
    * @param name        Collection name.
@@ -134,9 +162,24 @@ export class Collection {
     if (this.documents.has(id)) {
       throw new DuplicateKeyError(['_id'], [id]);
     }
-    this.assertUnique([...this.documents.values(), doc]);
+
+    // Check every unique index against its key map (O(1) each) before mutating
+    // anything, so a collision on a later index leaves the collection unchanged.
+    const additions: Array<{ entries: Map<string, DocumentId>; key: string }> = [];
+    for (const index of this.indexes) {
+      if (!index.entries) continue;
+      const values = index.valuesOf(doc);
+      const key = canonicalKey(values);
+      if (index.entries.has(key)) {
+        throw new DuplicateKeyError(index.fields, values);
+      }
+      additions.push({ entries: index.entries, key });
+    }
 
     this.documents.set(id, doc);
+    for (const { entries, key } of additions) {
+      entries.set(key, id);
+    }
     return doc;
   }
 
@@ -183,27 +226,45 @@ export class Collection {
     );
     if (matches.length === 0) return 0;
 
-    // Build the proposed documents first; validate before committing anything.
-    // applyChanges may throw (bad increment / path) — because nothing is written
-    // to storage until the loop completes, a failure leaves the collection intact.
-    const proposed = new Map<DocumentId, Document>();
+    // Build the updated documents up front. applyChanges may throw (bad
+    // increment / path); since nothing is committed until every check below
+    // passes, a failure leaves the collection unchanged.
+    const updated = new Map<DocumentId, Document>();
     for (const doc of matches) {
-      const updated = clone(doc);
-      applyChanges(updated, changes);
-      proposed.set(doc._id, updated);
+      const next = clone(doc);
+      applyChanges(next, changes);
+      updated.set(doc._id, next);
     }
 
-    const resulting = new Map(this.documents);
-    for (const [id, doc] of proposed) resulting.set(id, doc);
-    this.assertUnique(resulting.values());
+    this.assertUpdateKeepsUniqueness(updated);
 
-    for (const [id, doc] of proposed) this.documents.set(id, doc);
-    return proposed.size;
+    // Commit. Per unique index, vacate the updated documents' old keys, then
+    // claim their new ones — so a swap of two values commits cleanly.
+    for (const index of this.indexes) {
+      if (!index.entries) continue;
+      for (const id of updated.keys()) {
+        index.entries.delete(index.keyOf(this.documents.get(id)!));
+      }
+      for (const [id, doc] of updated) {
+        index.entries.set(index.keyOf(doc), id);
+      }
+    }
+    for (const [id, doc] of updated) {
+      this.documents.set(id, doc);
+    }
+    return updated.size;
   }
 
   /** Deletes the document with the given id. Returns whether it existed. */
   async delete(id: DocumentId): Promise<boolean> {
-    return this.documents.delete(id);
+    const doc = this.documents.get(id);
+    if (!doc) return false;
+
+    this.documents.delete(id);
+    for (const index of this.indexes) {
+      index.entries?.delete(index.keyOf(doc));
+    }
+    return true;
   }
 
   /**
@@ -233,7 +294,7 @@ export class Collection {
     }
 
     const unique = Boolean(options.unique);
-    const existing = this.indexes.find((def) => sameFields(def.fields, fields));
+    const existing = this.indexes.find((index) => sameFields(index.fields, fields));
     if (existing) {
       if (existing.unique !== unique) {
         throw new DatabaseError(
@@ -243,42 +304,59 @@ export class Collection {
       return; // idempotent
     }
 
-    const def: IndexDefinition = { fields, spec: { ...spec }, unique };
-    this.indexes.push(def);
-    try {
-      this.assertUnique(this.documents.values());
-    } catch (error) {
-      this.indexes.pop(); // roll back the half-created index
-      throw error;
+    const index = new Index(fields, { ...spec }, unique);
+    // Populate a unique index from existing data, rejecting any duplicate. The
+    // index is only registered once it has been built successfully.
+    if (index.entries) {
+      for (const doc of this.documents.values()) {
+        const values = index.valuesOf(doc);
+        const key = canonicalKey(values);
+        if (index.entries.has(key)) {
+          throw new DuplicateKeyError(index.fields, values);
+        }
+        index.entries.set(key, doc._id);
+      }
     }
+    this.indexes.push(index);
   }
 
   /** Describes the indexes defined on this collection. */
   listIndexes(): IndexDescription[] {
-    return this.indexes.map((def) => ({
-      fields: [...def.fields],
-      spec: { ...def.spec },
-      unique: def.unique,
+    return this.indexes.map((index) => ({
+      fields: [...index.fields],
+      spec: { ...index.spec },
+      unique: index.unique,
     }));
   }
 
-  /** Verifies that the given documents satisfy every unique index. */
-  private assertUnique(docs: Iterable<Document>): void {
-    const uniqueIndexes = this.indexes.filter((def) => def.unique);
-    if (uniqueIndexes.length === 0) return;
-
-    const seen = uniqueIndexes.map(() => new Set<string>());
-    for (const doc of docs) {
-      uniqueIndexes.forEach((def, i) => {
-        // A missing indexed field indexes as `null` (so a missing field and an
-        // explicit `null` collide, mirroring MongoDB).
-        const values = def.fields.map((field) => getPath(doc, field) ?? null);
+  /**
+   * Verifies that applying `updated` (id -> new document) violates no unique
+   * index, consulting each index's key map for O(matches) lookups rather than
+   * rescanning the whole collection. Throws on the first conflict.
+   */
+  private assertUpdateKeepsUniqueness(
+    updated: Map<DocumentId, Document>,
+  ): void {
+    for (const index of this.indexes) {
+      if (!index.entries) continue;
+      const claimed = new Map<string, DocumentId>();
+      for (const [id, doc] of updated) {
+        const values = index.valuesOf(doc);
         const key = canonicalKey(values);
-        if (seen[i]!.has(key)) {
-          throw new DuplicateKeyError(def.fields, values);
+
+        // Two updated documents cannot end up with the same key…
+        if (claimed.has(key)) {
+          throw new DuplicateKeyError(index.fields, values);
+        }
+        claimed.set(key, id);
+
+        // …nor can an updated document collide with one that is staying put. A
+        // holder that is itself being updated is vacating this key, so allow it.
+        const holder = index.entries.get(key);
+        if (holder !== undefined && !updated.has(holder)) {
+          throw new DuplicateKeyError(index.fields, values);
         }
-        seen[i]!.add(key);
-      });
+      }
     }
   }
 }

tests/cursor.test.ts

@@ -96,5 +96,36 @@ describe('list cursor: toArray + limit', () => {
       const rest = await cursor.toArray();
       expect(rest.map((d) => d.n)).toEqual([2, 3, 4, 5]);
     });
+
+    it('serves independent cursors from the same collection', async () => {
+      const a = people.list();
+      const b = people.list();
+      await a.next(); // advance only a
+      expect(await b.toArray()).toHaveLength(5); // b is independent
+      expect(await a.toArray()).toHaveLength(4); // a resumes after its first item
+    });
+
+    it('is unaffected by a delete after the cursor was created', async () => {
+      const cursor = people.list();
+      await people.delete('1');
+      expect(await cursor.toArray()).toHaveLength(5); // snapshot still holds it
+    });
+
+    it('yields the pre-update version of a document (snapshot)', async () => {
+      const cursor = people.list({ _id: '1' });
+      await people.update({ _id: '1' }, { n: 999 });
+      const [doc] = await cursor.toArray();
+      expect(doc?.n).toBe(1); // old value, not 999
+    });
+
+    it('supports early termination with break', async () => {
+      let seen = 0;
+      for await (const doc of people.list()) {
+        void doc;
+        seen += 1;
+        if (seen === 2) break;
+      }
+      expect(seen).toBe(2);
+    });
   });
 });

tests/database.test.ts

@@ -55,4 +55,14 @@ describe('Database', () => {
     await db.newCollection('b');
     expect(db.listCollections()).toEqual(['a', 'b']);
   });
+
+  it('drops all data when a collection is removed and re-created', async () => {
+    const db = new Database();
+    const c = await db.newCollection('c');
+    await c.insert({ _id: 'x' });
+    await db.removeCollection('c');
+    const fresh = await db.newCollection('c'); // same name is free again
+    expect(fresh.size).toBe(0);
+    expect(await fresh.get('x')).toBeNull();
+  });
 });

tests/ensure-index.test.ts

@@ -190,4 +190,54 @@ describe('ensureIndex', () => {
     );
     expect(users.size).toBe(before);
   });
+
+  it('treats compound indexes with different field order as distinct', async () => {
+    await users.ensureIndex({ a: 1, b: 1 }, { unique: true });
+    await users.ensureIndex({ b: 1, a: 1 }, { unique: true });
+    expect(users.listIndexes()).toHaveLength(2);
+  });
+
+  it('produces a stable key for out-of-domain nested values (undefined)', async () => {
+    await users.ensureIndex({ meta: 1 }, { unique: true });
+    await users.insert({ _id: 'a', meta: { x: undefined } });
+    await expect(
+      users.insert({ _id: 'b', meta: { x: undefined } }),
+    ).rejects.toBeInstanceOf(DuplicateKeyError);
+  });
+
+  it('indexes Date values by instant', async () => {
+    await users.ensureIndex({ at: 1 }, { unique: true });
+    await users.insert({ at: new Date('2024-01-01T00:00:00.000Z') });
+    await users.insert({ at: new Date('2024-02-01T00:00:00.000Z') }); // distinct
+    await expect(
+      users.insert({ at: new Date('2024-01-01T00:00:00.000Z') }),
+    ).rejects.toBeInstanceOf(DuplicateKeyError);
+  });
+
+  it('indexes array values by element and order', async () => {
+    await users.ensureIndex({ tags: 1 }, { unique: true });
+    await users.insert({ tags: ['a', 'b'] });
+    await users.insert({ tags: ['b', 'a'] }); // different order → distinct
+    await expect(users.insert({ tags: ['a', 'b'] })).rejects.toBeInstanceOf(
+      DuplicateKeyError,
+    );
+  });
+
+  it('deletes cleanly on a collection that has only a non-unique index', async () => {
+    await users.ensureIndex({ city: 1 }); // non-unique → no key map to maintain
+    await users.insert({ _id: 'x', city: 'NYC' });
+    expect(await users.delete('x')).toBe(true);
+    expect(users.size).toBe(0);
+  });
+
+  it('returns defensive copies from listIndexes', async () => {
+    await users.ensureIndex({ a: 1, b: -1 });
+    const first = users.listIndexes();
+    (first[0]!.fields as string[]).push('zzz');
+    (first[0]!.spec as Record<string, number>).a = 999;
+    // A later call is unaffected by mutating an earlier result.
+    const second = users.listIndexes();
+    expect(second[0]!.fields).toEqual(['a', 'b']);
+    expect(second[0]!.spec).toEqual({ a: 1, b: -1 });
+  });
 });

tests/filter.test.ts

@@ -64,6 +64,13 @@ describe('deepEqual', () => {
     expect(deepEqual([1, 2], { 0: 1, 1: 2 })).toBe(false);
   });
 
+  it('compares arrays of objects deeply (element key order independent)', () => {
+    expect(deepEqual([{ a: 1 }, { b: 2 }], [{ a: 1 }, { b: 2 }])).toBe(true);
+    expect(deepEqual([{ a: 1, b: 2 }], [{ b: 2, a: 1 }])).toBe(true);
+    expect(deepEqual([{ a: 1 }], [{ a: 2 }])).toBe(false);
+    expect(deepEqual([{ a: 1 }, { b: 2 }], [{ b: 2 }, { a: 1 }])).toBe(false);
+  });
+
   it('compares objects regardless of key order', () => {
     expect(deepEqual({ a: 1, b: 2 }, { b: 2, a: 1 })).toBe(true);
     expect(deepEqual({ a: 1 }, { a: 1, b: 2 })).toBe(false);