Merge pull request #12 from gpsanant/add-indexed-proof

Add indexed proof

merkletree.go

@@ -108,6 +108,19 @@ 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)
 
@@ -123,17 +136,31 @@ func (t *MerkleTree) GenerateProof(data []byte, height int) (*Proof, error) {
 
 // GenerateMultiProof generates the proof for multiple pieces of data.
 func (t *MerkleTree) GenerateMultiProof(data [][]byte) (*MultiProof, error) {
-	hashes := make([][][]byte, len(data))
 	indices := make([]uint64, len(data))
 
 	// Step 1: generate individual proofs.
 	for dataIndex := range data {
-		tmpProof, err := t.GenerateProof(data[dataIndex], 0)
+		index, err := t.indexOf(data[dataIndex])
+		if err != nil {
+			return nil, err
+		}
+		indices[dataIndex] = index
+	}
+
+	return t.GenerateMultiProofWithIndices(indices)
+}
+
+// GenerateMultiProof generates the proof for multiple pieces of data.
+func (t *MerkleTree) GenerateMultiProofWithIndices(indices []uint64) (*MultiProof, error) {
+	hashes := make([][][]byte, len(indices))
+
+	// Step 1: generate individual proofs.
+	for i, leafIndex := range indices {
+		tmpProof, err := t.GenerateProofWithIndex(leafIndex, 0)
 		if err != nil {
 			return nil, err
 		}
-		hashes[dataIndex] = tmpProof.Hashes
-		indices[dataIndex] = tmpProof.Index
+		hashes[i] = tmpProof.Hashes
 	}
 
 	// Step 2: combine the hashes across all proofs and highlight all calculated indices.

multiproof_test.go

@@ -24,6 +24,39 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func TestMultiProofWithIndices(t *testing.T) {
+	for i, test := range tests {
+		if test.createErr == nil {
+			tree, err := NewTree(
+				WithData(test.data),
+				WithHashType(test.hashType),
+				WithSalt(test.salt),
+				WithSorted(test.sorted),
+			)
+			assert.Nil(t, err, fmt.Sprintf("failed to create tree at test %d", i))
+
+			// Test proof for all combinations of data.
+			var proof *MultiProof
+			combinations := 1<<len(test.data) - 1
+			for j := 1; j <= combinations; j++ {
+				indices := make([]uint64, 0)
+				items := make([][]byte, 0)
+				for k := 0; k < len(test.data); k++ {
+					if (j>>k)&1 == 1 {
+						indices = append(indices, uint64(k))
+						items = append(items, test.data[k])
+					}
+				}
+				proof, err = tree.GenerateMultiProofWithIndices(indices)
+				assert.Nil(t, err, fmt.Sprintf("failed to create multiproof at test %d data %d", i, j))
+				proven, err := proof.Verify(items, tree.Root())
+				assert.Nil(t, err, fmt.Sprintf("error verifying multiproof at test %d data %d", i, j))
+				assert.True(t, proven, fmt.Sprintf("failed to verify multiproof at test %d data %d", i, j))
+			}
+		}
+	}
+}
+
 func TestMultiProof(t *testing.T) {
 	for i, test := range tests {
 		if test.createErr == nil {

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 {