X-Git-Url: http://git.osdn.net/view?p=bytom%2Fvapor.git;a=blobdiff_plain;f=protocol%2Fstate%2Fblockindex.go;h=e2e13548fa8e1244d82a3a6d8febd8360e5afbbd;hp=3262d8d856a4562875e652627f741ad00bdac82f;hb=c25a6856e885c46a1086d0328d7fc528ca19086c;hpb=2cf5801b2e693a45de9b51ec9aa9c1f787d57105

diff --git a/protocol/state/blockindex.go b/protocol/state/blockindex.go
index 3262d8d8..e2e13548 100644
--- a/protocol/state/blockindex.go
+++ b/protocol/state/blockindex.go
@@ -24,6 +24,7 @@ type BlockNode struct {
 	Version                uint64
 	Height                 uint64
 	Timestamp              uint64
+	BlockWitness           *common.BitMap
 	TransactionsMerkleRoot bc.Hash
 	TransactionStatusHash  bc.Hash
 }
@@ -42,6 +43,15 @@ func NewBlockNode(bh *types.BlockHeader, parent *BlockNode) (*BlockNode, error)
 		TransactionsMerkleRoot: bh.TransactionsMerkleRoot,
 		TransactionStatusHash:  bh.TransactionStatusHash,
 	}
+
+	node.BlockWitness = common.NewBitMap(uint32(len(bh.Witness)))
+	for i, witness := range bh.Witness {
+		if len(witness) != 0 {
+			if err := node.BlockWitness.Set(uint32(i)); err != nil {
+				return nil, err
+			}
+		}
+	}
 	return node, nil
 }
 
@@ -79,15 +89,17 @@ func (node *BlockNode) CalcPastMedianTime() uint64 {
 type BlockIndex struct {
 	sync.RWMutex
 
-	index     map[bc.Hash]*BlockNode
-	mainChain []*BlockNode
+	index       map[bc.Hash]*BlockNode
+	heightIndex map[uint64][]*BlockNode
+	mainChain   []*BlockNode
 }
 
 // NewBlockIndex will create a empty BlockIndex
 func NewBlockIndex() *BlockIndex {
 	return &BlockIndex{
-		index:     make(map[bc.Hash]*BlockNode),
-		mainChain: make([]*BlockNode, 0, approxNodesPerDay),
+		index:       make(map[bc.Hash]*BlockNode),
+		heightIndex: make(map[uint64][]*BlockNode),
+		mainChain:   make([]*BlockNode, 0, approxNodesPerDay),
 	}
 }
 
@@ -95,6 +107,7 @@ func NewBlockIndex() *BlockIndex {
 func (bi *BlockIndex) AddNode(node *BlockNode) {
 	bi.Lock()
 	bi.index[node.Hash] = node
+	bi.heightIndex[node.Height] = append(bi.heightIndex[node.Height], node)
 	bi.Unlock()
 }
 
@@ -145,6 +158,13 @@ func (bi *BlockIndex) NodeByHeight(height uint64) *BlockNode {
 	return bi.nodeByHeight(height)
 }
 
+// NodesByHeight return all block nodes at the specified height.
+func (bi *BlockIndex) NodesByHeight(height uint64) []*BlockNode {
+	bi.RLock()
+	defer bi.RUnlock()
+	return bi.heightIndex[height]
+}
+
 // SetMainChain will set the the mainChain array
 func (bi *BlockIndex) SetMainChain(node *BlockNode) {
 	bi.Lock()