fix logic of length calculation

package.json

@@ -1,6 +1,6 @@
 {
     "name": "@auxo-dev/zkapp-offchain-storage",
-    "version": "1.0.2",
+    "version": "1.0.3",
     "description": "",
     "author": "",
     "license": "Apache-2.0",

src/address-map.ts

@@ -7,7 +7,7 @@ import { OneLevelStorage } from './one-level.js';
 export { AddressMap, ZkAppRef };
 
 const ZkAppRef = (num_addr: number) => {
-    const [MTWitness, ,] = getBestHeight(num_addr);
+    const [MTWitness, ,] = getBestHeight(BigInt(num_addr));
     class _ZkAppRef extends Struct({
         address: PublicKey,
         witness: MTWitness,
@@ -16,7 +16,7 @@ const ZkAppRef = (num_addr: number) => {
 };
 
 const AddressMap = (num_addr: number) => {
-    const [MTWitness, NewMTWitness, EmptyMT] = getBestHeight(num_addr);
+    const [MTWitness, NewMTWitness, EmptyMT] = getBestHeight(BigInt(num_addr));
     class _ZkAppRef extends ZkAppRef(num_addr) {}
 
     return class _AddressMap extends OneLevelStorage<

src/key-value.ts

@@ -1,6 +1,6 @@
 import { Field, MerkleTree } from 'o1js';
 import { BaseStorage } from './base-storage.js';
-import { EmptyMT254, MTWitness254 } from './merkle-tree.js';
+import { EmptyMT255, MTWitness255 } from './merkle-tree.js';
 
 export abstract class KeyValueStorage<RawValue>
     implements BaseStorage<RawValue>
@@ -17,7 +17,7 @@ export abstract class KeyValueStorage<RawValue>
             isRaw: boolean;
         }[]
     ) {
-        this._mapping = EmptyMT254();
+        this._mapping = EmptyMT255();
         this._leafs = {};
         if (leafs) {
             for (let i = 0; i < leafs.length; i++) {
@@ -32,7 +32,7 @@ export abstract class KeyValueStorage<RawValue>
     }
 
     get height(): number {
-        return MTWitness254.height;
+        return MTWitness255.height;
     }
 
     get size(): bigint {
@@ -65,8 +65,8 @@ export abstract class KeyValueStorage<RawValue>
     // 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()));
+    getWitness(key: Field): MTWitness255 {
+        return new MTWitness255(this._mapping.getWitness(key.toBigInt()));
     }
 
     updateLeaf({ index }: { index: Field }, leaf: Field) {

src/merkle-tree.ts

@@ -1,9 +1,7 @@
 import { MerkleTree, MerkleWitness } from 'o1js';
 import { Witness } from './base-storage';
 
-export const SUPPORTED_HEIGHTS = [
-    2, 3, 4, 5, 6, 8, 10, 16, 32, 64, 128, 254, 256,
-];
+export const SUPPORTED_HEIGHTS = [2, 3, 4, 5, 6, 8, 10, 16, 32, 64, 128, 255];
 
 export class MTWitness2 extends MerkleWitness(2) {}
 export const NewMTWitness2 = (wtn: Witness) => new MTWitness2(wtn);
@@ -49,49 +47,13 @@ export class MTWitness128 extends MerkleWitness(128) {}
 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 MTWitness255 extends MerkleWitness(255) {}
+export const NewMTWitness255 = (wtn: Witness) => new MTWitness255(wtn);
+export const EmptyMT255 = () => new MerkleTree(255);
 
-export class MTWitness256 extends MerkleWitness(256) {}
-export const NewMTWitness256 = (wtn: Witness) => new MTWitness256(wtn);
-export const EmptyMT256 = () => new MerkleTree(256);
-
-// export function getBestHeight(size: number) {
-//     if (size <= 2) {
-//         return [MTWitness2, NewMTWitness2, EmptyMT2];
-//     } else if (size <= 3) {
-//         return [MTWitness3, NewMTWitness3, EmptyMT3];
-//     } else if (size <= 4) {
-//         return [MTWitness4, NewMTWitness4, EmptyMT4];
-//     } else if (size <= 5) {
-//         return [MTWitness5, NewMTWitness5, EmptyMT5];
-//     } else if (size <= 6) {
-//         return [MTWitness6, NewMTWitness6, EmptyMT6];
-//     } else if (size <= 8) {
-//         return [MTWitness8, NewMTWitness8, EmptyMT8];
-//     } else if (size <= 10) {
-//         return [MTWitness10, NewMTWitness10, EmptyMT10];
-//     } else if (size <= 16) {
-//         return [MTWitness16, NewMTWitness16, EmptyMT16];
-//     } else if (size <= 32) {
-//         return [MTWitness32, NewMTWitness32, EmptyMT32];
-//     } else if (size <= 64) {
-//         return [MTWitness64, NewMTWitness64, EmptyMT64];
-//     } else if (size <= 128) {
-//         return [MTWitness128, NewMTWitness128, EmptyMT128];
-//     } else if (size <= 254) {
-//         return [MTWitness254, NewMTWitness254, EmptyMT254];
-//     } else if (size <= 256) {
-//         return [MTWitness256, NewMTWitness256, EmptyMT256];
-//     } else {
-//         throw new Error(`Unsupported tree height: ${size}`);
-//     }
-// }
-
-export function getBestHeight(size: number) {
+export function getBestHeight(size: bigint) {
     for (const height of SUPPORTED_HEIGHTS) {
-        if (size <= height) {
+        if (size <= BigInt(2 ** (height - 1))) {
             return [
                 eval(`MTWitness${height}`),
                 eval(`NewMTWitness${height}`),

src/one-level.test.ts

@@ -3,7 +3,8 @@ import { OneLevelStorage } from './one-level.js';
 import { getBestHeight } from './merkle-tree.js';
 
 describe('Single Level Storage', () => {
-    const [MTWitness, NewMTWitness, EmptyMT] = getBestHeight(4);
+    const size = BigInt(2 ** 3);
+    const [MTWitness, NewMTWitness, EmptyMT] = getBestHeight(size);
     class TestStorage extends OneLevelStorage<Field, typeof MTWitness> {
         static readonly height = MTWitness.height;
 
@@ -53,6 +54,15 @@ describe('Single Level Storage', () => {
         );
     });
 
+    it('Should have correct size', async () => {
+        for (let i = 0; i < Number(size); i++) {
+            testStorage.updateLeaf({ level1Index: Field(i) }, Field(i));
+        }
+        expect(() =>
+            testStorage.updateLeaf({ level1Index: Field(size) }, Field(size))
+        ).toThrowError();
+    });
+
     it('Should get public data', async () => {
         testStorage.root,
             testStorage.leafs,

src/two-level.test.ts

@@ -3,8 +3,8 @@ import { TwoLevelStorage } from './two-level.js';
 import { getBestHeight } from './merkle-tree.js';
 
 describe('Two Level Storage', () => {
-    const [MTWitnessL1, NewMTWitnessL1, EmptyMTL1] = getBestHeight(4);
-    const [MTWitnessL2, NewMTWitnessL2, EmptyMTL2] = getBestHeight(6);
+    const [MTWitnessL1, NewMTWitnessL1, EmptyMTL1] = getBestHeight(16n);
+    const [MTWitnessL2, NewMTWitnessL2, EmptyMTL2] = getBestHeight(64n);
 
     class TestStorage extends TwoLevelStorage<
         Field,
@@ -79,9 +79,7 @@ describe('Two Level Storage', () => {
             testStorage.leafs,
             testStorage.level1,
             testStorage.level2s,
-            testStorage.level2(Field(0)),
-            testStorage.height1,
-            testStorage.height2,
-            testStorage.size;
+            testStorage.level2(Field(0));
+        console.log(testStorage.height1, testStorage.height2, testStorage.size);
     });
 });