diff --git a/merkletree.go b/merkletree.go index 88e9f81..aaa4afa 100644 --- a/merkletree.go +++ b/merkletree.go @@ -135,17 +135,31 @@ func (t *MerkleTree) GenerateProofWithIndex(index uint64, height int) (*Proof, e // 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 } - hashes[dataIndex] = tmpProof.Hashes - indices[dataIndex] = tmpProof.Index + 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[i] = tmpProof.Hashes } // Step 2: combine the hashes across all proofs and highlight all calculated indices. diff --git a/multiproof_test.go b/multiproof_test.go index 4b295b3..5b32e07 100644 --- a/multiproof_test.go +++ b/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<>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 {