diff --git a/boc/merkle_proof.go b/boc/merkle_proof.go index 538b8569..fe9ea4f4 100644 --- a/boc/merkle_proof.go +++ b/boc/merkle_proof.go @@ -2,8 +2,6 @@ package boc type MerkleProver struct { root *immutableCell - - pruned map[*immutableCell]struct{} } func NewMerkleProver(root *Cell) (*MerkleProver, error) { @@ -11,20 +9,20 @@ func NewMerkleProver(root *Cell) (*MerkleProver, error) { if err != nil { return nil, err } - return &MerkleProver{root: immRoot, pruned: make(map[*immutableCell]struct{})}, nil + return &MerkleProver{root: immRoot}, nil } type Cursor struct { cell *immutableCell - prover *MerkleProver + pruned map[*immutableCell]struct{} } func (p *MerkleProver) Cursor() *Cursor { - return &Cursor{cell: p.root, prover: p} + return &Cursor{cell: p.root, pruned: make(map[*immutableCell]struct{})} } -func (p *MerkleProver) CreateProof() ([]byte, error) { - immRoot, err := p.root.pruneCells(p.pruned) +func (p *MerkleProver) CreateProof(cursor *Cursor) ([]byte, error) { + immRoot, err := p.root.pruneCells(cursor.pruned) if err != nil { return nil, err } @@ -47,9 +45,9 @@ func (p *MerkleProver) CreateProof() ([]byte, error) { } func (c *Cursor) Prune() { - c.prover.pruned[c.cell] = struct{}{} + c.pruned[c.cell] = struct{}{} } func (c *Cursor) Ref(ref int) *Cursor { - return &Cursor{cell: c.cell.refs[ref], prover: c.prover} + return &Cursor{cell: c.cell.refs[ref], pruned: c.pruned} } diff --git a/tlb/hashmap.go b/tlb/hashmap.go index 1f29d019..aed32271 100644 --- a/tlb/hashmap.go +++ b/tlb/hashmap.go @@ -171,43 +171,44 @@ func countLeafs(keySize, leftKeySize int, c *boc.Cell) (int, error) { return 1, nil } -func ProveKeyInHashmap(cell *boc.Cell, key boc.BitString) ([]byte, error) { +func ProveKeyInHashmap[T any](prover *boc.MerkleProver, cell *boc.Cell, key boc.BitString) (T, []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 + cursor := prover.Cursor() + var t T + remaining := keySize for { var err error var size int - size, prefix, err = loadLabel(keySize, cell, prefix) + size, prefix, err = loadLabel(remaining, cell, prefix) if err != nil { - return nil, err + return t, nil, err } - if keySize <= size { + _ = prefix + if remaining <= size { break } if _, err = key.ReadBits(size); err != nil { - return nil, err + return t, nil, err } - isRight, err := key.ReadBit() if err != nil { - return nil, err + return t, nil, err + } + if err := prefix.WriteBit(isRight); err != nil { + return t, nil, err } - keySize = keySize - size - 1 + remaining = remaining - size - 1 next, err := cell.NextRef() if err != nil { - return nil, err + return t, nil, err } if isRight { cursor.Ref(0).Prune() next, err = cell.NextRef() if err != nil { - return nil, err + return t, nil, err } cursor = cursor.Ref(1) } else { @@ -215,8 +216,24 @@ func ProveKeyInHashmap(cell *boc.Cell, key boc.BitString) ([]byte, error) { cursor = cursor.Ref(0) } cell = next + cell.ResetCounters() + } + if err := Unmarshal(cell, &t); err != nil { + return t, nil, err + } + constructedKey, err := prefix.ReadBits(keySize) + if err != nil { + return t, nil, err } - return prover.CreateProof() + if constructedKey.ToFiftHex() != key.ToFiftHex() { + return t, nil, errors.New("key is not found") + } + proof, err := prover.CreateProof(cursor) + if err != nil { + return t, nil, err + } + return t, proof, nil + } func (h *Hashmap[keyT, T]) mapInner(keySize, leftKeySize int, c *boc.Cell, keyPrefix *boc.BitString, decoder *Decoder) error {