add GenerateProofWithIndex

merkletree.go

@@ -108,6 +108,18 @@ func (t *MerkleTree) GenerateProof(data []byte, height int) (*Proof, error) {
 		return nil, err
 	}
 
+	return t.GenerateProofWithIndex(index, height)
+}
+
+// GenerateProofWithIndex generates the proof for the data at the given index. It is faster than GenerateProof() if the index is already known.
+// Height is the height of the pollard to verify the proof.  If using the Merkle root to verify this should be 0.
+// If the index is out of range this will return an error.
+// If the data is present in the tree this will return the hashes for each level in the tree and the index of the value in the tree.
+func (t *MerkleTree) GenerateProofWithIndex(index uint64, height int) (*Proof, error) {
+	if index >= uint64(len(t.Data)) {
+		return nil, errors.New("index out of range")
+	}
+
 	proofLen := int(math.Ceil(math.Log2(float64(len(t.Data))))) - height
 	hashes := make([][]byte, proofLen)
 

proof_test.go

@@ -21,6 +21,25 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+func TestProofWithIndex(t *testing.T) {
+	for i, test := range tests {
+		if test.createErr == nil {
+			tree, err := NewTree(
+				WithData(test.data),
+				WithHashType(test.hashType),
+			)
+			assert.Nil(t, err, fmt.Sprintf("failed to create tree at test %d", i))
+			for j, data := range test.data {
+				proof, err := tree.GenerateProofWithIndex(uint64(j), 0)
+				assert.Nil(t, err, fmt.Sprintf("failed to create proof at test %d data %d", i, j))
+				proven, err := VerifyProofUsing(data, false, proof, [][]byte{tree.Root()}, test.hashType)
+				assert.Nil(t, err, fmt.Sprintf("error verifying proof at test %d", i))
+				assert.True(t, proven, fmt.Sprintf("failed to verify proof at test %d data %d", i, j))
+			}
+		}
+	}
+}
+
 func TestProof(t *testing.T) {
 	for i, test := range tests {
 		if test.createErr == nil {
@@ -100,6 +119,20 @@ func TestMissingProof(t *testing.T) {
 	}
 }
 
+func TestProveInvalidIndex(t *testing.T) {
+	for i, test := range tests {
+		if test.createErr == nil {
+			tree, err := NewTree(
+				WithData(test.data),
+				WithHashType(test.hashType),
+			)
+			assert.Nil(t, err, fmt.Sprintf("failed to create tree at test %d", i))
+			_, err = tree.GenerateProofWithIndex(uint64(len(test.data)+i), 0)
+			assert.Equal(t, err.Error(), "index out of range")
+		}
+	}
+}
+
 func TestBadProof(t *testing.T) {
 	for i, test := range tests {
 		if test.createErr == nil && len(test.data) > 1 {