ryan feedback

api/proof.go

@@ -68,7 +68,8 @@ func (s *Server) GetProof(w http.ResponseWriter, r *http.Request) {
 	}
 
 	var (
-		p    = merkle.New(leaves).Proof(target)
+		mt   = merkle.New(leaves)
+		p    = mt.Proof(mt.Index(target))
 		phex = []hexutil.Bytes{}
 	)
 

api/tree.go

@@ -147,7 +147,7 @@ func (s *Server) CreateTree(w http.ResponseWriter, r *http.Request) {
 
 	var (
 		proofHashes = make([][]any, 0, len(leaves))
-		allProofs   = tree.Proofs()
+		allProofs   = tree.LeafProofs()
 	)
 
 	for _, p := range allProofs {

merkle/tree.go

@@ -74,6 +74,18 @@ func (t Tree) Root() []byte {
 	return t[len(t)-1][0]
 }
 
+// Returns the index of the target in the tree.
+// If the target is not in the tree, returns -1.
+func (t Tree) Index(target []byte) int {
+	ht := crypto.Keccak256(target)
+	for i, h := range t[0] {
+		if bytes.Equal(ht, h) {
+			return i
+		}
+	}
+	return -1
+}
+
 // Returns a list of hashes such that
 // cumulatively hashing the list pairwise
 // will yield the root hash of the tree. Example:
@@ -90,22 +102,7 @@ func (t Tree) Root() []byte {
 // [cd]
 //
 // The result of this func will be used in [Valid]
-func (t Tree) Proof(target []byte) [][]byte {
-	var (
-		ht    = crypto.Keccak256(target)
-		index int
-	)
-	for i, h := range t[0] {
-		if bytes.Equal(ht, h) {
-			index = i
-			break
-		}
-	}
-
-	return t.proofForEdge(index)
-}
-
-func (t Tree) proofForEdge(index int) [][]byte {
+func (t Tree) Proof(index int) [][]byte {
 	var proof [][]byte
 	for _, level := range t {
 		var i int
@@ -123,13 +120,13 @@ func (t Tree) proofForEdge(index int) [][]byte {
 	return proof
 }
 
-// Returns proofs for all edges in the tree.
+// Returns proofs for all leafs in the tree.
 // For details on how an individual proof is calculated, see [Tree.Proof].
-func (t Tree) Proofs() [][][]byte {
+func (t Tree) LeafProofs() [][][]byte {
 	proofs := make([][][]byte, len(t[0]))
 
 	for i := range t[0] {
-		proofs[i] = t.proofForEdge(i)
+		proofs[i] = t.Proof(i)
 	}
 
 	return proofs

merkle/tree_test.go

@@ -81,15 +81,47 @@ func TestProof(t *testing.T) {
 
 	for _, tc := range cases {
 		mt := New(tc.leaves)
-		for _, l := range tc.leaves {
-			pf := mt.Proof(l)
+		for i, l := range tc.leaves {
+			pf := mt.Proof(i)
 			if !Valid(mt.Root(), pf, l) {
 				t.Error("invalid proof")
 			}
 		}
 	}
 }
 
+func TestIndex(t *testing.T) {
+	leaves := [][]byte{
+		[]byte("a"),
+		[]byte("b"),
+		[]byte("c"),
+		[]byte("d"),
+		[]byte("e"),
+	}
+	mt := New(leaves)
+	for i, l := range leaves {
+		r := mt.Index(l)
+		if r != i {
+			t.Errorf("incorrect index, expected %d, got %d", i, r)
+		}
+	}
+}
+
+func TestMissingIndex(t *testing.T) {
+	mt := New([][]byte{
+		[]byte("a"),
+		[]byte("b"),
+		[]byte("c"),
+		[]byte("d"),
+		[]byte("e"),
+	})
+
+	pf := mt.Index([]byte("f"))
+	if pf != -1 {
+		t.Errorf("incorrect index, expected %d, got %d", -1, pf)
+	}
+}
+
 func BenchmarkNew(b *testing.B) {
 	var leaves [][]byte
 	for i := 0; i < 50000; i++ {
@@ -109,10 +141,10 @@ func BenchmarkProof(b *testing.B) {
 	mt := New(leaves)
 	for i := 0; i < b.N; i++ {
 		var eg errgroup.Group
-		for _, l := range leaves {
-			l := l
+		for i := range leaves {
+			i := i
 			eg.Go(func() error {
-				mt.Proof(l)
+				mt.Proof(i)
 				return nil
 			})
 		}
@@ -127,6 +159,6 @@ func BenchmarkProofs(b *testing.B) {
 	}
 	mt := New(leaves)
 	for i := 0; i < b.N; i++ {
-		mt.Proofs()
+		mt.LeafProofs()
 	}
 }