core: vm: improve MPK reverse lookup performance

By storing array offset in the original DKGMasterPublicKeyProposed map,
we can locate MPK by proposer ID easily.

3 files changed, 92 insertions, 82 deletions

diff --git a/core/vm/oracle_contract_abi.go b/core/vm/oracle_contract_abi.go
index 799d0f28f..bd99c0cc7 100644
--- a/core/vm/oracle_contract_abi.go
+++ b/core/vm/oracle_contract_abi.go
@@ -38,25 +38,6 @@ const GovernanceABIJSON = `
   },
   {
     "constant": true,
-    "inputs": [
-      {
-        "name": "",
-        "type": "bytes32"
-      }
-    ],
-    "name": "dkgMasterPublicKeyProposed",
-    "outputs": [
-      {
-        "name": "",
-        "type": "bool"
-      }
-    ],
-    "payable": false,
-    "stateMutability": "view",
-    "type": "function"
-  },
-  {
-    "constant": true,
     "inputs": [],
     "name": "totalSupply",
     "outputs": [
@@ -368,6 +349,25 @@ const GovernanceABIJSON = `
     "inputs": [
       {
         "name": "",
+        "type": "bytes32"
+      }
+    ],
+    "name": "dkgMasterPublicKeyOffset",
+    "outputs": [
+      {
+        "name": "",
+        "type": "uint256"
+      }
+    ],
+    "payable": false,
+    "stateMutability": "view",
+    "type": "function"
+  },
+  {
+    "constant": true,
+    "inputs": [
+      {
+        "name": "",
         "type": "address"
       }
     ],
diff --git a/core/vm/oracle_contracts.go b/core/vm/oracle_contracts.go
index 06b9c2a89..4878fd325 100644
--- a/core/vm/oracle_contracts.go
+++ b/core/vm/oracle_contracts.go
@@ -71,7 +71,7 @@ const (
 	dkgRoundLoc
 	dkgResetCountLoc
 	dkgMasterPublicKeysLoc
-	dkgMasterPublicKeyProposedLoc
+	dkgMasterPublicKeyOffsetLoc
 	dkgComplaintsLoc
 	dkgComplaintsProposedLoc
 	dkgReadyLoc
@@ -603,6 +603,20 @@ func (s *GovernanceState) IncDKGResetCount(round *big.Int) {
 func (s *GovernanceState) LenDKGMasterPublicKeys() *big.Int {
 	return s.getStateBigInt(big.NewInt(dkgMasterPublicKeysLoc))
 }
+func (s *GovernanceState) DKGMasterPublicKey(offset *big.Int) []byte {
+	loc := big.NewInt(dkgMasterPublicKeysLoc)
+	dataLoc := s.getSlotLoc(loc)
+	elementLoc := new(big.Int).Add(dataLoc, offset)
+	return s.readBytes(elementLoc)
+}
+func (s *GovernanceState) DKGMasterPublicKeyItem(offset *big.Int) *dkgTypes.MasterPublicKey {
+	element := s.DKGMasterPublicKey(offset)
+	x := new(dkgTypes.MasterPublicKey)
+	if err := rlp.DecodeBytes(element, x); err != nil {
+		panic(err)
+	}
+	return x
+}
 func (s *GovernanceState) DKGMasterPublicKeys() [][]byte {
 	return s.read1DByteArray(big.NewInt(dkgMasterPublicKeysLoc))
 }
@@ -621,44 +635,22 @@ func (s *GovernanceState) DKGMasterPublicKeyItems() []*dkgTypes.MasterPublicKey
 	}
 	return dkgMasterPKs
 }
-func (s *GovernanceState) GetDKGMasterPublicKeyByProposerID(
-	proposerID coreTypes.NodeID) (*dkgTypes.MasterPublicKey, error) {
-
-	for _, mpk := range s.DKGMasterPublicKeys() {
-		x := new(dkgTypes.MasterPublicKey)
-		if err := rlp.DecodeBytes(mpk, x); err != nil {
-			panic(err)
-		}
-		if x.ProposerID.Equal(proposerID) {
-			return x, nil
-		}
-	}
-	return nil, errors.New("not found")
-}
 func (s *GovernanceState) ClearDKGMasterPublicKeys() {
 	s.erase1DByteArray(big.NewInt(dkgMasterPublicKeysLoc))
 }
 
-// mapping(bytes32 => bool) public dkgMasterPublicKeyProposed;
-func (s *GovernanceState) DKGMasterPublicKeyProposed(id Bytes32) bool {
-	loc := s.getMapLoc(big.NewInt(dkgMasterPublicKeyProposedLoc), id[:])
-	return s.getStateBigInt(loc).Cmp(big.NewInt(0)) > 0
+// mapping(bytes32 => uint256) public dkgMasterPublicKeyOffset;
+func (s *GovernanceState) DKGMasterPublicKeyOffset(id Bytes32) *big.Int {
+	loc := s.getMapLoc(big.NewInt(dkgMasterPublicKeyOffsetLoc), id[:])
+	return new(big.Int).Sub(s.getStateBigInt(loc), big.NewInt(1))
 }
-func (s *GovernanceState) PutDKGMasterPublicKeyProposed(id Bytes32, status bool) {
-	loc := s.getMapLoc(big.NewInt(dkgMasterPublicKeyProposedLoc), id[:])
-	val := big.NewInt(0)
-	if status {
-		val = big.NewInt(1)
-	}
-	s.setStateBigInt(loc, val)
+func (s *GovernanceState) PutDKGMasterPublicKeyOffset(id Bytes32, offset *big.Int) {
+	loc := s.getMapLoc(big.NewInt(dkgMasterPublicKeyOffsetLoc), id[:])
+	s.setStateBigInt(loc, new(big.Int).Add(offset, big.NewInt(1)))
 }
-func (s *GovernanceState) ClearDKGMasterPublicKeyProposed() {
-	for _, mpk := range s.DKGMasterPublicKeys() {
-		x := new(dkgTypes.MasterPublicKey)
-		if err := rlp.DecodeBytes(mpk, x); err != nil {
-			panic(err)
-		}
-		s.PutDKGMasterPublicKeyProposed(getDKGMasterPublicKeyID(x), false)
+func (s *GovernanceState) ClearDKGMasterPublicKeyOffset() {
+	for _, mpk := range s.DKGMasterPublicKeyItems() {
+		s.PutDKGMasterPublicKeyOffset(getDKGMasterPublicKeyID(mpk), big.NewInt(-1))
 	}
 }
 
@@ -666,6 +658,12 @@ func (s *GovernanceState) ClearDKGMasterPublicKeyProposed() {
 func (s *GovernanceState) LenDKGComplaints() *big.Int {
 	return s.getStateBigInt(big.NewInt(dkgComplaintsLoc))
 }
+func (s *GovernanceState) DKGComplaint(offset *big.Int) []byte {
+	loc := big.NewInt(dkgComplaintsLoc)
+	dataLoc := s.getSlotLoc(loc)
+	elementLoc := new(big.Int).Add(dataLoc, offset)
+	return s.readBytes(elementLoc)
+}
 func (s *GovernanceState) DKGComplaints() [][]byte {
 	return s.read1DByteArray(big.NewInt(dkgComplaintsLoc))
 }
@@ -701,12 +699,8 @@ func (s *GovernanceState) PutDKGComplaintProposed(id Bytes32, status bool) {
 	s.setStateBigInt(loc, val)
 }
 func (s *GovernanceState) ClearDKGComplaintProposed() {
-	for _, comp := range s.DKGComplaints() {
-		x := new(dkgTypes.Complaint)
-		if err := rlp.DecodeBytes(comp, x); err != nil {
-			panic(err)
-		}
-		s.PutDKGComplaintProposed(getDKGComplaintID(x), false)
+	for _, comp := range s.DKGComplaintItems() {
+		s.PutDKGComplaintProposed(getDKGComplaintID(comp), false)
 	}
 }
 
@@ -1267,10 +1261,8 @@ func (c *defaultCoreDKGUtils) NewGroupPublicKey(
 		}
 	}
 
-	// Prepare DKGMasterPublicKeys.
 	mpks := state.DKGMasterPublicKeyItems()
 	comps := state.DKGComplaintItems()
-
 	gpk, err = dkgTypes.NewGroupPublicKey(round.Uint64(), mpks, comps, threshold)
 	if err != nil {
 		return nil, err
@@ -1357,7 +1349,7 @@ func (g *GovernanceContract) inNotarySet(round *big.Int, nodeID coreTypes.NodeID
 
 func (g *GovernanceContract) clearDKG() {
 	dkgSet := g.getNotarySet(g.state.DKGRound())
-	g.state.ClearDKGMasterPublicKeyProposed()
+	g.state.ClearDKGMasterPublicKeyOffset()
 	g.state.ClearDKGMasterPublicKeys()
 	g.state.ClearDKGComplaintProposed()
 	g.state.ClearDKGComplaints()
@@ -1371,7 +1363,8 @@ func (g *GovernanceContract) fineFailStopDKG(threshold int) {
 	fineNode := make(map[coreTypes.NodeID]struct{})
 	dkgSet := g.getNotarySet(g.state.DKGRound())
 	for id := range dkgSet {
-		if g.state.DKGMasterPublicKeyProposed(Bytes32(id.Hash)) {
+		offset := g.state.DKGMasterPublicKeyOffset(Bytes32(id.Hash))
+		if offset.Cmp(big.NewInt(0)) >= 0 {
 			continue
 		}
 		fineNode[id] = struct{}{}
@@ -1465,10 +1458,8 @@ func (g *GovernanceContract) addDKGComplaint(comp []byte) ([]byte, error) {
 		return nil, errExecutionReverted
 	}
 
-	mpk, err := g.state.GetDKGMasterPublicKeyByProposerID(dkgComplaint.PrivateShare.ProposerID)
-	if err != nil {
-		return nil, errExecutionReverted
-	}
+	mpkOffset := g.state.DKGMasterPublicKeyOffset(Bytes32(dkgComplaint.PrivateShare.ProposerID.Hash))
+	mpk := g.state.DKGMasterPublicKeyItem(mpkOffset)
 
 	// Verify DKG complaint is correct.
 	ok, err := coreUtils.VerifyDKGComplaint(&dkgComplaint, mpk)
@@ -1514,7 +1505,8 @@ func (g *GovernanceContract) addDKGMasterPublicKey(mpk []byte) ([]byte, error) {
 		g.state.SetDKGRound(round)
 	}
 
-	if g.state.DKGMasterPublicKeyProposed(getDKGMasterPublicKeyID(&dkgMasterPK)) {
+	mpkOffset := g.state.DKGMasterPublicKeyOffset(getDKGMasterPublicKeyID(&dkgMasterPK))
+	if mpkOffset.Cmp(big.NewInt(0)) >= 0 {
 		return nil, errExecutionReverted
 	}
 
@@ -1553,8 +1545,10 @@ func (g *GovernanceContract) addDKGMasterPublicKey(mpk []byte) ([]byte, error) {
 		return nil, errExecutionReverted
 	}
 
+	mpkOffset = g.state.LenDKGMasterPublicKeys()
 	g.state.PushDKGMasterPublicKey(mpk)
-	g.state.PutDKGMasterPublicKeyProposed(getDKGMasterPublicKeyID(&dkgMasterPK), true)
+	g.state.PutDKGMasterPublicKeyOffset(getDKGMasterPublicKeyID(&dkgMasterPK), mpkOffset)
+
 	return g.useGas(GovernanceActionGasCost)
 }
 
@@ -2233,16 +2227,23 @@ func (g *GovernanceContract) Run(evm *EVM, input []byte, contract *Contract) (re
 		}
 		return res, nil
 	case "dkgComplaints":
-		index := new(big.Int)
-		if err := method.Inputs.Unpack(&index, arguments); err != nil {
+		offset := new(big.Int)
+		if err := method.Inputs.Unpack(&offset, arguments); err != nil {
 			return nil, errExecutionReverted
 		}
-		complaints := g.state.DKGComplaints()
-		if int(index.Uint64()) >= len(complaints) {
+		complaint := g.state.DKGComplaint(offset)
+		res, err := method.Outputs.Pack(complaint)
+		if err != nil {
 			return nil, errExecutionReverted
 		}
-		complaint := complaints[index.Uint64()]
-		res, err := method.Outputs.Pack(complaint)
+		return res, nil
+	case "dkgComplaintsProposed":
+		id := Bytes32{}
+		if err := method.Inputs.Unpack(&id, arguments); err != nil {
+			return nil, errExecutionReverted
+		}
+		proposed := g.state.DKGComplaintProposed(id)
+		res, err := method.Outputs.Pack(proposed)
 		if err != nil {
 			return nil, errExecutionReverted
 		}
@@ -2266,16 +2267,23 @@ func (g *GovernanceContract) Run(evm *EVM, input []byte, contract *Contract) (re
 		}
 		return res, nil
 	case "dkgMasterPublicKeys":
-		index := new(big.Int)
-		if err := method.Inputs.Unpack(&index, arguments); err != nil {
+		offset := new(big.Int)
+		if err := method.Inputs.Unpack(&offset, arguments); err != nil {
 			return nil, errExecutionReverted
 		}
-		mpks := g.state.DKGMasterPublicKeys()
-		if int(index.Uint64()) >= len(mpks) {
+		mpk := g.state.DKGMasterPublicKey(offset)
+		res, err := method.Outputs.Pack(mpk)
+		if err != nil {
 			return nil, errExecutionReverted
 		}
-		mpk := mpks[index.Uint64()]
-		res, err := method.Outputs.Pack(mpk)
+		return res, nil
+	case "dkgMasterPublicKeyOffset":
+		id := Bytes32{}
+		if err := method.Inputs.Unpack(&id, arguments); err != nil {
+			return nil, errExecutionReverted
+		}
+		offset := g.state.DKGMasterPublicKeyOffset(id)
+		res, err := method.Outputs.Pack(offset)
 		if err != nil {
 			return nil, errExecutionReverted
 		}
diff --git a/core/vm/oracle_contracts_test.go b/core/vm/oracle_contracts_test.go
index 5037a7cc0..269421717 100644
--- a/core/vm/oracle_contracts_test.go
+++ b/core/vm/oracle_contracts_test.go
@@ -977,7 +977,7 @@ func (g *OracleContractsTestSuite) TestResetDKG() {
 
 			// Clear DKG states for next round.
 			dkgSet := dkgSets[round-1]
-			g.s.ClearDKGMasterPublicKeyProposed()
+			g.s.ClearDKGMasterPublicKeyOffset()
 			g.s.ClearDKGMasterPublicKeys()
 			g.s.ClearDKGComplaintProposed()
 			g.s.ClearDKGComplaints()
@@ -1003,6 +1003,7 @@ func (g *OracleContractsTestSuite) TestResetDKG() {
 		g.Require().Len(dkgSet, int(g.config.NotarySetSize))
 		dkgSets[round] = dkgSet
 
+		i := 0
 		for id := range dkgSet {
 			offset := g.s.NodesOffsetByNodeKeyAddress(IdToAddress(id))
 			if offset.Cmp(big.NewInt(0)) < 0 {
@@ -1016,7 +1017,7 @@ func (g *OracleContractsTestSuite) TestResetDKG() {
 				panic(err)
 			}
 			g.s.PushDKGMasterPublicKey(b)
-			g.s.PutDKGMasterPublicKeyProposed(Bytes32(id.Hash), true)
+			g.s.PutDKGMasterPublicKeyOffset(Bytes32(id.Hash), big.NewInt(int64(i)))
 			// Prepare Complaint.
 			y := dkgTypes.Complaint{}
 			b, err = rlp.EncodeToBytes(&y)
@@ -1034,6 +1035,7 @@ func (g *OracleContractsTestSuite) TestResetDKG() {
 				g.s.PutDKGFinalized(addr, true)
 				g.s.IncDKGFinalizedsCount()
 			}
+			i += 1
 		}
 		dkgSetSize := len(dkgSet)
 		g.Require().Len(g.s.DKGMasterPublicKeys(), dkgSetSize)