Repository: celestiaorg/nmt
Pull request #40 (2021-05-23 16:44:12)

Reuse hash 
After reviewing https://github.com/lazyledger/lazyledger-core/pull/351 I've realized that #38 introduced an unnecessary alloc per leaf and inner node. 

This PR reverts the changes of #38, caches the root s.t. it only gets computed when necessary but keeps that test that aims to reproduce the issue in #37. (draft as I want to confirm this really does not re-introduce the issue in test withing ll-core, too https://github.com/lazyledger/lazyledger-core/commit/152202137455450f971df5c7294380714dcc7ac6)


Patch:
diff --git a/hasher.go b/hasher.go
index a85b90b..7091713 100644
--- a/hasher.go
+++ b/hasher.go
@@ -15,11 +15,9 @@ const (
 	DefaultNamespaceIDLen = 8
 )
 
-type NewHash = func() hash.Hash
-
 // defaultHasher uses sha256 as a base-hasher, 8 bytes
 // for the namespace IDs and ignores the maximum possible namespace.
-var defaultHasher = NewNmtHasher(sha256.New, DefaultNamespaceIDLen, true)
+var defaultHasher = NewNmtHasher(sha256.New(), DefaultNamespaceIDLen, true)
 
 // Sha256Namespace8FlaggedLeaf uses sha256 as a base-hasher, 8 bytes
 // for the namespace IDs and ignores the maximum possible namespace.
