Implement proving that key is in hashmap (merkle proof)

Author: aleksej-paschenko

boc/cell.go

@@ -424,3 +424,17 @@ func (c *Cell) bocReprWithoutRefs(mask levelMask) []byte {
 func (c *Cell) Level() int {
 	return c.mask.Level()
 }
+
+func (c *Cell) GetMerkleRoot() ([32]byte, error) {
+	if c.CellType() != MerkleProofCell {
+		return [32]byte{}, errors.New("not merkle proof cell")
+	}
+	bytes, err := c.ReadBytes(33)
+	if err != nil {
+		return [32]byte{}, err
+	}
+	var hash [32]byte
+	copy(hash[:], bytes[1:])
+	return hash, nil
+
+}

boc/immutable_cell.go

@@ -3,6 +3,7 @@ package boc
 import (
 	"crypto/sha256"
 	"encoding/binary"
+	"fmt"
 )
 
 // immutableCell provides a convenient way to calculate a cell's hash and depth.
@@ -13,9 +14,12 @@ type immutableCell struct {
 	hashes   [][]byte
 	depths   []int
 
-	// bitsBuf points to a content of this cell.
+	// bitsBuf points to the content of this cell.
 	// Used if this cell is a pruned branch cell to extract hashes and depth of an original cell.
 	bitsBuf []byte
+	// bitsLen is a length of the content of this cell in bits.
+	// it's not always equal to len(bitsBuf) because bitsBuf can contain more bytes than the actual content.
+	bitsLen int
 }
 
 // newImmutableCell returns a new instance of immutable cell.
@@ -31,6 +35,7 @@ func newImmutableCell(c *Cell, cache map[*Cell]*immutableCell) (*immutableCell,
 		hashes:   make([][]byte, 0, c.mask.HashesCount()),
 		depths:   make([]int, 0, c.mask.HashesCount()),
 		bitsBuf:  c.bits.buf,
+		bitsLen:  c.bits.len,
 	}
 	for _, ref := range c.refs {
 		if ref == nil {
@@ -154,3 +159,59 @@ func (ic *immutableCell) Depth(level int) int {
 	}
 	return ic.depths[index]
 }
+
+// pruneCells return the current subtree (which this cell represents) with pruned cells.
+// if this cell is pruned, "pruneCells" returns a new pruned branch cell instead of this cell.
+// As of now, this function doesn't work with MerkleProofCell and MerkleUpdateCell.
+func (ic *immutableCell) pruneCells(pruned map[*immutableCell]struct{}) (*Cell, error) {
+	if ic.cellType == MerkleProofCell || ic.cellType == MerkleUpdateCell {
+		return nil, fmt.Errorf("unsupported cell type: %v", ic.cellType)
+	}
+	if _, ok := pruned[ic]; ok {
+		// we are pruned
+		// let's replace this cell with a pruned branch cell
+		prunedCell := NewCell()
+		prunedCell.mask = 1
+		prunedCell.cellType = PrunedBranchCell
+		if err := prunedCell.WriteUint(1, 8); err != nil {
+			return nil, err
+		}
+		if err := prunedCell.WriteUint(1, 8); err != nil {
+			return nil, err
+		}
+		if err := prunedCell.WriteBytes(ic.Hash(0)); err != nil {
+			return nil, err
+		}
+		if err := prunedCell.WriteUint(uint64(ic.Depth(0)), 16); err != nil {
+			return nil, err
+		}
+		prunedCell.ResetCounters()
+		return prunedCell, nil
+	}
+	// all good,
+	// going down the tree
+	bits := BitString{
+		buf: ic.bitsBuf,
+		cap: ic.bitsLen,
+		len: ic.bitsLen,
+	}
+	res := Cell{
+		bits:     bits,
+		refs:     [4]*Cell{},
+		cellType: ic.cellType,
+		mask:     ic.mask,
+	}
+	mask := ic.mask
+	for i, ref := range ic.refs {
+		cell, err := ref.pruneCells(pruned)
+		if err != nil {
+			return nil, err
+		}
+		if cell.mask > 0 {
+			mask |= cell.mask
+		}
+		res.refs[i] = cell
+	}
+	res.mask = mask
+	return &res, nil
+}

boc/merkle_proof.go

@@ -0,0 +1,55 @@
+package boc
+
+type MerkleProver struct {
+	root *immutableCell
+
+	pruned map[*immutableCell]struct{}
+}
+
+func NewMerkleProver(root *Cell) (*MerkleProver, error) {
+	immRoot, err := newImmutableCell(root, make(map[*Cell]*immutableCell))
+	if err != nil {
+		return nil, err
+	}
+	return &MerkleProver{root: immRoot, pruned: make(map[*immutableCell]struct{})}, nil
+}
+
+type Cursor struct {
+	cell   *immutableCell
+	prover *MerkleProver
+}
+
+func (p *MerkleProver) Cursor() *Cursor {
+	return &Cursor{cell: p.root, prover: p}
+}
+
+func (p *MerkleProver) CreateProof() ([]byte, error) {
+	immRoot, err := p.root.pruneCells(p.pruned)
+	if err != nil {
+		return nil, err
+	}
+	mp := NewCell()
+	mp.cellType = MerkleProofCell
+	if err := mp.WriteUint(3, 8); err != nil {
+		return nil, err
+	}
+	if err := mp.WriteBytes(p.root.Hash(0)); err != nil {
+		return nil, err
+	}
+	if err := mp.WriteUint(uint64(p.root.Depth(0)), 16); err != nil {
+		return nil, err
+	}
+	if err := mp.AddRef(immRoot); err != nil {
+		return nil, err
+	}
+	mp.ResetCounters()
+	return SerializeBoc(mp, false, false, false, 0)
+}
+
+func (c *Cursor) Prune() {
+	c.prover.pruned[c.cell] = struct{}{}
+}
+
+func (c *Cursor) Ref(ref int) *Cursor {
+	return &Cursor{cell: c.cell.refs[ref], prover: c.prover}
+}

tlb/hashmap.go

@@ -171,6 +171,54 @@ func countLeafs(keySize, leftKeySize int, c *boc.Cell) (int, error) {
 	return 1, nil
 }
 
+func ProveKeyInHashmap(cell *boc.Cell, key boc.BitString) ([]byte, error) {
+	keySize := key.BitsAvailableForRead()
+	prover, err := boc.NewMerkleProver(cell)
+	if err != nil {
+		return nil, err
+	}
+	cursor := prover.Cursor()
+	bitString := boc.NewBitString(keySize)
+	prefix := &bitString
+	for {
+		var err error
+		var size int
+		size, prefix, err = loadLabel(keySize, cell, prefix)
+		if err != nil {
+			return nil, err
+		}
+		if keySize <= size {
+			break
+		}
+		if _, err = key.ReadBits(size); err != nil {
+			return nil, err
+		}
+
+		isRight, err := key.ReadBit()
+		if err != nil {
+			return nil, err
+		}
+		keySize = keySize - size - 1
+		next, err := cell.NextRef()
+		if err != nil {
+			return nil, err
+		}
+		if isRight {
+			cursor.Ref(0).Prune()
+			next, err = cell.NextRef()
+			if err != nil {
+				return nil, err
+			}
+			cursor = cursor.Ref(1)
+		} else {
+			cursor.Ref(1).Prune()
+			cursor = cursor.Ref(0)
+		}
+		cell = next
+	}
+	return prover.CreateProof()
+}
+
 func (h *Hashmap[keyT, T]) mapInner(keySize, leftKeySize int, c *boc.Cell, keyPrefix *boc.BitString, decoder *Decoder) error {
 	var err error
 	var size int