Skip to content

Commit

Permalink
Merge pull request #16 from lazyledger/ismail/add_node_visitor
Browse files Browse the repository at this point in the history
Add possibility to traverse inner nodes
  • Loading branch information
liamsi authored Jan 18, 2021
2 parents 4b4c58a + 8340e76 commit 6e8a6a5
Show file tree
Hide file tree
Showing 2 changed files with 56 additions and 4 deletions.
25 changes: 21 additions & 4 deletions nmt.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,12 +21,16 @@ const (
var (
ErrMismatchedNamespaceSize = errors.New("mismatching namespace sizes")
ErrInvalidPushOrder = errors.New("pushed data has to be lexicographically ordered by namespace IDs")
noOp = func(hash []byte, children ...[]byte) {}
)

type NodeVisitorFn = func(hash []byte, children ...[]byte)

type Options struct {
InitialCapacity int
NamespaceIDSize namespace.IDSize
IgnoreMaxNamespace bool
NodeVisitor NodeVisitorFn
}

type Option func(*Options)
Expand Down Expand Up @@ -63,8 +67,15 @@ func IgnoreMaxNamespace(ignore bool) Option {
}
}

func NodeVisitor(nodeVisitorFn NodeVisitorFn) Option {
return func(opts *Options) {
opts.NodeVisitor = nodeVisitorFn
}
}

type NamespacedMerkleTree struct {
treeHasher internal.NmtHasher
visit NodeVisitorFn

// just cache stuff until we pass in a store and keep all nodes in there
// currently, only leaves and leafHashes are stored:
Expand All @@ -89,6 +100,7 @@ func New(h hash.Hash, setters ...Option) *NamespacedMerkleTree {
InitialCapacity: 128,
NamespaceIDSize: 8,
IgnoreMaxNamespace: true,
NodeVisitor: noOp,
}

for _, setter := range setters {
Expand All @@ -97,6 +109,7 @@ func New(h hash.Hash, setters ...Option) *NamespacedMerkleTree {
treeHasher := internal.NewNmtHasher(h, opts.NamespaceIDSize, opts.IgnoreMaxNamespace)
return &NamespacedMerkleTree{
treeHasher: treeHasher,
visit: opts.NodeVisitor,
leaves: make([][]byte, 0, opts.InitialCapacity),
leafHashes: make([][]byte, 0, opts.InitialCapacity),
namespaceRanges: make(map[string]leafRange),
Expand Down Expand Up @@ -260,19 +273,23 @@ func (n *NamespacedMerkleTree) Root() namespace.IntervalDigest {
func (n NamespacedMerkleTree) computeRoot(start, end int) []byte {
switch end - start {
case 0:
return n.treeHasher.EmptyRoot()
rootHash := n.treeHasher.EmptyRoot()
n.visit(rootHash)
return rootHash
case 1:
leafHash := n.treeHasher.HashLeaf(n.leaves[start])
if len(n.leafHashes) < len(n.leaves) {
n.leafHashes = append(n.leafHashes, leafHash)
}
n.visit(leafHash, n.leaves[start])
return leafHash
default:
k := getSplitPoint(end - start)
left := n.computeRoot(start, start+k)
right := n.computeRoot(start+k, end)

return n.treeHasher.HashNode(left, right)
hash := n.treeHasher.HashNode(left, right)
n.visit(hash, left, right)
return hash
}
}

Expand All @@ -283,7 +300,7 @@ func getSplitPoint(length int) int {
}
uLength := uint(length)
bitlen := bits.Len(uLength)
k := 1 << uint(bitlen-1)
k := 1 << (bitlen - 1)
if k == length {
k >>= 1
}
Expand Down
35 changes: 35 additions & 0 deletions nmt_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -471,6 +471,41 @@ func TestIgnoreMaxNamespace(t *testing.T) {
}
}

func TestNodeVisitor(t *testing.T) {
const (
numLeaves = 4
nidSize = 2
leafSize = 6
)
nodeHashes := make([][]byte, 0)
collectNodeHashes := func(hash []byte, _children ...[]byte) {
nodeHashes = append(nodeHashes, hash)
}

data := generateRandNamespacedRawData(numLeaves, nidSize, leafSize)
n := New(sha256.New(), NamespaceIDSize(nidSize), NodeVisitor(collectNodeHashes))
for j := 0; j < numLeaves; j++ {
if err := n.Push(data[j][:nidSize], data[j][nidSize:]); err != nil {
t.Errorf("err: %v", err)
}
}
root := n.Root()
last := nodeHashes[len(nodeHashes)-1]
if !bytes.Equal(root.Hash(), last[nidSize*2:]) {
t.Fatalf("last visited node's digest does not match the tree root's.")
}
if !bytes.Equal(root.Min(), last[:nidSize]) {
t.Fatalf("last visited node's min namespace does not match the tree root's.")
}
if !bytes.Equal(root.Max(), last[nidSize:nidSize*2]) {
t.Fatalf("last visited node's max namespace does not match the tree root's.")
}
t.Log("printing nodes in visiting order") // postorder DFS
for _, nodeHash := range nodeHashes {
t.Logf("|min: %x, max: %x, digest: %x...|\n", nodeHash[:nidSize], nodeHash[nidSize:nidSize*2], nodeHash[nidSize*2:nidSize*2+3])
}
}

func TestNamespacedMerkleTree_ProveErrors(t *testing.T) {
tests := []struct {
name string
Expand Down

0 comments on commit 6e8a6a5

Please sign in to comment.