From eb35c201f33b36446678bf311ab68688a5475439 Mon Sep 17 00:00:00 2001 From: Charlie Lye Date: Sat, 30 Mar 2024 10:54:11 +0000 Subject: [PATCH] feat: in memory kv store. --- yarn-project/kv-store/package.json | 1 + yarn-project/kv-store/src/lmdb/array.test.ts | 81 +-------- yarn-project/kv-store/src/lmdb/map.test.ts | 89 +--------- yarn-project/kv-store/src/mem/array.test.ts | 7 + yarn-project/kv-store/src/mem/array.ts | 70 ++++++++ yarn-project/kv-store/src/mem/counter.test.ts | 7 + yarn-project/kv-store/src/mem/counter.ts | 50 ++++++ yarn-project/kv-store/src/mem/index.ts | 1 + yarn-project/kv-store/src/mem/map.test.ts | 7 + yarn-project/kv-store/src/mem/map.ts | 163 ++++++++++++++++++ yarn-project/kv-store/src/mem/mem_db.ts | 53 ++++++ .../kv-store/src/mem/singleton.test.ts | 7 + yarn-project/kv-store/src/mem/singleton.ts | 27 +++ yarn-project/kv-store/src/mem/store.test.ts | 66 +++++++ yarn-project/kv-store/src/mem/store.ts | 54 ++++++ .../kv-store/src/tests/aztec_array_tests.ts | 81 +++++++++ .../kv-store/src/tests/aztec_counter_tests.ts | 120 +++++++++++++ .../kv-store/src/tests/aztec_map_tests.ts | 108 ++++++++++++ .../src/tests/aztec_singleton_tests.ts | 28 +++ yarn-project/kv-store/src/utils.ts | 5 +- 20 files changed, 857 insertions(+), 168 deletions(-) create mode 100644 yarn-project/kv-store/src/mem/array.test.ts create mode 100644 yarn-project/kv-store/src/mem/array.ts create mode 100644 yarn-project/kv-store/src/mem/counter.test.ts create mode 100644 yarn-project/kv-store/src/mem/counter.ts create mode 100644 yarn-project/kv-store/src/mem/index.ts create mode 100644 yarn-project/kv-store/src/mem/map.test.ts create mode 100644 yarn-project/kv-store/src/mem/map.ts create mode 100644 yarn-project/kv-store/src/mem/mem_db.ts create mode 100644 yarn-project/kv-store/src/mem/singleton.test.ts create mode 100644 yarn-project/kv-store/src/mem/singleton.ts create mode 100644 yarn-project/kv-store/src/mem/store.test.ts create mode 100644 yarn-project/kv-store/src/mem/store.ts create mode 100644 yarn-project/kv-store/src/tests/aztec_array_tests.ts create mode 100644 yarn-project/kv-store/src/tests/aztec_counter_tests.ts create mode 100644 yarn-project/kv-store/src/tests/aztec_map_tests.ts create mode 100644 yarn-project/kv-store/src/tests/aztec_singleton_tests.ts diff --git a/yarn-project/kv-store/package.json b/yarn-project/kv-store/package.json index 55808c673125..32aaaca24cf5 100644 --- a/yarn-project/kv-store/package.json +++ b/yarn-project/kv-store/package.json @@ -5,6 +5,7 @@ "exports": { ".": "./dest/interfaces/index.js", "./lmdb": "./dest/lmdb/index.js", + "./mem": "./dest/mem/index.js", "./utils": "./dest/utils.js" }, "scripts": { diff --git a/yarn-project/kv-store/src/lmdb/array.test.ts b/yarn-project/kv-store/src/lmdb/array.test.ts index 3058302e87f0..3be13b108849 100644 --- a/yarn-project/kv-store/src/lmdb/array.test.ts +++ b/yarn-project/kv-store/src/lmdb/array.test.ts @@ -1,91 +1,14 @@ import { Database, open } from 'lmdb'; +import { addArrayTests } from '../tests/aztec_array_tests.js'; import { LmdbAztecArray } from './array.js'; describe('LmdbAztecArray', () => { let db: Database; - let arr: LmdbAztecArray; beforeEach(() => { db = open({} as any); - arr = new LmdbAztecArray(db, 'test'); }); - it('should be able to push and pop values', async () => { - await arr.push(1); - await arr.push(2); - await arr.push(3); - - expect(arr.length).toEqual(3); - expect(await arr.pop()).toEqual(3); - expect(await arr.pop()).toEqual(2); - expect(await arr.pop()).toEqual(1); - expect(await arr.pop()).toEqual(undefined); - }); - - it('should be able to get values by index', async () => { - await arr.push(1); - await arr.push(2); - await arr.push(3); - - expect(arr.at(0)).toEqual(1); - expect(arr.at(1)).toEqual(2); - expect(arr.at(2)).toEqual(3); - expect(arr.at(3)).toEqual(undefined); - expect(arr.at(-1)).toEqual(3); - expect(arr.at(-2)).toEqual(2); - expect(arr.at(-3)).toEqual(1); - expect(arr.at(-4)).toEqual(undefined); - }); - - it('should be able to set values by index', async () => { - await arr.push(1); - await arr.push(2); - await arr.push(3); - - expect(await arr.setAt(0, 4)).toEqual(true); - expect(await arr.setAt(1, 5)).toEqual(true); - expect(await arr.setAt(2, 6)).toEqual(true); - - expect(await arr.setAt(3, 7)).toEqual(false); - - expect(arr.at(0)).toEqual(4); - expect(arr.at(1)).toEqual(5); - expect(arr.at(2)).toEqual(6); - expect(arr.at(3)).toEqual(undefined); - - expect(await arr.setAt(-1, 8)).toEqual(true); - expect(await arr.setAt(-2, 9)).toEqual(true); - expect(await arr.setAt(-3, 10)).toEqual(true); - - expect(await arr.setAt(-4, 11)).toEqual(false); - - expect(arr.at(-1)).toEqual(8); - expect(arr.at(-2)).toEqual(9); - expect(arr.at(-3)).toEqual(10); - expect(arr.at(-4)).toEqual(undefined); - }); - - it('should be able to iterate over values', async () => { - await arr.push(1); - await arr.push(2); - await arr.push(3); - - expect([...arr.values()]).toEqual([1, 2, 3]); - expect([...arr.entries()]).toEqual([ - [0, 1], - [1, 2], - [2, 3], - ]); - }); - - it('should be able to restore state', async () => { - await arr.push(1); - await arr.push(2); - await arr.push(3); - - const arr2 = new LmdbAztecArray(db, 'test'); - expect(arr2.length).toEqual(3); - expect([...arr2.values()]).toEqual([...arr.values()]); - }); + addArrayTests(() => new LmdbAztecArray(db, 'test')); }); diff --git a/yarn-project/kv-store/src/lmdb/map.test.ts b/yarn-project/kv-store/src/lmdb/map.test.ts index 007b4c4eb8fa..0261eabc5bf4 100644 --- a/yarn-project/kv-store/src/lmdb/map.test.ts +++ b/yarn-project/kv-store/src/lmdb/map.test.ts @@ -1,99 +1,14 @@ import { Database, open } from 'lmdb'; +import { addMapTests } from '../tests/aztec_map_tests.js'; import { LmdbAztecMap } from './map.js'; describe('LmdbAztecMap', () => { let db: Database; - let map: LmdbAztecMap; beforeEach(() => { db = open({ dupSort: true } as any); - map = new LmdbAztecMap(db, 'test'); }); - it('should be able to set and get values', async () => { - await map.set('foo', 'bar'); - await map.set('baz', 'qux'); - - expect(map.get('foo')).toEqual('bar'); - expect(map.get('baz')).toEqual('qux'); - expect(map.get('quux')).toEqual(undefined); - }); - - it('should be able to set values if they do not exist', async () => { - expect(await map.setIfNotExists('foo', 'bar')).toEqual(true); - expect(await map.setIfNotExists('foo', 'baz')).toEqual(false); - - expect(map.get('foo')).toEqual('bar'); - }); - - it('should be able to delete values', async () => { - await map.set('foo', 'bar'); - await map.set('baz', 'qux'); - - expect(await map.delete('foo')).toEqual(true); - - expect(map.get('foo')).toEqual(undefined); - expect(map.get('baz')).toEqual('qux'); - }); - - it('should be able to iterate over entries', async () => { - await map.set('foo', 'bar'); - await map.set('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', 'quux'); - - 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(['baz', 'foo']); - }); - - it('should be able to get multiple values for a single key', async () => { - await map.set('foo', 'bar'); - await map.set('foo', 'baz'); - - 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']); - }); + addMapTests(() => new LmdbAztecMap(db, 'test')); }); diff --git a/yarn-project/kv-store/src/mem/array.test.ts b/yarn-project/kv-store/src/mem/array.test.ts new file mode 100644 index 000000000000..e01dd4c57119 --- /dev/null +++ b/yarn-project/kv-store/src/mem/array.test.ts @@ -0,0 +1,7 @@ +import { addArrayTests } from '../tests/aztec_array_tests.js'; +import { MemAztecArray } from './array.js'; +import { MemDb } from './mem_db.js'; + +describe('MemAztecArray', () => { + addArrayTests(() => new MemAztecArray('test', new MemDb())); +}); diff --git a/yarn-project/kv-store/src/mem/array.ts b/yarn-project/kv-store/src/mem/array.ts new file mode 100644 index 000000000000..3350fda7f51e --- /dev/null +++ b/yarn-project/kv-store/src/mem/array.ts @@ -0,0 +1,70 @@ +import { type AztecArray } from '../interfaces/array.js'; +import type { MemDb } from './mem_db.js'; + +/** + * An persistent array backed by mem. + */ +export class MemAztecArray implements AztecArray { + private slot: string; + + constructor(private name: string, private db: MemDb) { + this.slot = JSON.stringify(['array', this.name]); + } + + get length(): number { + return this.db.get(this.slot)?.length || 0; + } + + push(...vals: T[]): Promise { + const arr = this.db.get(this.slot); + if (arr) { + this.db.set(this.slot, [...arr, ...vals]); + } else { + this.db.set(this.slot, [...vals]); + } + return Promise.resolve(this.length); + } + + pop(): Promise { + const arr = [...this.db.get(this.slot)]; + const result = arr.pop(); + this.db.set(this.slot, arr); + return Promise.resolve(result); + } + + at(index: number): T | undefined { + const arr = this.db.get(this.slot) || []; + if (index < 0) { + return arr[arr.length + index]; + } else { + return arr[index]; + } + } + + setAt(index: number, val: T): Promise { + if (index < 0) { + index = this.length + index; + } + + if (index < 0 || index >= this.length) { + return Promise.resolve(false); + } + + const arr = [...this.db.get(this.slot)]; + arr[index] = val; + this.db.set(this.slot, arr); + return Promise.resolve(true); + } + + entries(): IterableIterator<[number, T]> { + return this.db.get(this.slot)?.entries() || []; + } + + values(): IterableIterator { + return this.db.get(this.slot)?.values() || []; + } + + [Symbol.iterator](): IterableIterator { + return this.values(); + } +} diff --git a/yarn-project/kv-store/src/mem/counter.test.ts b/yarn-project/kv-store/src/mem/counter.test.ts new file mode 100644 index 000000000000..17df911b542a --- /dev/null +++ b/yarn-project/kv-store/src/mem/counter.test.ts @@ -0,0 +1,7 @@ +import { addCounterTests } from '../tests/aztec_counter_tests.js'; +import { MemAztecCounter } from './counter.js'; +import { MemDb } from './mem_db.js'; + +describe('MemAztecCounter', () => { + addCounterTests(() => new MemAztecCounter('test', new MemDb())); +}); diff --git a/yarn-project/kv-store/src/mem/counter.ts b/yarn-project/kv-store/src/mem/counter.ts new file mode 100644 index 000000000000..13a71b65e9cc --- /dev/null +++ b/yarn-project/kv-store/src/mem/counter.ts @@ -0,0 +1,50 @@ +import type { Key, Range } from '../interfaces/common.js'; +import type { AztecCounter } from '../interfaces/counter.js'; +import { MemAztecMap } from './map.js'; +import type { MemDb } from './mem_db.js'; + +export class MemAztecCounter implements AztecCounter { + private map: MemAztecMap; + + constructor(name: string, db: MemDb) { + this.map = new MemAztecMap(name, db); + } + + async set(key: Key, value: number): Promise { + if (value) { + return this.map.set(key, value); + } else { + await this.map.delete(key); + return true; + } + } + + async update(key: Key, delta = 1): Promise { + const current = this.map.get(key) ?? 0; + const next = current + delta; + + if (next < 0) { + throw new Error(`Cannot update ${key} in counter below zero`); + } + + await this.map.delete(key); + + if (next > 0) { + await this.map.set(key, next); + } + + return true; + } + + get(key: Key): number { + return this.map.get(key) ?? 0; + } + + entries(range: Range = {}): IterableIterator<[Key, number]> { + return this.map.entries(range); + } + + keys(range: Range = {}): IterableIterator { + return this.map.keys(range); + } +} diff --git a/yarn-project/kv-store/src/mem/index.ts b/yarn-project/kv-store/src/mem/index.ts new file mode 100644 index 000000000000..f5aee46d6d71 --- /dev/null +++ b/yarn-project/kv-store/src/mem/index.ts @@ -0,0 +1 @@ +export { AztecMemStore } from './store.js'; diff --git a/yarn-project/kv-store/src/mem/map.test.ts b/yarn-project/kv-store/src/mem/map.test.ts new file mode 100644 index 000000000000..c9c437db8d48 --- /dev/null +++ b/yarn-project/kv-store/src/mem/map.test.ts @@ -0,0 +1,7 @@ +import { addMapTests } from '../tests/aztec_map_tests.js'; +import { MemAztecMap } from './map.js'; +import { MemDb } from './mem_db.js'; + +describe('MemAztecMap', () => { + addMapTests(() => new MemAztecMap('test', new MemDb())); +}); diff --git a/yarn-project/kv-store/src/mem/map.ts b/yarn-project/kv-store/src/mem/map.ts new file mode 100644 index 000000000000..c4f7fa916757 --- /dev/null +++ b/yarn-project/kv-store/src/mem/map.ts @@ -0,0 +1,163 @@ +import type { Key, Range } from '../interfaces/common.js'; +import type { AztecMultiMap } from '../interfaces/map.js'; +import type { MemDb } from './mem_db.js'; + +// Comparator function for keys already parsed from JSON +function compareKeys(a: Key, b: Key) { + // Handle array (tuple) comparison + if (Array.isArray(a) && Array.isArray(b)) { + for (let i = 0; i < Math.min(a.length, b.length); i++) { + if (a[i] < b[i]) { + return -1; + } + if (a[i] > b[i]) { + return 1; + } + } + return a.length - b.length; + } + + // Fallback to normal comparison for non-array keys + if (a < b) { + return -1; + } + if (a > b) { + return 1; + } + return 0; +} + +/** + * A map backed by mem. + */ +export class MemAztecMap implements AztecMultiMap { + constructor(private name: string, private db: MemDb) {} + + close(): Promise { + return Promise.resolve(); + } + + get(key: Key): V | undefined { + const r = this.db.get(this.slot(key)); + return r ? r[r.length - 1] : undefined; + } + + getValues(key: Key): IterableIterator { + const r = this.db.get(this.slot(key)); + return r ? r.values() : new Array().values(); + } + + has(key: Key): boolean { + const r = this.db.get(this.slot(key)); + return r ? r.length > 0 : false; + } + + set(key: Key, val: V): Promise { + const r = this.db.get(this.slot(key)); + if (r) { + this.db.set(this.slot(key), [...r, val]); + } else { + this.db.set(this.slot(key), [val]); + } + return Promise.resolve(true); + } + + swap(key: Key, fn: (val: V | undefined) => V): Promise { + const entry = this.get(key); + const newValue = fn(entry); + this.db.set(this.slot(key), [newValue]); + return Promise.resolve(true); + } + + async setIfNotExists(key: Key, val: V): Promise { + const r = this.get(key); + if (!r) { + await this.set(key, val); + return true; + } + return false; + } + + delete(key: Key): Promise { + const r = this.db.get(this.slot(key)); + if (r?.length) { + this.db.set(this.slot(key), []); + return Promise.resolve(true); + } + return Promise.resolve(false); + } + + deleteValue(key: Key, val: V): Promise { + const r = this.db.get(this.slot(key)); + if (r) { + const i = r.indexOf(val); + if (i != -1) { + this.db.set(this.slot(key), [...r.slice(0, i), ...r.slice(i + 1)]); + } + } + return Promise.resolve(); + } + + *entries(range: Range = {}): IterableIterator<[Key, V]> { + let { limit } = range; + const { start, end, reverse = false } = range; + + // TODO: Horrifically inefficient as backing db is not an ordered map. + // Make it so. + const keys = this.db + .keys() + .map(key => JSON.parse(key)) + .filter(key => key[0] == 'map' && key[1] == this.name) + .map(key => key[2]) + .sort(compareKeys); + + if (reverse) { + keys.reverse(); + } + + for (const key of keys) { + if (reverse) { + if (end !== undefined && compareKeys(key, end) > 0) { + continue; + } + if (start !== undefined && compareKeys(key, start) <= 0) { + break; + } + } else { + if (start !== undefined && compareKeys(key, start) < 0) { + continue; + } + if (end !== undefined && compareKeys(key, end) >= 0) { + break; + } + } + + const values = this.db.get(this.slot(key)); + if (!values) { + return; + } + for (const value of values) { + yield [key, value]; + if (limit && --limit <= 0) { + return; + } + } + } + } + + *values(range: Range = {}): IterableIterator { + for (const [_, value] of this.entries(range)) { + yield value; + } + } + + *keys(range: Range = {}): IterableIterator { + for (const [key, _] of this.entries(range)) { + yield key; + } + } + + private slot(key: Key): string { + return JSON.stringify(['map', this.name, key]); + } +} diff --git a/yarn-project/kv-store/src/mem/mem_db.ts b/yarn-project/kv-store/src/mem/mem_db.ts new file mode 100644 index 000000000000..aca1567fd21b --- /dev/null +++ b/yarn-project/kv-store/src/mem/mem_db.ts @@ -0,0 +1,53 @@ +export class MemDb { + private data: { [key: string]: any } = {}; + private tx: { [key: string]: any } | undefined; + + set(key: string, value: any) { + if (this.tx) { + this.tx[key] = value; + } else { + this.data[key] = value; + } + } + + get(key: string) { + if (this.tx && this.tx[key] !== undefined) { + return this.tx[key]; + } + return this.data[key]; + } + + del(key: string) { + if (this.tx) { + this.tx[key] = undefined; + return; + } + delete this.data[key]; + } + + keys() { + return Array.from(new Set([...Object.keys(this.data), ...Object.keys(this.tx || {})])); + } + + commit() { + if (!this.tx) { + throw new Error('tx not in progress.'); + } + Object.assign(this.data, this.tx); + this.tx = undefined; + } + + rollback() { + if (!this.tx) { + throw new Error('tx not in progress.'); + } + this.tx = undefined; + } + + startTx() { + if (this.tx) { + throw new Error('MemDb can only handle 1 tx at a time.'); + } + this.tx = {}; + } +} diff --git a/yarn-project/kv-store/src/mem/singleton.test.ts b/yarn-project/kv-store/src/mem/singleton.test.ts new file mode 100644 index 000000000000..9ea8533bde70 --- /dev/null +++ b/yarn-project/kv-store/src/mem/singleton.test.ts @@ -0,0 +1,7 @@ +import { addSingletonTests } from '../tests/aztec_singleton_tests.js'; +import { MemDb } from './mem_db.js'; +import { MemAztecSingleton } from './singleton.js'; + +describe('MemAztecSingleton', () => { + addSingletonTests(() => new MemAztecSingleton('test', new MemDb())); +}); diff --git a/yarn-project/kv-store/src/mem/singleton.ts b/yarn-project/kv-store/src/mem/singleton.ts new file mode 100644 index 000000000000..147fc94ac4db --- /dev/null +++ b/yarn-project/kv-store/src/mem/singleton.ts @@ -0,0 +1,27 @@ +import type { AztecSingleton } from '../interfaces/singleton.js'; +import type { MemDb } from './mem_db.js'; + +/** + * Stores a single value in mem. + */ +export class MemAztecSingleton implements AztecSingleton { + private slot: string; + + constructor(private name: string, private db: MemDb) { + this.slot = JSON.stringify(['array', this.name]); + } + + get(): T | undefined { + return this.db.get(this.slot); + } + + set(val: T): Promise { + this.db.set(this.slot, val); + return Promise.resolve(true); + } + + delete(): Promise { + this.db.del(this.slot); + return Promise.resolve(true); + } +} diff --git a/yarn-project/kv-store/src/mem/store.test.ts b/yarn-project/kv-store/src/mem/store.test.ts new file mode 100644 index 000000000000..7ffd3b87a8b1 --- /dev/null +++ b/yarn-project/kv-store/src/mem/store.test.ts @@ -0,0 +1,66 @@ +import type { AztecArray, AztecCounter, AztecMap, AztecSingleton } from '../interfaces/index.js'; +import { AztecMemStore } from './store.js'; + +describe('AztecMemStore', () => { + let store: AztecMemStore; + let array: AztecArray; + let multimap: AztecMap; + let counter: AztecCounter; + let singleton: AztecSingleton; + + beforeEach(async () => { + store = new AztecMemStore(); + + array = store.openArray('test-array'); + multimap = store.openMultiMap('test-multimap'); + counter = store.openCounter('test-counter'); + singleton = store.openSingleton('test-singleton'); + + await array.push(1, 2, 3); + await multimap.set('key-1', 1); + await multimap.set('key-2', 2); + await counter.set('counter-1', 3); + await singleton.set(4); + }); + + it('check initial state', () => { + expect(array.at(2)).toBe(3); + expect(multimap.get('key-2')).toBe(2); + expect(counter.get('counter-1')).toBe(3); + expect(singleton.get()).toBe(4); + }); + + it('state should update with successful tx', async () => { + await store.transaction(() => { + void array.setAt(2, 10); + void multimap.set('key-2', 20); + void counter.set('counter-1', 30); + void singleton.set(40); + }); + void multimap.set('key-2', 20); + + expect(array.at(2)).toBe(10); + expect(multimap.get('key-2')).toBe(20); + expect(counter.get('counter-1')).toBe(30); + expect(singleton.get()).toBe(40); + }); + + it('state should rollback with unsuccessful tx', async () => { + try { + await store.transaction(() => { + void array.setAt(2, 10); + void multimap.set('key-2', 20); + void counter.set('counter-1', 30); + void singleton.set(40); + throw new Error(); + }); + } catch (err) { + // swallow + } + + expect(array.at(2)).toBe(3); + expect(multimap.get('key-2')).toBe(2); + expect(counter.get('counter-1')).toBe(3); + expect(singleton.get()).toBe(4); + }); +}); diff --git a/yarn-project/kv-store/src/mem/store.ts b/yarn-project/kv-store/src/mem/store.ts new file mode 100644 index 000000000000..c35f47a5718d --- /dev/null +++ b/yarn-project/kv-store/src/mem/store.ts @@ -0,0 +1,54 @@ +import type { AztecArray } from '../interfaces/array.js'; +import type { AztecCounter } from '../interfaces/counter.js'; +import type { AztecMap, AztecMultiMap } from '../interfaces/map.js'; +import type { AztecSingleton } from '../interfaces/singleton.js'; +import type { AztecKVStore } from '../interfaces/store.js'; +import { MemAztecArray } from './array.js'; +import { MemAztecCounter } from './counter.js'; +import { MemAztecMap } from './map.js'; +import { MemDb } from './mem_db.js'; +import { MemAztecSingleton } from './singleton.js'; + +/** + * A key-value store backed by mem. + */ +export class AztecMemStore implements AztecKVStore { + private data = new MemDb(); + + openMap(name: string): AztecMap { + return new MemAztecMap(name, this.data) as any; + } + + openMultiMap(name: string): AztecMultiMap { + return new MemAztecMap(name, this.data) as any; + } + + openCounter>(name: string): AztecCounter { + return new MemAztecCounter(name, this.data) as any; + } + + openArray(name: string): AztecArray { + return new MemAztecArray(name, this.data) as any; + } + + openSingleton(name: string): AztecSingleton { + return new MemAztecSingleton(name, this.data) as any; + } + + transaction(callback: () => T): Promise { + this.data.startTx(); + try { + const result = callback(); + this.data.commit(); + return Promise.resolve(result); + } catch (err) { + this.data.rollback(); + throw err; + } + } + + clear() { + this.data = new MemDb(); + return Promise.resolve(); + } +} diff --git a/yarn-project/kv-store/src/tests/aztec_array_tests.ts b/yarn-project/kv-store/src/tests/aztec_array_tests.ts new file mode 100644 index 000000000000..19fa70b1c3d0 --- /dev/null +++ b/yarn-project/kv-store/src/tests/aztec_array_tests.ts @@ -0,0 +1,81 @@ +import { beforeEach, describe, expect, it } from '@jest/globals'; + +import { type AztecArray } from '../interfaces/index.js'; + +export function addArrayTests(getArray: () => AztecArray) { + let arr: AztecArray; + + describe('AztecArray', () => { + beforeEach(() => { + arr = getArray(); + }); + + it('should be able to push and pop values', async () => { + await arr.push(1); + await arr.push(2); + await arr.push(3); + + expect(arr.length).toEqual(3); + expect(await arr.pop()).toEqual(3); + expect(await arr.pop()).toEqual(2); + expect(await arr.pop()).toEqual(1); + expect(await arr.pop()).toEqual(undefined); + }); + + it('should be able to get values by index', async () => { + await arr.push(1); + await arr.push(2); + await arr.push(3); + + expect(arr.at(0)).toEqual(1); + expect(arr.at(1)).toEqual(2); + expect(arr.at(2)).toEqual(3); + expect(arr.at(3)).toEqual(undefined); + expect(arr.at(-1)).toEqual(3); + expect(arr.at(-2)).toEqual(2); + expect(arr.at(-3)).toEqual(1); + expect(arr.at(-4)).toEqual(undefined); + }); + + it('should be able to set values by index', async () => { + await arr.push(1); + await arr.push(2); + await arr.push(3); + + expect(await arr.setAt(0, 4)).toEqual(true); + expect(await arr.setAt(1, 5)).toEqual(true); + expect(await arr.setAt(2, 6)).toEqual(true); + + expect(await arr.setAt(3, 7)).toEqual(false); + + expect(arr.at(0)).toEqual(4); + expect(arr.at(1)).toEqual(5); + expect(arr.at(2)).toEqual(6); + expect(arr.at(3)).toEqual(undefined); + + expect(await arr.setAt(-1, 8)).toEqual(true); + expect(await arr.setAt(-2, 9)).toEqual(true); + expect(await arr.setAt(-3, 10)).toEqual(true); + + expect(await arr.setAt(-4, 11)).toEqual(false); + + expect(arr.at(-1)).toEqual(8); + expect(arr.at(-2)).toEqual(9); + expect(arr.at(-3)).toEqual(10); + expect(arr.at(-4)).toEqual(undefined); + }); + + it('should be able to iterate over values', async () => { + await arr.push(1); + await arr.push(2); + await arr.push(3); + + expect([...arr.values()]).toEqual([1, 2, 3]); + expect([...arr.entries()]).toEqual([ + [0, 1], + [1, 2], + [2, 3], + ]); + }); + }); +} diff --git a/yarn-project/kv-store/src/tests/aztec_counter_tests.ts b/yarn-project/kv-store/src/tests/aztec_counter_tests.ts new file mode 100644 index 000000000000..0a5a1b3aa2ad --- /dev/null +++ b/yarn-project/kv-store/src/tests/aztec_counter_tests.ts @@ -0,0 +1,120 @@ +import { randomBytes } from '@aztec/foundation/crypto'; + +import { beforeEach, describe, expect, it } from '@jest/globals'; + +import { type AztecCounter } from '../interfaces/index.js'; + +export function addCounterTests(get: () => AztecCounter) { + describe('AztecCounter', () => { + describe.each([ + ['floating point number', () => Math.random()], + ['integers', () => (Math.random() * 1000) | 0], + ['strings', () => randomBytes(8).toString('hex')], + ['strings', () => [Math.random(), randomBytes(8).toString('hex')]], + ])('counts occurrences of %s values', (_, genKey) => { + let counter: AztecCounter; + + beforeEach(() => { + counter = get(); + }); + + it('returns 0 for unknown keys', () => { + expect(counter.get(genKey())).toEqual(0); + }); + + it('increments values', async () => { + const key = genKey(); + await counter.update(key, 1); + + expect(counter.get(key)).toEqual(1); + }); + + it('decrements values', async () => { + const key = genKey(); + await counter.update(key, 1); + await counter.update(key, -1); + + 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); + await counter.update(key, 2); + + expect(counter.get(key)).toEqual(3); + }); + + it('resets the counter', async () => { + const key = genKey(); + await counter.update(key, 1); + await counter.update(key, 2); + await counter.set(key, 0); + + expect(counter.get(key)).toEqual(0); + }); + + it('iterates over entries', async () => { + const key = genKey(); + await counter.update(key, 1); + await counter.update(key, 2); + + expect([...counter.entries({})]).toEqual([[key, 3]]); + }); + }); + + it.each([ + [ + [ + ['c', 2342], + ['a', 8], + ['b', 1], + ], + [ + ['a', 8], + ['b', 1], + ['c', 2342], + ], + ], + [ + [ + [10, 2], + [18, 1], + [1, 2], + ], + [ + [1, 2], + [10, 2], + [18, 1], + ], + ], + [ + [ + [[10, 'a'], 1], + [[10, 'c'], 2], + [[11, 'b'], 1], + [[9, 'f'], 1], + [[10, 'b'], 1], + ], + [ + [[9, 'f'], 1], + [[10, 'a'], 1], + [[10, 'b'], 1], + [[10, 'c'], 2], + [[11, 'b'], 1], + ], + ], + ])('iterates in key order', async (insertOrder, expectedOrder) => { + const counter = get(); + await Promise.all(insertOrder.map(([key, value]) => counter.update(key, value as number))); + expect([...counter.entries({})]).toEqual(expectedOrder); + }); + }); +} diff --git a/yarn-project/kv-store/src/tests/aztec_map_tests.ts b/yarn-project/kv-store/src/tests/aztec_map_tests.ts new file mode 100644 index 000000000000..5ba3685ea729 --- /dev/null +++ b/yarn-project/kv-store/src/tests/aztec_map_tests.ts @@ -0,0 +1,108 @@ +import { beforeEach, describe, expect, it } from '@jest/globals'; + +import { type Key } from '../interfaces/common.js'; +import { type AztecMultiMap } from '../interfaces/index.js'; + +export function addMapTests(get: () => AztecMultiMap) { + describe('AztecMap', () => { + let map: AztecMultiMap; + + beforeEach(() => { + map = get(); + }); + + it('should be able to set and get values', async () => { + await map.set('foo', 'bar'); + await map.set('baz', 'qux'); + + expect(map.get('foo')).toEqual('bar'); + expect(map.get('baz')).toEqual('qux'); + expect(map.get('quux')).toEqual(undefined); + }); + + it('should be able to update values', async () => { + await map.set('foo', 'bar'); + expect(map.get('foo')).toEqual('bar'); + + await map.set('foo', 'qux'); + expect(map.get('foo')).toEqual('qux'); + }); + + it('should be able to set values if they do not exist', async () => { + expect(await map.setIfNotExists('foo', 'bar')).toEqual(true); + expect(await map.setIfNotExists('foo', 'baz')).toEqual(false); + + expect(map.get('foo')).toEqual('bar'); + }); + + it('should be able to delete values', async () => { + await map.set('foo', 'bar'); + await map.set('baz', 'qux'); + + expect(await map.delete('foo')).toEqual(true); + + expect(map.get('foo')).toEqual(undefined); + expect(map.get('baz')).toEqual('qux'); + }); + + it('should be able to iterate over entries', async () => { + await map.set('foo', 'bar'); + await map.set('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', 'quux'); + + 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(['baz', 'foo']); + }); + + it('should be able to get multiple values for a single key', async () => { + await map.set('foo', 'bar'); + await map.set('foo', 'baz'); + + expect([...map.getValues('foo')]).toEqual(['bar', 'baz']); + }); + + it('supports tuple keys', async () => { + const map = get(); + + 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']); + }); + }); +} diff --git a/yarn-project/kv-store/src/tests/aztec_singleton_tests.ts b/yarn-project/kv-store/src/tests/aztec_singleton_tests.ts new file mode 100644 index 000000000000..5247f1ad3d49 --- /dev/null +++ b/yarn-project/kv-store/src/tests/aztec_singleton_tests.ts @@ -0,0 +1,28 @@ +import { beforeEach, describe, expect, it } from '@jest/globals'; + +import { type Key } from '../interfaces/common.js'; +import { type AztecSingleton } from '../interfaces/singleton.js'; + +export function addSingletonTests(get: () => AztecSingleton) { + describe('AztecSingleton', () => { + let singleton: AztecSingleton; + beforeEach(() => { + singleton = get(); + }); + + it('returns undefined if the value is not set', () => { + expect(singleton.get()).toEqual(undefined); + }); + + it('should be able to set and get values', async () => { + expect(await singleton.set('foo')).toEqual(true); + expect(singleton.get()).toEqual('foo'); + }); + + it('overwrites the value if it is set again', async () => { + expect(await singleton.set('foo')).toEqual(true); + expect(await singleton.set('bar')).toEqual(true); + expect(singleton.get()).toEqual('bar'); + }); + }); +} diff --git a/yarn-project/kv-store/src/utils.ts b/yarn-project/kv-store/src/utils.ts index a85f01e454c9..50664e8e41d6 100644 --- a/yarn-project/kv-store/src/utils.ts +++ b/yarn-project/kv-store/src/utils.ts @@ -2,7 +2,7 @@ import { EthAddress } from '@aztec/foundation/eth-address'; import { type Logger } from '@aztec/foundation/log'; import { type AztecKVStore } from './interfaces/store.js'; -import { AztecLmdbStore } from './lmdb/store.js'; +import { AztecMemStore } from './mem/store.js'; /** * Clears the store if the rollup address does not match the one stored in the database. @@ -38,5 +38,6 @@ export async function initStoreForRollup( * @returns A new store */ export function openTmpStore(ephemeral: boolean = false): AztecKVStore { - return AztecLmdbStore.open(undefined, ephemeral); + return new AztecMemStore(); + // return AztecLmdbStore.open(undefined, ephemeral); }