feat: add custom hasher option

celestiaorg · Jun 26, 2023 · 6dc07a6 · 6dc07a6
1 parent 2f7f25e
commit 6dc07a6
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/nmt.go b/nmt.go
@@ -83,6 +83,13 @@ func NodeVisitor(nodeVisitorFn NodeVisitorFn) Option {
 	}
 }
 
+// CustomHasher replaces the default hasher.
+func CustomHasher(h Hasher) Option {
+	return func(o *Options) {
+		o.Hasher = h
+	}
+}
+
 type NamespacedMerkleTree struct {
 	treeHasher Hasher
 	visit      NodeVisitorFn

diff --git a/nmt_test.go b/nmt_test.go
@@ -553,6 +553,19 @@ func TestNodeVisitor(t *testing.T) {
 	}
 }
 
+func TestCustomHasher(t *testing.T) {
+	type customHasher struct {
+		*NmtHasher
+	}
+
+	h := customHasher{NewNmtHasher(sha256.New(), namespace.IDSize(8), true)}
+
+	tree := New(sha256.New(), NamespaceIDSize(8), IgnoreMaxNamespace(true), CustomHasher(h))
+
+	_, ok := tree.treeHasher.(customHasher)
+	require.True(t, ok)
+}
+
 func TestNamespacedMerkleTree_ProveErrors(t *testing.T) {
 	tests := []struct {
 		name     string