core: Add NotarySet and DKGSet to cache (#176)

3 files changed, 94 insertions, 27 deletions

diff --git a/core/nodeset-cache.go b/core/nodeset-cache.go
index 49521ab..8cfe470 100644
--- a/core/nodeset-cache.go
+++ b/core/nodeset-cache.go
@@ -30,11 +30,17 @@ var (
 	ErrRoundNotReady = errors.New("round is not ready")
 )
 
+type sets struct {
+	nodeSet   *types.NodeSet
+	notarySet []map[types.NodeID]struct{}
+	dkgSet    map[types.NodeID]struct{}
+}
+
 // NodeSetCache caches node set information from governance contract.
 type NodeSetCache struct {
 	lock    sync.RWMutex
 	gov     Governance
-	rounds  map[uint64]*types.NodeSet
+	rounds  map[uint64]*sets
 	keyPool map[types.NodeID]*struct {
 		pubKey crypto.PublicKey
 		refCnt int
@@ -45,7 +51,7 @@ type NodeSetCache struct {
 func NewNodeSetCache(gov Governance) *NodeSetCache {
 	return &NodeSetCache{
 		gov:    gov,
-		rounds: make(map[uint64]*types.NodeSet),
+		rounds: make(map[uint64]*sets),
 		keyPool: make(map[types.NodeID]*struct {
 			pubKey crypto.PublicKey
 			refCnt int
@@ -63,7 +69,7 @@ func (cache *NodeSetCache) Exists(
 			return
 		}
 	}
-	_, exists = nIDs.IDs[nodeID]
+	_, exists = nIDs.nodeSet.IDs[nodeID]
 	return
 }
 
@@ -91,7 +97,50 @@ func (cache *NodeSetCache) GetNodeSet(
 			return
 		}
 	}
-	nIDs = IDs.Clone()
+	nIDs = IDs.nodeSet.Clone()
+	return
+}
+
+// GetNotarySet returns of notary set of this round.
+func (cache *NodeSetCache) GetNotarySet(
+	round uint64, chainID uint32) (map[types.NodeID]struct{}, error) {
+	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
+}
+
+// GetDKGSet returns of DKG set of this round.
+func (cache *NodeSetCache) GetDKGSet(
+	round uint64) (map[types.NodeID]struct{}, error) {
+	IDs, err := cache.getOrUpdate(round)
+	if err != nil {
+		return nil, err
+	}
+	return cache.cloneMap(IDs.dkgSet), nil
+}
+
+func (cache *NodeSetCache) cloneMap(
+	nIDs map[types.NodeID]struct{}) map[types.NodeID]struct{} {
+	nIDsCopy := make(map[types.NodeID]struct{}, len(nIDs))
+	for k := range nIDs {
+		nIDsCopy[k] = struct{}{}
+	}
+	return nIDsCopy
+}
+
+func (cache *NodeSetCache) getOrUpdate(round uint64) (nIDs *sets, err error) {
+	s, exists := cache.get(round)
+	if !exists {
+		if s, err = cache.update(round); err != nil {
+			return
+		}
+	}
+	nIDs = s
 	return
 }
 
@@ -100,7 +149,7 @@ func (cache *NodeSetCache) GetNodeSet(
 // This cache would maintain 10 rounds before the updated round and purge
 // rounds not in this range.
 func (cache *NodeSetCache) update(
-	round uint64) (nIDs *types.NodeSet, err error) {
+	round uint64) (nIDs *sets, err error) {
 
 	cache.lock.Lock()
 	defer cache.lock.Unlock()
@@ -113,10 +162,10 @@ func (cache *NodeSetCache) update(
 		return
 	}
 	// Cache new round.
-	nIDs = types.NewNodeSet()
+	nodeSet := types.NewNodeSet()
 	for _, key := range keySet {
 		nID := types.NewNodeID(key)
-		nIDs.Add(nID)
+		nodeSet.Add(nID)
 		if rec, exists := cache.keyPool[nID]; exists {
 			rec.refCnt++
 		} else {
@@ -126,13 +175,27 @@ func (cache *NodeSetCache) update(
 			}{key, 1}
 		}
 	}
+	cfg := cache.gov.Configuration(round)
+	crs := cache.gov.CRS(round)
+	nIDs = &sets{
+		nodeSet:   nodeSet,
+		notarySet: make([]map[types.NodeID]struct{}, cfg.NumChains),
+		dkgSet: nodeSet.GetSubSet(
+			cfg.NumDKGSet, types.NewDKGSetTarget(crs)),
+	}
+	for i := range nIDs.notarySet {
+		nIDs.notarySet[i] = nodeSet.GetSubSet(
+			cfg.NumNotarySet, types.NewNotarySetTarget(crs, uint32(i)))
+	}
+
 	cache.rounds[round] = nIDs
 	// Purge older rounds.
 	for rID, nIDs := range cache.rounds {
+		nodeSet := nIDs.nodeSet
 		if round-rID <= 5 {
 			continue
 		}
-		for nID := range nIDs.IDs {
+		for nID := range nodeSet.IDs {
 			rec := cache.keyPool[nID]
 			if rec.refCnt--; rec.refCnt == 0 {
 				delete(cache.keyPool, nID)
@@ -144,7 +207,7 @@ func (cache *NodeSetCache) update(
 }
 
 func (cache *NodeSetCache) get(
-	round uint64) (nIDs *types.NodeSet, exists bool) {
+	round uint64) (nIDs *sets, exists bool) {
 
 	cache.lock.RLock()
 	defer cache.lock.RUnlock()
diff --git a/core/nodeset-cache_test.go b/core/nodeset-cache_test.go
index d7e7dba..993b3ec 100644
--- a/core/nodeset-cache_test.go
+++ b/core/nodeset-cache_test.go
@@ -29,12 +29,19 @@ import (
 
 type testGov struct {
 	s       *NodeSetCacheTestSuite
+	crs     common.Hash
 	curKeys []crypto.PublicKey
 }
 
-func (g *testGov) Configuration(round uint64) (cfg *types.Config) { return }
-func (g *testGov) CRS(round uint64) (b common.Hash)               { return }
-func (g *testGov) ProposeCRS(uint64, []byte)                      {}
+func (g *testGov) Configuration(round uint64) (cfg *types.Config) {
+	return &types.Config{
+		NumNotarySet: 7,
+		NumDKGSet:    7,
+		NumChains:    4,
+	}
+}
+func (g *testGov) CRS(round uint64) (b common.Hash) { return g.crs }
+func (g *testGov) ProposeCRS(uint64, []byte)        {}
 func (g *testGov) NodeSet(round uint64) []crypto.PublicKey {
 	// Randomly generating keys, and check them for verification.
 	g.curKeys = []crypto.PublicKey{}
@@ -64,7 +71,10 @@ type NodeSetCacheTestSuite struct {
 
 func (s *NodeSetCacheTestSuite) TestBasicUsage() {
 	var (
-		gov   = &testGov{s: s}
+		gov = &testGov{
+			s:   s,
+			crs: common.NewRandomHash(),
+		}
 		cache = NewNodeSetCache(gov)
 		req   = s.Require()
 	)
@@ -88,6 +98,12 @@ func (s *NodeSetCacheTestSuite) TestBasicUsage() {
 	nodeSet0, err := cache.GetNodeSet(0)
 	req.NoError(err)
 	chk(cache, 0, nodeSet0.IDs)
+	notarySet, err := cache.GetNotarySet(0, 0)
+	req.NoError(err)
+	chk(cache, 0, notarySet)
+	dkgSet, err := cache.GetDKGSet(0)
+	req.NoError(err)
+	chk(cache, 0, dkgSet)
 	// Try to get round 1.
 	nodeSet1, err := cache.GetNodeSet(1)
 	req.NoError(err)
diff --git a/core/types/nodeset.go b/core/types/nodeset.go
index 83c64a7..eb83f19 100644
--- a/core/types/nodeset.go
+++ b/core/types/nodeset.go
@@ -38,7 +38,6 @@ type subSetTargetType byte
 
 const (
 	targetNotarySet subSetTargetType = iota
-	targetWitnessSet
 	targetDKGSet
 )
 
@@ -79,20 +78,9 @@ func NewNotarySetTarget(crs common.Hash, chainID uint32) SubSetTarget {
 	return newTarget(targetNotarySet, crs[:], binaryChainID)
 }
 
-// NewWitnessSetTarget  is the target for getting DKG Set.
-func NewWitnessSetTarget(crs common.Hash, round uint64) SubSetTarget {
-	binaryRound := make([]byte, 8)
-	binary.LittleEndian.PutUint64(binaryRound, round)
-
-	return newTarget(targetWitnessSet, crs[:], binaryRound)
-}
-
 // NewDKGSetTarget  is the target for getting DKG Set.
-func NewDKGSetTarget(crs common.Hash, round uint64) SubSetTarget {
-	binaryRound := make([]byte, 8)
-	binary.LittleEndian.PutUint64(binaryRound, round)
-
-	return newTarget(targetDKGSet, crs[:], binaryRound)
+func NewDKGSetTarget(crs common.Hash) SubSetTarget {
+	return newTarget(targetDKGSet, crs[:])
 }
 
 // Add a NodeID to the set.