Parallelizing the DFS construction (huge performance gain!)

1 files changed, 111 insertions, 2 deletions

diff --git a/merkle.go b/merkle.go
index ec936c6..6135b5a 100644
--- a/merkle.go
+++ b/merkle.go
@@ -114,6 +114,7 @@ func (m *BFSMerkle) parallelbuild() []byte {
 		}
 		i++
 	}
+	wg.Wait()
 
 	for height := m.height - 1; height >= 1; height-- {
 		left = 1 << uint64(height-1)
@@ -123,7 +124,7 @@ func (m *BFSMerkle) parallelbuild() []byte {
 			go func(id int64) {
 				defer wg.Done()
 				buf := make([]byte, 2*hsize)
-				m.ReadAt(buf, (id*2+1)*hsize)
+				m.ReadAt(buf[:hsize], (id*2+1)*hsize)
 				m.ReadAt(buf[hsize:], (id*2+2)*hsize)
 				hash := sha3.Sum256(buf)
 				m.WriteAt(hash[:], id*hsize)
@@ -133,6 +134,7 @@ func (m *BFSMerkle) parallelbuild() []byte {
 			}
 			i++
 		}
+		wg.Wait()
 	}
 	buf := make([]byte, hsize)
 	m.ReadAt(buf, 0)
@@ -224,9 +226,17 @@ func (m *PostMerkle) Read(buf []byte, id int64) {
 	m.ReadAt(buf, pos*hashSize)
 }
 
+func (m *PostMerkle) Build(parallel bool) []byte {
+	if parallel {
+		return m.parallelbuild()
+	} else {
+		return m.build()
+	}
+}
+
 // Iterative post-order depth-first construction of the Merkle tree
 // disk access is optimal and forward
-func (m *PostMerkle) Build(parallel bool) []byte {
+func (m *PostMerkle) build() []byte {
 	size := m.Size()
 	h := sha3.New256()
 	hsize := int64(h.Size())
@@ -270,6 +280,105 @@ func (m *PostMerkle) Build(parallel bool) []byte {
 	return stack[0]
 }
 
+func (m *PostMerkle) parallelbuild() (r []byte) {
+	var wg sync.WaitGroup
+	wg.Add(4)
+	go func(id int64) {
+		defer wg.Done()
+		m.buildsubtree(id)
+	}(3)
+	go func(id int64) {
+		defer wg.Done()
+		m.buildsubtree(id)
+	}(4)
+	go func(id int64) {
+		defer wg.Done()
+		m.buildsubtree(id)
+	}(5)
+	go func(id int64) {
+		defer wg.Done()
+		m.buildsubtree(id)
+	}(6)
+	wg.Wait()
+	buf := make([]byte, hashSize)
+	h := sha3.New256()
+	m.Read(buf, 3)
+	h.Write(buf)
+	m.Read(buf, 4)
+	h.Write(buf)
+	h.Sum(buf[:0])
+	m.Put(1, buf)
+
+	h.Reset()
+	m.Read(buf, 5)
+	h.Write(buf)
+	m.Read(buf, 6)
+	h.Write(buf)
+	h.Sum(buf[:0])
+	m.Put(2, buf)
+
+	h.Reset()
+	m.Read(buf, 1)
+	h.Write(buf)
+	m.Read(buf, 2)
+	h.Write(buf)
+	h.Sum(buf[:0])
+	m.Put(0, buf)
+
+	return buf
+}
+
+func (m *PostMerkle) buildsubtree(id int64) []byte {
+	size := m.Size()
+	h := sha3.New256()
+	hsize := int64(h.Size())
+	height := m.height - Log(id+1)
+
+	// pre-allocating the hash stack
+	stack := make([][]byte, height)
+	for i := 0; i < len(stack); i++ {
+		stack[i] = make([]byte, hsize)
+	}
+
+	var cur int64 = id // current node (bfs id)
+	for cur < size/2 {
+		cur = (cur << 1) + 1
+	}
+	var count int64 = Post(size, m.height, cur) // post-order id of current node
+	var l = 0                                   // length of the stack
+
+	for cur > id {
+		if cur >= size/2 { // leaf node
+			l++
+			h.Reset()
+			binary.Write(h, binary.LittleEndian, cur)
+			h.Sum(stack[l-1][:0])
+			m.WriteAt(stack[l-1], count*hsize)
+
+			for cur&1 == 0 && cur > id {
+				// we just completed a right node, moving up to the parent
+				cur = (cur - 1) >> 1
+				count++
+				h.Reset()
+				h.Write(stack[l-2])
+				h.Write(stack[l-1])
+				l-- // pop two items, add one item
+				h.Sum(stack[l-1][:0])
+				m.WriteAt(stack[l-1], count*hsize)
+			}
+			if cur <= id {
+				break
+			}
+			// we just completed a left node, moving to its sibling
+			cur++
+			count++
+		} else {
+			cur = (cur << 1) + 1 // moving to the left child
+		}
+	}
+	return stack[0]
+}
+
 func (m *PostMerkle) Proof(id int64) [][]byte {
 	// traversing the tree from the root to the leaf reading the siblings along
 	// the path and filling the proof from right to left