Merge pull request #30 from akifoq/main

More flexible merkle proof generation

Author: dvlkv

src/boc/cell/exoticMerkleProof.ts

@@ -9,6 +9,7 @@
 import { BitReader } from "../BitReader";
 import { BitString } from "../BitString";
 import { Cell } from "../Cell";
+import { beginCell } from '../Builder';
 
 export function exoticMerkleProof(bits: BitString, refs: Cell[]) {
     const reader = new BitReader(bits);
@@ -48,4 +49,13 @@ export function exoticMerkleProof(bits: BitString, refs: Cell[]) {
         proofDepth,
         proofHash
     };
+}
+
+export function convertToMerkleProof(c: Cell): Cell {
+    return beginCell()
+        .storeUint(3, 8)
+        .storeBuffer(c.hash(0))
+        .storeUint(c.depth(0), 16)
+        .storeRef(c)
+        .endCell({ exotic: true });
 }

src/dict/Dictionary.spec.ts

@@ -121,14 +121,22 @@ describe('Dictionary', () => {
         d.set(4, 44);
         d.set(5, 55);
 
+        const dictHash = beginCell().storeDictDirect(d).endCell().hash();
         for (let k = 1; k <= 5; k++) {
-            const proof = d.generateMerkleProof(k);
+            const proof = d.generateMerkleProof([k]);
             Cell.fromBoc(proof.toBoc());
             expect(exoticMerkleProof(proof.bits, proof.refs).proofHash).toEqual(
-                Buffer.from(
-                    'ee41b86bd71f8224ebd01848b4daf4cd46d3bfb3e119d8b865ce7c2802511de3',
-                    'hex'
-                )
+                dictHash
+            );
+
+            // todo: parse the pruned dictionary and check the presence of the keys
+        }
+
+        for (let k = 1; k <= 3; k++) {
+            const proof = d.generateMerkleProof([k, k + 1, k + 2]);
+            Cell.fromBoc(proof.toBoc());
+            expect(exoticMerkleProof(proof.bits, proof.refs).proofHash).toEqual(
+                dictHash
             );
         }
     });

src/dict/Dictionary.ts

@@ -12,7 +12,7 @@ import { Cell } from "../boc/Cell";
 import { Slice } from "../boc/Slice";
 import { BitString } from "../boc/BitString";
 import { Maybe } from "../utils/maybe";
-import { generateMerkleProof } from "./generateMerkleProof";
+import { generateMerkleProof, generateMerkleProofDirect } from "./generateMerkleProof";
 import { generateMerkleUpdate } from "./generateMerkleUpdate";
 import { parseDict } from "./parseDict";
 import { serializeDict } from "./serializeDict";
@@ -392,8 +392,23 @@ export class Dictionary<K extends DictionaryKeyTypes, V> {
         serializeDict(prepared, resolvedKey.bits, resolvedValue.serialize, builder);
     }
 
