Merge pull request #14 from Gusarich/patch2

add `generateMerkleProof` and `generateMerkleUpdate` functions for dictionaries

src/dict/Dictionary.spec.ts

@@ -8,6 +8,8 @@
 
 import { beginCell, Builder } from "../boc/Builder";
 import { Cell } from "../boc/Cell";
+import { exoticMerkleProof } from "../boc/cell/exoticMerkleProof";
+import { exoticMerkleUpdate } from "../boc/cell/exoticMerkleUpdate";
 import { Dictionary } from "./Dictionary";
 import fs from 'fs';
 
@@ -83,4 +85,53 @@ describe('Dictionary', () => {
             // console.warn(oracles);
         }
     });
+
+    it('should generate merkle proofs', () => {
+        let d = Dictionary.empty(
+            Dictionary.Keys.Uint(8),
+            Dictionary.Values.Uint(32)
+        );
+        d.set(1, 11);
+        d.set(2, 22);
+        d.set(3, 33);
+        d.set(4, 44);
+        d.set(5, 55);
+
+        for (let k = 1; k <= 5; k++) {
+            const proof = d.generateMerkleProof(k);
+            Cell.fromBoc(proof.toBoc());
+            expect(exoticMerkleProof(proof.bits, proof.refs).proofHash).toEqual(
+                Buffer.from(
+                    'ee41b86bd71f8224ebd01848b4daf4cd46d3bfb3e119d8b865ce7c2802511de3',
+                    'hex'
+                )
+            );
+        }
+    });
+
+    it('should generate merkle updates', () => {
+        let d = Dictionary.empty(
+            Dictionary.Keys.Uint(8),
+            Dictionary.Values.Uint(32)
+        );
+        d.set(1, 11);
+        d.set(2, 22);
+        d.set(3, 33);
+        d.set(4, 44);
+        d.set(5, 55);
+
+        for (let k = 1; k <= 5; k++) {
+            const update = d.generateMerkleUpdate(k, d.get(k)! * 2);
+            Cell.fromBoc(update.toBoc());
+            expect(
+                exoticMerkleUpdate(update.bits, update.refs).proofHash1
+            ).toEqual(
+                Buffer.from(
+                    'ee41b86bd71f8224ebd01848b4daf4cd46d3bfb3e119d8b865ce7c2802511de3',
+                    'hex'
+                )
+            );
+            d.set(k, Math.floor(d.get(k)! / 2));
+        }
+    });
 });

src/dict/Dictionary.ts

@@ -11,6 +11,8 @@ import { beginCell, Builder } from "../boc/Builder";
 import { Cell } from "../boc/Cell";
 import { Slice } from "../boc/Slice";
 import { Maybe } from "../utils/maybe";
+import { generateMerkleProof } from "./generateMerkleProof";
+import { generateMerkleUpdate } from "./generateMerkleUpdate";
 import { parseDict } from "./parseDict";
 import { serializeDict } from "./serializeDict";
 import { deserializeInternalKey, serializeInternalKey } from "./utils/internalKeySerializer";
@@ -364,6 +366,14 @@ export class Dictionary<K extends DictionaryKeyTypes, V> {
         // Store
         serializeDict(prepared, resolvedKey.bits, resolvedValue.serialize, builder);
     }
+
+    generateMerkleProof(key: K): Cell {
+        return generateMerkleProof(this, key, this._key!)
+    }
+
+    generateMerkleUpdate(key: K, newValue: V): Cell {
+        return generateMerkleUpdate(this, key, this._key!, newValue);
+    }
 }
 
 //

src/dict/generateMerkleProof.ts

