Add batch proof extraction: sorting helps more than batching, will have to investigate page faults

1 files changed, 38 insertions, 1 deletions

diff --git a/merkle.go b/merkle.go
index 1ef55ce..9257c64 100644
--- a/merkle.go
+++ b/merkle.go
@@ -85,7 +85,7 @@ func (m *BFSMerkle) Build() []byte {
 func (m *BFSMerkle) Proof(id int64) [][]byte {
 	proof := make([][]byte, m.height)
 	proof[0] = make([]byte, hashSize)
-	m.Read(proof[0], id)
+	m.ReadAt(proof[0], id*hashSize)
 	for i := 1; id > 0; i++ {
 		proof[i] = make([]byte, hashSize)
 		if id&1 == 0 {
@@ -98,6 +98,43 @@ func (m *BFSMerkle) Proof(id int64) [][]byte {
 	return proof
 }
 
+func (m *BFSMerkle) Proofs(ids []int64) [][][]byte {
+	proofs := make([][][]byte, len(ids))
+	for i, id := range ids {
+		proofs[i] = m.Proof(id)
+	}
+	return proofs
+}
+
+func (m *BFSMerkle) BatchProofs(ids []int64) [][][]byte {
+	proofs := make([][][]byte, len(ids))
+	tids := make([]int64, len(ids))
+	copy(tids, ids)
+	var id int64
+
+	// pre-allocating proofs + first step
+	for i := range proofs {
+		proofs[i] = make([][]byte, m.height)
+		proofs[i][0] = make([]byte, hashSize)
+		m.ReadAt(proofs[i][0], tids[i]*hashSize)
+	}
+
+	// constructing proofs level by level
+	for height := int64(1); height < m.height; height++ {
+		for i := range proofs {
+			proofs[i][height] = make([]byte, hashSize)
+			id = tids[i]
+			if id&1 == 0 {
+				m.ReadAt(proofs[i][height], (id-1)*hashSize)
+			} else {
+				m.ReadAt(proofs[i][height], (id+1)*hashSize)
+			}
+			tids[i] = (id - 1) >> 1
+		}
+	}
+	return proofs
+}
+
 // nodes are stored in depth-first post order
 type PostMerkle struct {
 	height int64