add key-value storage

src/key-value.test.ts

@@ -0,0 +1,42 @@
+import { Field, Provable } from 'o1js';
+import { KeyValueStorage } from './key-value.js';
+
+class TestStorage extends KeyValueStorage<Field> {
+    calculateLeaf(leaf: Field): Field {
+        return leaf;
+    }
+
+    calculateKey(index: Field): Field {
+        return index;
+    }
+}
+
+describe('Key Value Storage', () => {
+    let testStorage: TestStorage;
+
+    it('Should create new storage', async () => {
+        testStorage = new TestStorage();
+    });
+
+    it('Should update raw value', async () => {
+        let value = Field(123);
+        testStorage.updateRawValue(Field(0), value);
+        let witness = testStorage.getWitness(Field(0));
+        expect(witness.calculateRoot(value).toBigInt()).toEqual(
+            testStorage.root.toBigInt()
+        );
+    });
+
+    it('Should update value', async () => {
+        let value = Field(123);
+        testStorage.updateValue(Field(1), value);
+        let witness = testStorage.getWitness(Field(1));
+        expect(witness.calculateRoot(value).toBigInt()).toEqual(
+            testStorage.root.toBigInt()
+        );
+    });
+
+    it('Should get public data', async () => {
+        Provable.log(testStorage.root, testStorage.mapping, testStorage.values);
+    });
+});

src/key-value.ts

@@ -0,0 +1,83 @@
+import { Field, MerkleTree } from 'o1js';
+import { BaseStorage } from './base-storage.js';
+import { EmptyMT254, MTWitness254 } from './merkle-tree.js';
+
+export abstract class KeyValueStorage<RawValue>
+    implements BaseStorage<RawValue>
+{
+    private _mapping: MerkleTree;
+    private _leafs: {
+        [key: string]: { raw: RawValue | undefined; leaf: Field };
+    };
+
+    constructor(
+        leafs?: {
+            key: Field;
+            value: RawValue | Field;
+            isRaw: boolean;
+        }[]
+    ) {
+        this._mapping = EmptyMT254();
+        this._leafs = {};
+        if (leafs) {
+            for (let i = 0; i < leafs.length; i++) {
+                let { key, value, isRaw } = leafs[i];
+                if (isRaw) {
+                    this.updateRawValue(key, value as RawValue);
+                } else {
+                    this.updateValue(key, value as Field);
+                }
+            }
+        }
+    }
+
+    get root(): Field {
+        return this._mapping.getRoot();
+    }
+
+    get leafs(): { [key: string]: { raw: RawValue | undefined; leaf: Field } } {
+        return this._leafs;
+    }
+
+    get mapping(): MerkleTree {
+        return this._mapping;
+    }
+
+    get values(): {
+        [key: string]: { raw: RawValue | undefined; leaf: Field };
+    } {
+        return this.leafs;
+    }
+
+    abstract calculateLeaf(args: RawValue): Field;
+    calculateValue(args: RawValue): Field {
+        return this.calculateLeaf(args);
+    }
+
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    abstract calculateKey(args: any): Field;
+
+    getWitness(key: Field): MTWitness254 {
+        return new MTWitness254(this._mapping.getWitness(key.toBigInt()));
+    }
+
+    updateLeaf({ index }: { index: Field }, leaf: Field) {
+        let leafId = index.toString();
+        this._mapping.setLeaf(index.toBigInt(), leaf);
+        this._leafs[leafId] = { raw: undefined, leaf };
+    }
+
+    updateValue(index: Field, leaf: Field) {
+        this.updateLeaf({ index }, leaf);
+    }
+
+    updateRawLeaf({ index }: { index: Field }, rawLeaf: RawValue) {
+        let leaf = this.calculateLeaf(rawLeaf);
+        this.updateLeaf({ index }, leaf);
+        this._leafs[index.toString()] = { raw: rawLeaf, leaf };
+    }
+
+    updateRawValue(key: Field, rawValue: RawValue) {
+        this.updateRawLeaf({ index: key }, rawValue);
+    }
+}

src/merkle-tree.ts

@@ -2,37 +2,53 @@ import { MerkleTree, MerkleWitness } from 'o1js';
 import { Witness } from './base-storage';
 
 export class MTWitness2 extends MerkleWitness(2) {}
-export const EmptyMTWitness2 = (wtn: Witness) => new MTWitness2(wtn);
+export const NewMTWitness2 = (wtn: Witness) => new MTWitness2(wtn);
 export const EmptyMT2 = () => new MerkleTree(2);
 
+export class MTWitness3 extends MerkleWitness(3) {}
+export const NewMTWitness3 = (wtn: Witness) => new MTWitness3(wtn);
+export const EmptyMT3 = () => new MerkleTree(3);
+
 export class MTWitness4 extends MerkleWitness(4) {}
-export const EmptyMTWitness4 = (wtn: Witness) => new MTWitness4(wtn);
+export const NewMTWitness4 = (wtn: Witness) => new MTWitness4(wtn);
 export const EmptyMT4 = () => new MerkleTree(4);
 
+export class MTWitness5 extends MerkleWitness(5) {}
+export const NewMTWitness5 = (wtn: Witness) => new MTWitness5(wtn);
+export const EmptyMT5 = () => new MerkleTree(5);
+
 export class MTWitness6 extends MerkleWitness(6) {}
