diff --git a/merkletree.go b/merkletree.go index cec53ad..48a3c7d 100644 --- a/merkletree.go +++ b/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. 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 { diff --git a/proof_test.go b/proof_test.go index 8719247..2d72f29 100644 --- a/proof_test.go +++ b/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 {