@@ -48,26 +46,19 @@ func Sha256Namespace8FlaggedLeaf(namespacedData []byte) []byte {
 // The output will also be of length 2*DefaultNamespaceIDLen+sha256.Size = 48 bytes.
 func Sha256Namespace8FlaggedInner(leftRight []byte) []byte {
 	const flagLen = DefaultNamespaceIDLen * 2
-	left := leftRight[:flagLen+sha256.Size]
-	right := leftRight[flagLen+sha256.Size:]
+	sha256Len := defaultHasher.Size()
+	left := leftRight[:flagLen+sha256Len]
+	right := leftRight[flagLen+sha256Len:]
 
 	return defaultHasher.HashNode(left, right)
 }
 
 type Hasher struct {
-	newHash      NewHash
+	hash.Hash
 	NamespaceLen namespace.IDSize
 
 	ignoreMaxNs      bool
 	precomputedMaxNs namespace.ID
-	size             int
-}
-
-func (n *Hasher) Size() int {
-	if n.size == 0 {
-		n.size = n.newHash().Size()
-	}
-	return n.size
 }
 
 func (n *Hasher) IsMaxNamespaceIDIgnored() bool {
@@ -78,9 +69,9 @@ func (n *Hasher) NamespaceSize() namespace.IDSize {
 	return n.NamespaceLen
 }
 
-func NewNmtHasher(baseHasher NewHash, nidLen namespace.IDSize, ignoreMaxNamespace bool) *Hasher {
+func NewNmtHasher(baseHasher hash.Hash, nidLen namespace.IDSize, ignoreMaxNamespace bool) *Hasher {
 	return &Hasher{
-		newHash:          baseHasher,
+		Hash:             baseHasher,
 		NamespaceLen:     nidLen,
 		ignoreMaxNs:      ignoreMaxNamespace,
 		precomputedMaxNs: bytes.Repeat([]byte{0xFF}, int(nidLen)),
@@ -88,10 +79,9 @@ func NewNmtHasher(baseHasher NewHash, nidLen namespace.IDSize, ignoreMaxNamespac
 }
 
 func (n *Hasher) EmptyRoot() []byte {
-	h := n.newHash()
 	emptyNs := bytes.Repeat([]byte{0}, int(n.NamespaceLen))
-	hash := h.Sum(nil)
-	digest := append(append(emptyNs, emptyNs...), hash...)
+	h := n.Sum(nil)
+	digest := append(append(emptyNs, emptyNs...), h...)
 
 	return digest
 }
@@ -105,7 +95,8 @@ func (n *Hasher) EmptyRoot() []byte {
 //Note that for leaves minNs = maxNs = ns(leaf) = leaf[:NamespaceLen].
 //nolint:errcheck
 func (n *Hasher) HashLeaf(leaf []byte) []byte {
-	h := n.newHash()
+	h := n.Hash
+	h.Reset()
 
 	nID := leaf[:n.NamespaceLen]
 	data := leaf[n.NamespaceLen:]
@@ -120,7 +111,8 @@ func (n *Hasher) HashLeaf(leaf []byte) []byte {
 // left and right child node bytes, including their respective min and max namespace IDs:
 // left = left.Min() || left.Max() || l.Hash().
 func (n *Hasher) HashNode(l, r []byte) []byte {
-	h := n.newHash()
+	h := n.Hash
+	h.Reset()
 
 	// the actual hash result of the children got extended (or flagged) by their
 	// children's minNs || maxNs; hence the flagLen = 2 * NamespaceLen:
diff --git a/nmt.go b/nmt.go
index a3abe7a..b0316d1 100644
--- a/nmt.go
+++ b/nmt.go
@@ -5,6 +5,7 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	"hash"
 	"math/bits"
 
 	"github.com/celestiaorg/merkletree"
@@ -82,13 +83,16 @@ type NamespacedMerkleTree struct {
 	namespaceRanges map[string]leafRange
 	minNID          namespace.ID
 	maxNID          namespace.ID
+
+	// cache the root
+	rawRoot []byte
 }
 
 // New initializes a namespaced Merkle tree using the given base hash function
 // and for the given namespace size (number of bytes).
 // If the namespace size is 0 this corresponds to a regular non-namespaced
 // Merkle tree.
-func New(h NewHash, setters ...Option) *NamespacedMerkleTree {
+func New(h hash.Hash, setters ...Option) *NamespacedMerkleTree {
 	// default options:
 	opts := &Options{
 		InitialCapacity:    128,
@@ -255,13 +259,17 @@ func (n *NamespacedMerkleTree) Push(namespacedData namespace.PrefixedData) error
 	n.leaves = append(n.leaves, namespacedData)
 	n.updateNamespaceRanges()
 	n.updateMinMaxID(nID)
+	n.rawRoot = nil
 	return nil
 }
 
 // Return the namespaced Merkle Tree's root together with the
 // min. and max. namespace ID.
 func (n *NamespacedMerkleTree) Root() namespace.IntervalDigest {
-	return mustIntervalDigestFromBytes(n.NamespaceSize(), n.computeRoot(0, len(n.leaves)))
+	if n.rawRoot == nil {
+		n.rawRoot = n.computeRoot(0, len(n.leaves))
+	}
+	return mustIntervalDigestFromBytes(n.NamespaceSize(), n.rawRoot)
 }
 
 func (n NamespacedMerkleTree) computeRoot(start, end int) []byte {
diff --git a/proof.go b/proof.go
index 893fba0..c99b6f1 100644
--- a/proof.go
+++ b/proof.go
@@ -2,6 +2,7 @@ package nmt
 
 import (
 	"bytes"
+	"hash"
 	"math"
 	"math/bits"
 
@@ -95,7 +96,7 @@ func NewAbsenceProof(proofStart, proofEnd int, proofNodes [][]byte, leafHash []b
 // VerifyNamespace verifies a whole namespace, i.e. it verifies inclusion of
 // the provided data in the tree. Additionally, it verifies that the namespace
 // is complete and no leaf of that namespace was left out in the proof.
-func (proof Proof) VerifyNamespace(h NewHash, nID namespace.ID, data [][]byte, root namespace.IntervalDigest) bool {
+func (proof Proof) VerifyNamespace(h hash.Hash, nID namespace.ID, data [][]byte, root namespace.IntervalDigest) bool {
 	nth := NewNmtHasher(h, nID.Size(), proof.isMaxNamespaceIDIgnored)
 	if nID.Size() != root.Min.Size() || nID.Size() != root.Max.Size() {
 		// conflicting namespace sizes
@@ -203,7 +204,7 @@ func (proof Proof) verifyLeafHashes(nth *Hasher, verifyCompleteness bool, nID na
 	return bytes.Equal(tree.Root(), root.Bytes())
 }
 
-func (proof Proof) VerifyInclusion(h NewHash, nid namespace.ID, data []byte, root namespace.IntervalDigest) bool {
+func (proof Proof) VerifyInclusion(h hash.Hash, nid namespace.ID, data []byte, root namespace.IntervalDigest) bool {
 	nth := NewNmtHasher(h, nid.Size(), proof.isMaxNamespaceIDIgnored)
 	leafData := append(nid, data...)
 	return proof.verifyLeafHashes(nth, false, nid, [][]byte{nth.HashLeaf(leafData)}, root)


Test patch:
diff --git a/fuzz_test.go b/fuzz_test.go
index 13ee277..e7fe8af 100644
--- a/fuzz_test.go
+++ b/fuzz_test.go
@@ -30,8 +30,8 @@ func TestFuzzProveVerifyNameSpace(t *testing.T) {
 	for _, size := range testNsSizes {
 		nidDataMap, sortedKeys := makeRandDataAndSortedKeys(size, minNumberOfNamespaces, maxNumberOfNamespaces, minElementsPerNamespace, maxElementsPerNamespace, emptyNamespaceProbability)
 		t.Logf("Generated %v namespaces for size: %v ...", len(nidDataMap), size)
-		newHash := sha256.New
-		tree := nmt.New(newHash, nmt.NamespaceIDSize(int(size)))
+		hash := sha256.New()
+		tree := nmt.New(hash, nmt.NamespaceIDSize(int(size)))
 
 		// push data in order:
 		for _, ns := range sortedKeys {
@@ -55,7 +55,7 @@ func TestFuzzProveVerifyNameSpace(t *testing.T) {
 				t.Fatalf("error on ProveNamespace(%x): %v", ns, err)
 			}
 
-			if ok := proof.VerifyNamespace(newHash, nid, data, treeRoot); !ok {
+			if ok := proof.VerifyNamespace(hash, nid, data, treeRoot); !ok {
 				t.Fatalf("expected VerifyNamespace() == true")
 			}
 
@@ -72,7 +72,7 @@ func TestFuzzProveVerifyNameSpace(t *testing.T) {
 				if err != nil {
 					t.Fatalf("error on Prove(%v): %v", i, err)
 				}
-				if ok := singleItemProof.VerifyInclusion(newHash, data[i][:size], data[i][size:], treeRoot); !ok {
+				if ok := singleItemProof.VerifyInclusion(hash, data[i][:size], data[i][size:], treeRoot); !ok {
 					t.Fatalf("expected VerifyInclusion() == true; data = %#v; proof = %#v", data[i], singleItemProof)
 				}
 				leafIdx++
@@ -89,7 +89,7 @@ func TestFuzzProveVerifyNameSpace(t *testing.T) {
 
 			if len(data) != 0 {
 				emptyProof := nmt.NewEmptyRangeProof(false)
-				if emptyProof.VerifyNamespace(newHash, nid, data, treeRoot) {
+				if emptyProof.VerifyNamespace(hash, nid, data, treeRoot) {
 					t.Fatalf("empty range proof on non-empty data verified to true")
 				}
 				nonEmptyNsCount++
diff --git a/hasher_test.go b/hasher_test.go
index a209720..f805314 100644
--- a/hasher_test.go
+++ b/hasher_test.go
@@ -38,7 +38,7 @@ func Test_namespacedTreeHasher_HashLeaf(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			n := NewNmtHasher(sha256.New, tt.nsLen, false)
+			n := NewNmtHasher(sha256.New(), tt.nsLen, false)
 			if got := n.HashLeaf(tt.leaf); !reflect.DeepEqual(got, tt.want) {
 				t.Errorf("HashLeaf() = %v, want %v", got, tt.want)
 			}
@@ -91,7 +91,7 @@ func Test_namespacedTreeHasher_HashNode(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			n := NewNmtHasher(sha256.New, tt.nidLen, false)
+			n := NewNmtHasher(sha256.New(), tt.nidLen, false)
 			if got := n.HashNode(tt.children.l, tt.children.r); !reflect.DeepEqual(got, tt.want) {
 				t.Errorf("HashNode() = %v, want %v", got, tt.want)
 			}
diff --git a/nmt_test.go b/nmt_test.go
index ba74284..6763b35 100644
--- a/nmt_test.go
+++ b/nmt_test.go
@@ -46,7 +46,7 @@ func ExampleNamespacedMerkleTree() {
 		append(namespace.ID{1}, []byte("leaf_3")...)}
 	// Init a tree with the namespace size as well as
 	// the underlying hash function:
-	tree := New(sha256.New, NamespaceIDSize(nidSize))
+	tree := New(sha256.New(), NamespaceIDSize(nidSize))
 	for _, d := range data {
 		if err := tree.Push(d); err != nil {
 			panic(fmt.Sprintf("unexpected error: %v", err))
@@ -74,11 +74,11 @@ func ExampleNamespacedMerkleTree() {
 		append(namespace.ID{0}, []byte("leaf_1")...),
 	}
 
-	if proof.VerifyNamespace(sha256.New, namespace.ID{0}, leafs, root) {
+	if proof.VerifyNamespace(sha256.New(), namespace.ID{0}, leafs, root) {
 		fmt.Printf("Successfully verified namespace: %x\n", namespace.ID{0})
 	}
 
-	if proof.VerifyNamespace(sha256.New, namespace.ID{2}, leafs, root) {
+	if proof.VerifyNamespace(sha256.New(), namespace.ID{2}, leafs, root) {
 		panic(fmt.Sprintf("Proof for namespace %x, passed for namespace: %x\n", namespace.ID{0}, namespace.ID{2}))
 	}
 	// Output:
@@ -108,7 +108,7 @@ func TestNamespacedMerkleTree_Push(t *testing.T) {
 		// too short though, it can't extract the namespace and hence will complain:
 		{"push with wrong namespace size: Err", []byte{1, 1}, true},
 	}
-	n := New(sha256.New, NamespaceIDSize(3))
+	n := New(sha256.New(), NamespaceIDSize(3))
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			if err := n.Push(tt.data); (err != nil) != tt.wantErr {
@@ -148,7 +148,7 @@ func TestNamespacedMerkleTreeRoot(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			n := New(sha256.New, NamespaceIDSize(tt.nidLen))
+			n := New(sha256.New(), NamespaceIDSize(tt.nidLen))
 			for _, d := range tt.pushedData {
 				if err := n.Push(namespace.PrefixedData(append(d.ID, d.Data...))); err != nil {
 					t.Errorf("Push() error = %v, expected no error", err)
@@ -233,7 +233,7 @@ func TestNamespacedMerkleTree_ProveNamespace_Ranges_And_Verify(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			n := New(sha256.New, NamespaceIDSize(tt.nidLen))
+			n := New(sha256.New(), NamespaceIDSize(tt.nidLen))
 			for _, d := range tt.pushData {
 				err := n.Push(namespace.PrefixedData(append(d.ID, d.Data...)))
 				if err != nil {
@@ -260,7 +260,7 @@ func TestNamespacedMerkleTree_ProveNamespace_Ranges_And_Verify(t *testing.T) {
 
 			// Verification round-trip should always pass:
 			gotGetLeaves := n.Get(tt.proveNID)
-			gotChecksOut := gotProof.VerifyNamespace(sha256.New, tt.proveNID, gotGetLeaves, n.Root())
+			gotChecksOut := gotProof.VerifyNamespace(sha256.New(), tt.proveNID, gotGetLeaves, n.Root())
 			if !gotChecksOut {
 				t.Errorf("Proof.VerifyNamespace() gotChecksOut: %v, want: true", gotChecksOut)
 			}
@@ -272,7 +272,7 @@ func TestNamespacedMerkleTree_ProveNamespace_Ranges_And_Verify(t *testing.T) {
 					if err != nil {
 						t.Fatalf("unexpected error on Prove(): %v", err)
 					}
-					gotChecksOut := gotSingleProof.VerifyInclusion(sha256.New, data.ID, data.Data, n.Root())
+					gotChecksOut := gotSingleProof.VerifyInclusion(sha256.New(), data.ID, data.Data, n.Root())
 					if !gotChecksOut {
 						t.Errorf("Proof.VerifyInclusion() gotChecksOut: %v, want: true", gotChecksOut)
 					}
@@ -297,7 +297,7 @@ func TestNamespacedMerkleTree_ProveNamespace_Ranges_And_Verify(t *testing.T) {
 
 func TestIgnoreMaxNamespace(t *testing.T) {
 	var (
-		hash      = sha256.New
+		hash      = sha256.New()
 		nidSize   = 8
 		minNID    = []byte{0, 0, 0, 0, 0, 0, 0, 0}
 		secondNID = []byte{0, 0, 0, 0, 0, 0, 0, 1}
@@ -485,7 +485,7 @@ func TestNodeVisitor(t *testing.T) {
 	}
 
 	data := generateRandNamespacedRawData(numLeaves, nidSize, leafSize)
-	n := New(sha256.New, NamespaceIDSize(nidSize), NodeVisitor(collectNodeHashes))
+	n := New(sha256.New(), NamespaceIDSize(nidSize), NodeVisitor(collectNodeHashes))
 	for j := 0; j < numLeaves; j++ {
 		if err := n.Push(data[j]); err != nil {
 			t.Errorf("err: %v", err)
@@ -522,7 +522,7 @@ func TestNamespacedMerkleTree_ProveErrors(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			n := New(sha256.New, NamespaceIDSize(tt.nidLen), InitialCapacity(len(tt.pushData)))
+			n := New(sha256.New(), NamespaceIDSize(tt.nidLen), InitialCapacity(len(tt.pushData)))
 			for _, d := range tt.pushData {
 				err := n.Push(namespace.PrefixedData(append(d.ID, d.Data...)))
 				if err != nil {
@@ -566,7 +566,7 @@ func TestNamespacedMerkleTree_calculateAbsenceIndex_Panic(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			n := New(sha256.New, NamespaceIDSize(2))
+			n := New(sha256.New(), NamespaceIDSize(2))
 			shouldPanic(t,
 				func() { n.calculateAbsenceIndex(tt.nID) })
 		})
@@ -575,13 +575,13 @@ func TestNamespacedMerkleTree_calculateAbsenceIndex_Panic(t *testing.T) {
 
 func TestInvalidOptions(t *testing.T) {
 	shouldPanic(t, func() {
-		_ = New(sha256.New, InitialCapacity(-1))
+		_ = New(sha256.New(), InitialCapacity(-1))
 	})
 	shouldPanic(t, func() {
-		_ = New(sha256.New, NamespaceIDSize(-1))
+		_ = New(sha256.New(), NamespaceIDSize(-1))
 	})
 	shouldPanic(t, func() {
-		_ = New(sha256.New, NamespaceIDSize(namespace.IDMaxSize+1))
+		_ = New(sha256.New(), NamespaceIDSize(namespace.IDMaxSize+1))
 	})
 }
 
@@ -603,7 +603,7 @@ func BenchmarkComputeRoot(b *testing.B) {
 		b.ResetTimer()
 		b.Run(tt.name, func(b *testing.B) {
 			for i := 0; i < b.N; i++ {
-				n := New(sha256.New)
+				n := New(sha256.New())
 				for j := 0; j < tt.numLeaves; j++ {
 					if err := n.Push(data[j]); err != nil {
 						b.Errorf("err: %v", err)
@@ -618,7 +618,7 @@ func BenchmarkComputeRoot(b *testing.B) {
 
 func Test_Root_RaceCondition(t *testing.T) {
 	// this is very similar to: https://github.com/HuobiRDCenter/huobi_Golang/pull/9
-	tree := New(sha256.New)
+	tree := New(sha256.New())
 	_ = tree.Push([]byte("some data is good enough here"))
 	numRoutines := 200
 	wg := sync.WaitGroup{}
diff --git a/proof_test.go b/proof_test.go
index a451151..98e7708 100644
--- a/proof_test.go
+++ b/proof_test.go
@@ -13,7 +13,7 @@ import (
 func TestProof_VerifyNamespace_False(t *testing.T) {
 	const testNidLen = 3
 
-	n := New(sha256.New, NamespaceIDSize(testNidLen))
+	n := New(sha256.New(), NamespaceIDSize(testNidLen))
 	data := append(append([]namespaceDataPair{
 		newNamespaceDataPair([]byte{0, 0, 0}, []byte("first leaf"))},
 		generateLeafData(testNidLen, 0, 9, []byte("data"))...,
@@ -73,7 +73,7 @@ func TestProof_VerifyNamespace_False(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got := tt.proof.VerifyNamespace(sha256.New, tt.args.nID, tt.args.data, tt.args.root)
+			got := tt.proof.VerifyNamespace(sha256.New(), tt.args.nID, tt.args.data, tt.args.root)
 			if got != tt.want {
 				t.Errorf("VerifyNamespace() got = %v, want %v", got, tt.want)
 			}