@@ -0,0 +1,130 @@
+import { beginCell } from '../boc/Builder';
+import { Cell } from '../boc/Cell';
+import { Slice } from '../boc/Slice';
+import { DictionaryKeyTypes, Dictionary, DictionaryKey } from './Dictionary';
+import { readUnaryLength } from './utils/readUnaryLength';
+
+function convertToPrunedBranch(c: Cell): Cell {
+    return new Cell({
+        exotic: true,
+        bits: beginCell()
+            .storeUint(1, 8)
+            .storeUint(1, 8)
+            .storeBuffer(c.hash(0))
+            .storeUint(c.depth(0), 16)
+            .endCell()
+            .beginParse()
+            .loadBits(288),
+    });
+}
+
+function convertToMerkleProof(c: Cell): Cell {
+    return new Cell({
+        exotic: true,
+        bits: beginCell()
+            .storeUint(3, 8)
+            .storeBuffer(c.hash(0))
+            .storeUint(c.depth(0), 16)
+            .endCell()
+            .beginParse()
+            .loadBits(280),
+        refs: [c],
+    });
+}
+
+function doGenerateMerkleProof(
+    prefix: string,
+    slice: Slice,
+    n: number,
+    key: string
+): Cell {
+    // Reading label
+    const originalCell = slice.asCell();
+
+    let lb0 = slice.loadBit() ? 1 : 0;
+    let prefixLength = 0;
+    let pp = prefix;
+
+    if (lb0 === 0) {
+        // Short label detected
+
+        // Read
+        prefixLength = readUnaryLength(slice);
+
+        // Read prefix
+        for (let i = 0; i < prefixLength; i++) {
+            pp += slice.loadBit() ? '1' : '0';
+        }
+    } else {
+        let lb1 = slice.loadBit() ? 1 : 0;
+        if (lb1 === 0) {
+            // Long label detected
+            prefixLength = slice.loadUint(Math.ceil(Math.log2(n + 1)));
+            for (let i = 0; i < prefixLength; i++) {
+                pp += slice.loadBit() ? '1' : '0';
+            }
+        } else {
+            // Same label detected
+            let bit = slice.loadBit() ? '1' : '0';
+            prefixLength = slice.loadUint(Math.ceil(Math.log2(n + 1)));
+            for (let i = 0; i < prefixLength; i++) {
+                pp += bit;
+            }
+        }
+    }
+
+    if (n - prefixLength === 0) {
+        return originalCell;
+    } else {
+        let sl = originalCell.beginParse();
+        let left = sl.loadRef();
+        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);
+            }
+        }
+        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);
+            }
+        }
+
+        return beginCell()
+            .storeSlice(sl)
+            .storeRef(left)
+            .storeRef(right)
+            .endCell();
+    }
+}
+
+export function generateMerkleProof<K extends DictionaryKeyTypes, V>(
+    dict: Dictionary<K, V>,
+    key: K,
+    keyObject: DictionaryKey<K>
+): Cell {
+    const s = beginCell().storeDictDirect(dict).endCell().beginParse();
+    return convertToMerkleProof(
+        doGenerateMerkleProof(
+            '',
+            s,
+            keyObject.bits,
+            keyObject.serialize(key).toString(2).padStart(keyObject.bits, '0')
+        )
+    );
+}

src/dict/generateMerkleUpdate.ts

@@ -0,0 +1,32 @@
+import { beginCell } from '../boc/Builder';
+import { Cell } from '../boc/Cell';
+import { DictionaryKeyTypes, Dictionary, DictionaryKey } from './Dictionary';
+import { generateMerkleProof } from './generateMerkleProof';
+
+function convertToMerkleUpdate(c1: Cell, c2: Cell): Cell {
+    return new Cell({
+        exotic: true,
+        bits: beginCell()
+            .storeUint(4, 8)
+            .storeBuffer(c1.hash(0))
+            .storeBuffer(c2.hash(0))
+            .storeUint(c1.depth(0), 16)
+            .storeUint(c2.depth(0), 16)
+            .endCell()
+            .beginParse()
+            .loadBits(552),
+        refs: [c1, c2],
+    });
+}
+
+export function generateMerkleUpdate<K extends DictionaryKeyTypes, V>(
+    dict: Dictionary<K, V>,
+    key: K,
+    keyObject: DictionaryKey<K>,
+    newValue: V
+): Cell {
+    const oldProof = generateMerkleProof(dict, key, keyObject).refs[0];
+    dict.set(key, newValue);
+    const newProof = generateMerkleProof(dict, key, keyObject).refs[0];
+    return convertToMerkleUpdate(oldProof, newProof);
+}

src/dict/utils/readUnaryLength.ts

@@ -0,0 +1,9 @@
+import { Slice } from '../../boc/Slice';
+
+export function readUnaryLength(slice: Slice) {
+    let res = 0;
+    while (slice.loadBit()) {
+        res++;
+    }
+    return res;
+}

src/index.ts

@@ -32,6 +32,10 @@ export { exoticMerkleProof } from './boc/cell/exoticMerkleProof';
 export { exoticMerkleUpdate } from './boc/cell/exoticMerkleUpdate';
 export { exoticPruned } from './boc/cell/exoticPruned';
 
+// Merkle trees
+export { generateMerkleProof } from './dict/generateMerkleProof'
+export { generateMerkleUpdate } from './dict/generateMerkleUpdate'
+
 // Tuples
 export { Tuple, TupleItem, TupleItemNull, TupleItemInt, TupleItemNaN, TupleItemCell, TupleItemSlice, TupleItemBuilder } from './tuple/tuple';
 export { parseTuple, serializeTuple } from './tuple/tuple';