1 files changed, 23 insertions, 32 deletions

diff --git a/vendor/github.com/dexon-foundation/dexon-consensus/core/utils/nodeset-cache.go b/vendor/github.com/dexon-foundation/dexon-consensus/core/utils/nodeset-cache.go
index 8a07c9d2b..83541283b 100644
--- a/vendor/github.com/dexon-foundation/dexon-consensus/core/utils/nodeset-cache.go
+++ b/vendor/github.com/dexon-foundation/dexon-consensus/core/utils/nodeset-cache.go
@@ -19,6 +19,7 @@ package utils
 
 import (
 	"errors"
+	"fmt"
 	"sync"
 
 	"github.com/dexon-foundation/dexon-consensus/common"
@@ -40,9 +41,9 @@ var (
 type sets struct {
 	crs        common.Hash
 	nodeSet    *types.NodeSet
-	notarySet  []map[types.NodeID]struct{}
+	notarySet  map[types.NodeID]struct{}
 	dkgSet     map[types.NodeID]struct{}
-	leaderNode []map[uint64]types.NodeID
+	leaderNode map[uint64]types.NodeID
 }
 
 // NodeSetCacheInterface interface specifies interface used by NodeSetCache.
@@ -112,30 +113,29 @@ func (cache *NodeSetCache) GetPublicKey(
 }
 
 // GetNodeSet returns IDs of nodes set of this round as map.
-func (cache *NodeSetCache) GetNodeSet(
-	round uint64) (nIDs *types.NodeSet, err error) {
-
+func (cache *NodeSetCache) GetNodeSet(round uint64) (*types.NodeSet, error) {
 	IDs, exists := cache.get(round)
 	if !exists {
+		var err error
 		if IDs, err = cache.update(round); err != nil {
-			return
+			return nil, err
 		}
 	}
-	nIDs = IDs.nodeSet.Clone()
-	return
+	return IDs.nodeSet.Clone(), nil
 }
 
 // GetNotarySet returns of notary set of this round.
+// TODO(mission): remove chainID parameter.
 func (cache *NodeSetCache) GetNotarySet(
 	round uint64, chainID uint32) (map[types.NodeID]struct{}, error) {
+	if chainID != 0 {
+		panic(fmt.Errorf("non-zero chainID found: %d", chainID))
+	}
 	IDs, err := cache.getOrUpdate(round)
 	if err != nil {
 		return nil, err
 	}
-	if chainID >= uint32(len(IDs.notarySet)) {
-		return nil, ErrInvalidChainID
-	}
-	return cache.cloneMap(IDs.notarySet[chainID]), nil
+	return cache.cloneMap(IDs.notarySet), nil
 }
 
 // GetDKGSet returns of DKG set of this round.
@@ -155,24 +155,21 @@ func (cache *NodeSetCache) GetLeaderNode(pos types.Position) (
 	if err != nil {
 		return types.NodeID{}, err
 	}
-	if pos.ChainID >= uint32(len(IDs.leaderNode)) {
-		return types.NodeID{}, ErrInvalidChainID
-	}
 	cache.lock.Lock()
 	defer cache.lock.Unlock()
-	if _, exist := IDs.leaderNode[pos.ChainID][pos.Height]; !exist {
-		notarySet := types.NewNodeSetFromMap(IDs.notarySet[pos.ChainID])
-		leader :=
-			notarySet.GetSubSet(1, types.NewNodeLeaderTarget(IDs.crs, pos))
+	if _, exist := IDs.leaderNode[pos.Height]; !exist {
+		notarySet := types.NewNodeSetFromMap(IDs.notarySet)
+		leader := notarySet.GetSubSet(1, types.NewNodeLeaderTarget(
+			IDs.crs, pos.Height))
 		if len(leader) != 1 {
 			panic(errors.New("length of leader is not one"))
 		}
 		for nID := range leader {
-			IDs.leaderNode[pos.ChainID][pos.Height] = nID
+			IDs.leaderNode[pos.Height] = nID
 			break
 		}
 	}
-	return IDs.leaderNode[pos.ChainID][pos.Height], nil
+	return IDs.leaderNode[pos.Height], nil
 }
 
 func (cache *NodeSetCache) cloneMap(
@@ -235,23 +232,17 @@ func (cache *NodeSetCache) update(
 		err = ErrConfigurationNotReady
 		return
 	}
+	nodesPerChain := cfg.RoundInterval / cfg.MinBlockInterval
 	nIDs = &sets{
 		crs:       crs,
 		nodeSet:   nodeSet,
-		notarySet: make([]map[types.NodeID]struct{}, cfg.NumChains),
+		notarySet: make(map[types.NodeID]struct{}),
 		dkgSet: nodeSet.GetSubSet(
 			int(cfg.DKGSetSize), types.NewDKGSetTarget(crs)),
-		leaderNode: make([]map[uint64]types.NodeID, cfg.NumChains),
-	}
-	for i := range nIDs.notarySet {
-		nIDs.notarySet[i] = nodeSet.GetSubSet(
-			int(cfg.NotarySetSize), types.NewNotarySetTarget(crs, uint32(i)))
+		leaderNode: make(map[uint64]types.NodeID, nodesPerChain),
 	}
-	nodesPerChain := cfg.RoundInterval / cfg.MinBlockInterval
-	for i := range nIDs.leaderNode {
-		nIDs.leaderNode[i] = make(map[uint64]types.NodeID, nodesPerChain)
-	}
-
+	nIDs.notarySet = nodeSet.GetSubSet(
+		int(cfg.NotarySetSize), types.NewNotarySetTarget(crs))
 	cache.rounds[round] = nIDs
 	// Purge older rounds.
 	for rID, nIDs := range cache.rounds {