refactor: counter uses map

yarn-project/kv-store/src/interfaces/map.ts

@@ -45,19 +45,19 @@ export interface AztecMap<K extends Key, V> {
   delete(key: K): Promise<boolean>;
 
   /**
-   * Iterates over the map's key-value entries
+   * Iterates over the map's key-value entries in the key's natural order
    * @param range - The range of keys to iterate over
    */
   entries(range?: Range<K>): IterableIterator<[K, V]>;
 
   /**
-   * Iterates over the map's values
+   * Iterates over the map's values in the key's natural order
    * @param range - The range of keys to iterate over
    */
   values(range?: Range<K>): IterableIterator<V>;
 
   /**
-   * Iterates over the map's keys
+   * Iterates over the map's keys in the key's natural order
    * @param range - The range of keys to iterate over
    */
   keys(range?: Range<K>): IterableIterator<K>;
@@ -66,7 +66,7 @@ export interface AztecMap<K extends Key, V> {
 /**
  * A map backed by a persistent store that can have multiple values for a single key.
  */
-export interface AztecMultiMap<K extends string | number, V> extends AztecMap<K, V> {
+export interface AztecMultiMap<K extends Key, V> extends AztecMap<K, V> {
   /**
    * Gets all the values at the given key.
    * @param key - The key to get the values from

yarn-project/kv-store/src/lmdb/counter.test.ts

@@ -40,6 +40,13 @@ describe('LmdbAztecCounter', () => {
       expect(counter.get(key)).toEqual(0);
     });
 
+    it('throws when decrementing below zero', async () => {
+      const key = genKey();
+      await counter.update(key, 1);
+
+      await expect(counter.update(key, -2)).rejects.toThrow();
+    });
+
     it('increments values by a delta', async () => {
       const key = genKey();
       await counter.update(key, 1);

yarn-project/kv-store/src/lmdb/counter.ts

@@ -2,81 +2,56 @@ import { Key as BaseKey, Database } from 'lmdb';
 
 import { Key, Range } from '../interfaces/common.js';
 import { AztecCounter } from '../interfaces/counter.js';
-
-/** The slot where a key-value entry would be stored */
-type CountMapKey<K> = ['count_map', string, 'slot', K];
+import { LmdbAztecMap } from './map.js';
 
 /**
  * A counter implementation backed by LMDB
  */
 export class LmdbAztecCounter<K extends Key> implements AztecCounter<K> {
-  #db: Database<[K, number], CountMapKey<K>>;
+  #db: Database;
   #name: string;
-
-  #startSentinel: CountMapKey<Buffer>;
-  #endSentinel: CountMapKey<Buffer>;
+  #map: LmdbAztecMap<K, number>;
 
   constructor(db: Database<unknown, BaseKey>, name: string) {
+    this.#db = db;
     this.#name = name;
-    this.#db = db as Database<[K, number], CountMapKey<K>>;
-
-    this.#startSentinel = ['count_map', this.#name, 'slot', Buffer.from([])];
-    this.#endSentinel = ['count_map', this.#name, 'slot', Buffer.from([255])];
+    this.#map = new LmdbAztecMap(db, name);
   }
 
   set(key: K, value: number): Promise<boolean> {
-    return this.#db.put(this.#slot(key), [key, value]);
+    return this.#map.set(key, value);
   }
 
   update(key: K, delta = 1): Promise<boolean> {
     return this.#db.childTransaction(() => {
-      const slot = this.#slot(key);
-      const [_, current] = this.#db.get(slot) ?? [key, 0];
+      const current = this.#map.get(key) ?? 0;
       const next = current + delta;
 
       if (next < 0) {
         throw new Error(`Cannot update ${key} in counter ${this.#name} below zero`);
       }
 
       if (next === 0) {
-        void this.#db.remove(slot);
+        void this.#map.delete(key);
       } else {
         // store the key inside the entry because LMDB might return an internal representation
         // of the key when iterating over the database
-        void this.#db.put(slot, [key, next]);
+        void this.#map.set(key, next);
       }
 
       return true;
     });
   }
 
   get(key: K): number {
-    return (this.#db.get(this.#slot(key)) ?? [key, 0])[1];
-  }
-
-  *entries(range: Range<K> = {}): IterableIterator<[K, number]> {
-    const { start, end, reverse, limit } = range;
-    const cursor = this.#db.getRange({
-      start: start ? this.#slot(start) : reverse ? this.#endSentinel : this.#startSentinel,
-      end: end ? this.#slot(end) : reverse ? this.#startSentinel : this.#endSentinel,
-      reverse,
-      limit,
-    });
-
-    for (const {
-      value: [key, value],
-    } of cursor) {
-      yield [key, value];
-    }
+    return this.#map.get(key) ?? 0;
   }
 
-  *keys(range: Range<K> = {}): IterableIterator<K> {
-    for (const [key] of this.entries(range)) {
-      yield key;
-    }
+  entries(range: Range<K> = {}): IterableIterator<[K, number]> {
+    return this.#map.entries(range);
   }
 
-  #slot(key: K): CountMapKey<K> {
-    return ['count_map', this.#name, 'slot', key];
+  keys(range: Range<K> = {}): IterableIterator<K> {
+    return this.#map.keys(range);
   }
 }

yarn-project/kv-store/src/lmdb/map.test.ts

@@ -41,26 +41,24 @@ describe('LmdbAztecMap', () => {
     await map.set('foo', 'bar');
     await map.set('baz', 'qux');
 
-    expect([...map.entries()]).toEqual(
-      expect.arrayContaining([
-        ['foo', 'bar'],
-        ['baz', 'qux'],
-      ]),
-    );
+    expect([...map.entries()]).toEqual([
+      ['baz', 'qux'],
+      ['foo', 'bar'],
+    ]);
   });
 
   it('should be able to iterate over values', async () => {
     await map.set('foo', 'bar');
-    await map.set('baz', 'qux');
+    await map.set('baz', 'quux');
 
-    expect([...map.values()]).toEqual(expect.arrayContaining(['bar', 'qux']));
+    expect([...map.values()]).toEqual(['quux', 'bar']);
   });
 
   it('should be able to iterate over keys', async () => {
     await map.set('foo', 'bar');
     await map.set('baz', 'qux');
 
-    expect([...map.keys()]).toEqual(expect.arrayContaining(['foo', 'baz']));
+    expect([...map.keys()]).toEqual(['baz', 'foo']);
   });
 
   it('should be able to get multiple values for a single key', async () => {
@@ -69,4 +67,33 @@ describe('LmdbAztecMap', () => {
 
     expect([...map.getValues('foo')]).toEqual(['bar', 'baz']);
   });
+
+  it('supports tuple keys', async () => {
+    const map = new LmdbAztecMap<[number, string], string>(db, 'test');
+
+    await map.set([5, 'bar'], 'val');
+    await map.set([0, 'foo'], 'val');
+
+    expect([...map.keys()]).toEqual([
+      [0, 'foo'],
+      [5, 'bar'],
+    ]);
+
+    expect(map.get([5, 'bar'])).toEqual('val');
+  });
+
+  it('supports range queries', async () => {
+    await map.set('a', 'a');
+    await map.set('b', 'b');
+    await map.set('c', 'c');
+    await map.set('d', 'd');
+
+    expect([...map.keys({ start: 'b', end: 'c' })]).toEqual(['b']);
+    expect([...map.keys({ start: 'b' })]).toEqual(['b', 'c', 'd']);
+    expect([...map.keys({ end: 'c' })]).toEqual(['a', 'b']);
+    expect([...map.keys({ start: 'b', end: 'c', reverse: true })]).toEqual(['c']);
+    expect([...map.keys({ start: 'b', limit: 1 })]).toEqual(['b']);
+    expect([...map.keys({ start: 'b', reverse: true })]).toEqual(['d', 'c']);
+    expect([...map.keys({ end: 'b', reverse: true })]).toEqual(['b', 'a']);
+  });
 });

yarn-project/kv-store/src/lmdb/map.ts

@@ -1,24 +1,24 @@
-import { Database, Key } from 'lmdb';
+import { Database, RangeOptions } from 'lmdb';
 
-import { Range } from '../interfaces/common.js';
+import { Key, Range } from '../interfaces/common.js';
 import { AztecMultiMap } from '../interfaces/map.js';
 
 /** The slot where a key-value entry would be stored */
-type MapKeyValueSlot<K extends string | number | Buffer> = ['map', string, 'slot', K];
+type MapValueSlot<K extends Key | Buffer> = ['map', string, 'slot', K];
 
 /**
  * A map backed by LMDB.
  */
-export class LmdbAztecMap<K extends string | number, V> implements AztecMultiMap<K, V> {
-  protected db: Database<V, MapKeyValueSlot<K>>;
+export class LmdbAztecMap<K extends Key, V> implements AztecMultiMap<K, V> {
+  protected db: Database<[K, V], MapValueSlot<K>>;
   protected name: string;
 
-  #startSentinel: MapKeyValueSlot<Buffer>;
-  #endSentinel: MapKeyValueSlot<Buffer>;
+  #startSentinel: MapValueSlot<Buffer>;
+  #endSentinel: MapValueSlot<Buffer>;
 
-  constructor(rootDb: Database<unknown, Key>, mapName: string) {
+  constructor(rootDb: Database, mapName: string) {
     this.name = mapName;
-    this.db = rootDb as Database<V, MapKeyValueSlot<K>>;
+    this.db = rootDb as Database<[K, V], MapValueSlot<K>>;
 
     // sentinels are used to define the start and end of the map
     // with LMDB's key encoding, no _primitive value_ can be "less than" an empty buffer or greater than Byte 255
@@ -32,13 +32,13 @@ export class LmdbAztecMap<K extends string | number, V> implements AztecMultiMap
   }
 
   get(key: K): V | undefined {
-    return this.db.get(this.#slot(key)) as V | undefined;
+    return this.db.get(this.#slot(key))?.[1];
   }
 
   *getValues(key: K): IterableIterator<V> {
     const values = this.db.getValues(this.#slot(key));
     for (const value of values) {
-      yield value;
+      yield value?.[1];
     }
   }
 
@@ -47,14 +47,14 @@ export class LmdbAztecMap<K extends string | number, V> implements AztecMultiMap
   }
 
   set(key: K, val: V): Promise<boolean> {
-    return this.db.put(this.#slot(key), val);
+    return this.db.put(this.#slot(key), [key, val]);
   }
 
   swap(key: K, fn: (val: V | undefined) => V): Promise<boolean> {
     return this.db.childTransaction(() => {
       const slot = this.#slot(key);
-      const val = this.db.get(slot);
-      void this.db.put(slot, fn(val));
+      const entry = this.db.get(slot);
+      void this.db.put(slot, [key, fn(entry?.[1])]);
 
       return true;
     });
@@ -63,7 +63,7 @@ export class LmdbAztecMap<K extends string | number, V> implements AztecMultiMap
   setIfNotExists(key: K, val: V): Promise<boolean> {
     const slot = this.#slot(key);
     return this.db.ifNoExists(slot, () => {
-      void this.db.put(slot, val);
+      void this.db.put(slot, [key, val]);
     });
   }
 
@@ -72,27 +72,42 @@ export class LmdbAztecMap<K extends string | number, V> implements AztecMultiMap
   }
 
   async deleteValue(key: K, val: V): Promise<void> {
-    await this.db.remove(this.#slot(key), val);
+    await this.db.remove(this.#slot(key), [key, val]);
   }
 
   *entries(range: Range<K> = {}): IterableIterator<[K, V]> {
-    const { start, end, reverse = false, limit } = range;
+    const { reverse = false, limit } = range;
     // LMDB has a quirk where it expects start > end when reverse=true
     // in that case, we need to swap the start and end sentinels
-    const iterator = this.db.getRange({
-      start: start ? this.#slot(start) : reverse ? this.#endSentinel : this.#startSentinel,
-      end: end ? this.#slot(end) : reverse ? this.#startSentinel : this.#endSentinel,
+    const start = reverse
+      ? range.end
+        ? this.#slot(range.end)
+        : this.#endSentinel
+      : range.start
+      ? this.#slot(range.start)
+      : this.#startSentinel;
+
+    const end = reverse
+      ? range.start
+        ? this.#slot(range.start)
+        : this.#startSentinel
+      : range.end
+      ? this.#slot(range.end)
+      : this.#endSentinel;
+
+    const lmdbRange: RangeOptions = {
+      start,
+      end,
       reverse,
       limit,
-    });
+    };
 
-    for (const { key, value } of iterator) {
-      if (key[0] !== 'map' || key[1] !== this.name) {
-        break;
-      }
+    const iterator = this.db.getRange(lmdbRange);
 
-      const originalKey = key[3];
-      yield [originalKey, value];
+    for (const {
+      value: [key, value],
+    } of iterator) {
+      yield [key, value];
     }
   }
 
@@ -108,7 +123,7 @@ export class LmdbAztecMap<K extends string | number, V> implements AztecMultiMap
     }
   }
 
-  #slot(key: K): MapKeyValueSlot<K> {
+  #slot(key: K): MapValueSlot<K> {
     return ['map', this.name, 'slot', key];
   }
 }