-export const EmptyMTWitness6 = (wtn: Witness) => new MTWitness6(wtn);
+export const NewMTWitness6 = (wtn: Witness) => new MTWitness6(wtn);
 export const EmptyMT6 = () => new MerkleTree(6);
 
 export class MTWitness8 extends MerkleWitness(8) {}
-export const EmptyMTWitness8 = (wtn: Witness) => new MTWitness8(wtn);
+export const NewMTWitness8 = (wtn: Witness) => new MTWitness8(wtn);
 export const EmptyMT8 = () => new MerkleTree(8);
 
 export class MTWitness10 extends MerkleWitness(10) {}
-export const EmptyMTWitness10 = (wtn: Witness) => new MTWitness10(wtn);
+export const NewMTWitness10 = (wtn: Witness) => new MTWitness10(wtn);
 export const EmptyMT10 = () => new MerkleTree(10);
 
 export class MTWitness16 extends MerkleWitness(16) {}
-export const EmptyMTWitness16 = (wtn: Witness) => new MTWitness16(wtn);
+export const NewMTWitness16 = (wtn: Witness) => new MTWitness16(wtn);
 export const EmptyMT16 = () => new MerkleTree(16);
 
 export class MTWitness32 extends MerkleWitness(32) {}
-export const EmptyMTWitness32 = (wtn: Witness) => new MTWitness32(wtn);
+export const NewMTWitness32 = (wtn: Witness) => new MTWitness32(wtn);
 export const EmptyMT32 = () => new MerkleTree(32);
 
 export class MTWitness64 extends MerkleWitness(64) {}
-export const EmptyMTWitness64 = (wtn: Witness) => new MTWitness64(wtn);
+export const NewMTWitness64 = (wtn: Witness) => new MTWitness64(wtn);
 export const EmptyMT64 = () => new MerkleTree(64);
 
 export class MTWitness128 extends MerkleWitness(128) {}
-export const EmptyMTWitness128 = (wtn: Witness) => new MTWitness128(wtn);
+export const NewMTWitness128 = (wtn: Witness) => new MTWitness128(wtn);
 export const EmptyMT128 = () => new MerkleTree(128);
+
+export class MTWitness254 extends MerkleWitness(254) {}
+export const NewMTWitness254 = (wtn: Witness) => new MTWitness254(wtn);
+export const EmptyMT254 = () => new MerkleTree(254);
+
+export class MTWitness256 extends MerkleWitness(256) {}
+export const NewMTWitness256 = (wtn: Witness) => new MTWitness256(wtn);
+export const EmptyMT256 = () => new MerkleTree(256);

src/one-level.test.ts

@@ -1,6 +1,6 @@
 import { Field, Provable } from 'o1js';
 import { OneLevelStorage } from './one-level.js';
-import { EmptyMT4, EmptyMTWitness4, MTWitness4 } from './merkle-tree.js';
+import { EmptyMT4, NewMTWitness4, MTWitness4 } from './merkle-tree.js';
 
 class TestStorage extends OneLevelStorage<Field, MTWitness4> {
     calculateLeaf(leaf: Field): Field {
@@ -16,7 +16,7 @@ describe('Single Level Storage', () => {
     let testStorage: TestStorage;
 
     it('Should create new storage', async () => {
-        testStorage = new TestStorage(EmptyMT4, EmptyMTWitness4);
+        testStorage = new TestStorage(EmptyMT4, NewMTWitness4);
     });
 
     it('Should update raw leaf', async () => {
@@ -40,29 +40,4 @@ describe('Single Level Storage', () => {
     it('Should get public data', async () => {
         Provable.log(testStorage.root, testStorage.leafs, testStorage.level1);
     });
-
-    // it('Should manage 2-level MT', async () => {
-    //     let publicKey = PrivateKey.random().toPublicKey();
-    //     publicKeyStorage.updateRawLeaf(
-    //         { level1Index: Field(1), level2Index: Field(2) },
-    //         publicKey.toGroup()
-    //     );
-    //     let witnesses = publicKeyStorage.getWitness(Field(1), Field(2));
-    //     expect(
-    //         witnesses.level2
-    //             .calculateRoot(
-    //                 PublicKeyStorage.calculateLeaf(publicKey.toGroup())
-    //             )
-    //             .toBigInt()
-    //     ).toEqual(
-    //         publicKeyStorage.level2s[Field(1).toString()].getRoot().toBigInt()
-    //     );
-    //     expect(
-    //         witnesses.level1
-    //             .calculateRoot(
-    //                 publicKeyStorage.level2s[Field(1).toString()].getRoot()
-    //             )
-    //             .toBigInt()
-    //     ).toEqual(publicKeyStorage.root.toBigInt());
-    // });
 });

src/two-level.test.ts

@@ -3,8 +3,8 @@ import { TwoLevelStorage } from './two-level.js';
 import {
     EmptyMT4,
     EmptyMT6,
-    EmptyMTWitness4,
-    EmptyMTWitness6,
+    NewMTWitness4,
+    NewMTWitness6,
     MTWitness4,
     MTWitness6,
 } from './merkle-tree.js';
@@ -29,9 +29,9 @@ describe('Two Level Storage', () => {
     it('Should create new storage', async () => {
         testStorage = new TestStorage(
             EmptyMT4,
-            EmptyMTWitness4,
+            NewMTWitness4,
             EmptyMT6,
-            EmptyMTWitness6
+            NewMTWitness6
         );
     });