-    generateMerkleProof(key: K): Cell {
-        return generateMerkleProof(this, key, this._key!)
+    /**
+     * Generate merkle proof for multiple keys in the dictionary
+     * @param keys an array of the keys
+     * @returns generated merkle proof cell
+     */
+    generateMerkleProof(keys: K[]): Cell {
+        return generateMerkleProof(this, keys, this._key!)
+    }
+
+    /**
+     * Low level method for generating pruned dictionary directly.
+     * The result can be used as a part of a bigger merkle proof
+     * @param keys an array of the keys
+     * @returns cell that contains the pruned dictionary
+     */
+    generateMerkleProofDirect(keys: K[]): Cell {
+        return generateMerkleProofDirect(this, keys, this._key!)
     }
 
     generateMerkleUpdate(key: K, newValue: V): Cell {

src/dict/generateMerkleProof.ts

@@ -3,6 +3,7 @@ import { Cell } from '../boc/Cell';
 import { Slice } from '../boc/Slice';
 import { DictionaryKeyTypes, Dictionary, DictionaryKey } from './Dictionary';
 import { readUnaryLength } from './utils/readUnaryLength';
+import { convertToMerkleProof } from '../boc/cell/exoticMerkleProof';
 
 function convertToPrunedBranch(c: Cell): Cell {
     return beginCell()
@@ -13,24 +14,20 @@ function convertToPrunedBranch(c: Cell): Cell {
         .endCell({ exotic: true });
 }
 
-function convertToMerkleProof(c: Cell): Cell {
-    return beginCell()
-        .storeUint(3, 8)
-        .storeBuffer(c.hash(0))
-        .storeUint(c.depth(0), 16)
-        .storeRef(c)
-        .endCell({ exotic: true });
-}
-
 function doGenerateMerkleProof(
     prefix: string,
     slice: Slice,
     n: number,
-    key: string
+    keys: string[]
 ): Cell {
     // Reading label
     const originalCell = slice.asCell();
 
+    if (keys.length == 0) {
+        // no keys to prove, prune the whole subdict
+        return convertToPrunedBranch(originalCell);
+    }
+
     let lb0 = slice.loadBit() ? 1 : 0;
     let prefixLength = 0;
     let pp = prefix;
@@ -71,28 +68,26 @@ function doGenerateMerkleProof(
         let right = sl.loadRef();
         // NOTE: Left and right branches are implicitly contain prefixes '0' and '1'
         if (!left.isExotic) {
-            if (pp + '0' === key.slice(0, pp.length + 1)) {
-                left = doGenerateMerkleProof(
-                    pp + '0',
-                    left.beginParse(),
-                    n - prefixLength - 1,
-                    key
-                );
-            } else {
-                left = convertToPrunedBranch(left);
-            }
+            const leftKeys = keys.filter((key) => {
+                return pp + '0' === key.slice(0, pp.length + 1);
+            });
+            left = doGenerateMerkleProof(
+                pp + '0',
+                left.beginParse(),
+                n - prefixLength - 1,
+                leftKeys
+            );
         }
         if (!right.isExotic) {
-            if (pp + '1' === key.slice(0, pp.length + 1)) {
-                right = doGenerateMerkleProof(
-                    pp + '1',
-                    right.beginParse(),
-                    n - prefixLength - 1,
-                    key
-                );
-            } else {
-                right = convertToPrunedBranch(right);
-            }
+            const rightKeys = keys.filter((key) => {
+                return pp + '1' === key.slice(0, pp.length + 1);
+            });
+            right = doGenerateMerkleProof(
+                pp + '1',
+                right.beginParse(),
+                n - prefixLength - 1,
+                rightKeys
+            );
         }
 
         return beginCell()
@@ -103,18 +98,25 @@ function doGenerateMerkleProof(
     }
 }
 
-export function generateMerkleProof<K extends DictionaryKeyTypes, V>(
+export function generateMerkleProofDirect<K extends DictionaryKeyTypes, V>(
     dict: Dictionary<K, V>,
-    key: K,
+    keys: K[],
     keyObject: DictionaryKey<K>
 ): Cell {
-    const s = beginCell().storeDictDirect(dict).endCell().beginParse();
-    return convertToMerkleProof(
-        doGenerateMerkleProof(
+    keys = keys.filter((key) => dict.has(key)); 
+    const s = beginCell().storeDictDirect(dict).asSlice();
+    return doGenerateMerkleProof(
             '',
             s,
             keyObject.bits,
-            keyObject.serialize(key).toString(2).padStart(keyObject.bits, '0')
-        )
+            keys.map((key) => keyObject.serialize(key).toString(2).padStart(keyObject.bits, '0'))
     );
 }
+
+export function generateMerkleProof<K extends DictionaryKeyTypes, V>(
+    dict: Dictionary<K, V>,
+    keys: K[],
+    keyObject: DictionaryKey<K>
+): Cell {
+    return convertToMerkleProof(generateMerkleProofDirect(dict, keys, keyObject));
+}

src/dict/generateMerkleUpdate.ts

@@ -21,8 +21,8 @@ export function generateMerkleUpdate<K extends DictionaryKeyTypes, V>(
     keyObject: DictionaryKey<K>,
     newValue: V
 ): Cell {
-    const oldProof = generateMerkleProof(dict, key, keyObject).refs[0];
+    const oldProof = generateMerkleProof(dict, [key], keyObject).refs[0];
     dict.set(key, newValue);
-    const newProof = generateMerkleProof(dict, key, keyObject).refs[0];
+    const newProof = generateMerkleProof(dict, [key], keyObject).refs[0];
     return convertToMerkleUpdate(oldProof, newProof);
 }