1 files changed, 116 insertions, 0 deletions
diff --git a/block_pool.go b/block_pool.go
new file mode 100644
index 000000000..3225bdff2
--- /dev/null
+++ b/block_pool.go
@@ -0,0 +1,116 @@
+package eth
+
+import (
+	"math"
+	"math/big"
+	"sync"
+
+	"github.com/ethereum/eth-go/ethchain"
+	"github.com/ethereum/eth-go/ethutil"
+)
+
+type block struct {
+	peer  *Peer
+	block *ethchain.Block
+}
+
+type BlockPool struct {
+	mut sync.Mutex
+
+	eth *Ethereum
+
+	hashPool [][]byte
+	pool     map[string]*block
+
+	td *big.Int
+}
+
+func NewBlockPool(eth *Ethereum) *BlockPool {
+	return &BlockPool{
+		eth:  eth,
+		pool: make(map[string]*block),
+		td:   ethutil.Big0,
+	}
+}
+
+func (self *BlockPool) HasLatestHash() bool {
+	return self.pool[string(self.eth.BlockChain().CurrentBlock.Hash())] != nil
+}
+
+func (self *BlockPool) HasCommonHash(hash []byte) bool {
+	return self.eth.BlockChain().GetBlock(hash) != nil
+}
+
+func (self *BlockPool) AddHash(hash []byte) {
+	if self.pool[string(hash)] == nil {
+		self.pool[string(hash)] = &block{nil, nil}
+
+		self.hashPool = append([][]byte{hash}, self.hashPool...)
+	}
+}
+
+func (self *BlockPool) SetBlock(b *ethchain.Block) {
+	hash := string(b.Hash())
+
+	if self.pool[string(hash)] == nil {
+		self.pool[hash] = &block{nil, nil}
+	}
+
+	self.pool[hash].block = b
+}
+
+func (self *BlockPool) CheckLinkAndProcess(f func(block *ethchain.Block)) bool {
+	self.mut.Lock()
+	defer self.mut.Unlock()
+
+	if self.IsLinked() {
+		for i, hash := range self.hashPool {
+			block := self.pool[string(hash)].block
+			if block != nil {
+				f(block)
+
+				delete(self.pool, string(hash))
+			} else {
+				self.hashPool = self.hashPool[i:]
+
+				return false
+			}
+		}
+
+		return true
+	}
+
+	return false
+}
+
+func (self *BlockPool) IsLinked() bool {
+	if len(self.hashPool) == 0 {
+		return false
+	}
+
+	block := self.pool[string(self.hashPool[0])].block
+	if block != nil {
+		return self.eth.BlockChain().HasBlock(block.PrevHash)
+	}
+
+	return false
+}
+
+func (self *BlockPool) Take(amount int, peer *Peer) (hashes [][]byte) {
+	self.mut.Lock()
+	defer self.mut.Unlock()
+
+	num := int(math.Min(float64(amount), float64(len(self.pool))))
+	j := 0
+	for i := 0; i < len(self.hashPool) && j < num; i++ {
+		hash := string(self.hashPool[i])
+		if self.pool[hash].peer == nil || self.pool[hash].peer == peer {
+			self.pool[hash].peer = peer
+
+			hashes = append(hashes, self.hashPool[i])
+			j++
+		}
+	}
+
+	return
+}