Pull apart writeCurrentStakers method - Part 1

vms/platformvm/state/state.go

@@ -1672,10 +1672,17 @@ func (s *state) initValidatorSets() error {
 }
 
 func (s *state) write(updateValidators bool, height uint64) error {
+	weightDiffs, blsKeyDiffs, err := s.processCurrentStakers()
+	if err != nil {
+		return err
+	}
+
 	errs := wrappers.Errs{}
 	errs.Add(
 		s.writeBlocks(),
-		s.writeCurrentStakers(updateValidators, height),
+		s.writeCurrentStakers(updateValidators),
+		s.writeWeightDiffs(height, weightDiffs),
+		s.writeBlsKeyDiffs(height, blsKeyDiffs),
 		s.writePendingStakers(),
 		s.WriteValidatorMetadata(s.currentValidatorList, s.currentSubnetValidatorList), // Must be called after writeCurrentStakers
 		s.writeTXs(),
@@ -1690,6 +1697,155 @@ func (s *state) write(updateValidators bool, height uint64) error {
 	return errs.Err
 }
 
+func (s *state) writeWeightDiffs(height uint64, weightDiffs map[subnetNodePair]*ValidatorWeightDiff) error {
+	for k, weightDiff := range weightDiffs {
+		if weightDiff.Amount == 0 {
+			continue
+		}
+
+		err := s.flatValidatorWeightDiffsDB.Put(
+			marshalDiffKey(k.subnetID, height, k.nodeID),
+			marshalWeightDiff(weightDiff),
+		)
+		if err != nil {
+			return err
+		}
+
+		// TODO: Remove this once we no longer support version rollbacks.
+		prefixStruct := heightWithSubnet{
+			Height:   height,
+			SubnetID: k.subnetID,
+		}
+		prefixBytes, err := block.GenesisCodec.Marshal(block.Version, prefixStruct)
+		if err != nil {
+			return fmt.Errorf("failed to create prefix bytes: %w", err)
+		}
+		rawNestedWeightDiffDB := prefixdb.New(prefixBytes, s.nestedValidatorWeightDiffsDB)
+		nestedWeightDiffDB := linkeddb.NewDefault(rawNestedWeightDiffDB)
+
+		weightDiffBytes, err := block.GenesisCodec.Marshal(block.Version, weightDiff)
+		if err != nil {
+			return fmt.Errorf("failed to serialize validator weight diff: %w", err)
+		}
+		if err := nestedWeightDiffDB.Put(k.nodeID[:], weightDiffBytes); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (s *state) writeBlsKeyDiffs(height uint64, blsKeyDiffs map[ids.NodeID]*bls.PublicKey) error {
+	for nodeID, blsKey := range blsKeyDiffs {
+		key := marshalDiffKey(constants.PrimaryNetworkID, height, nodeID)
+		blsKeyBytes := []byte{}
+		if blsKey != nil {
+			// Note: in flatValidatorPublicKeyDiffsDB we store the
+			// uncompressed public key here as it is
+			// significantly more efficient to parse when applying diffs.
+			blsKeyBytes = blsKey.Serialize()
+		}
+		if err := s.flatValidatorPublicKeyDiffsDB.Put(key, blsKeyBytes); err != nil {
+			return fmt.Errorf("failed to add bls key diffs: %w", err)
+		}
+
+		// TODO: Remove this once we no longer support version rollbacks.
+		if blsKey != nil {
+			heightBytes := database.PackUInt64(height)
+			rawNestedPublicKeyDiffDB := prefixdb.New(heightBytes, s.nestedValidatorPublicKeyDiffsDB)
+			nestedPKDiffDB := linkeddb.NewDefault(rawNestedPublicKeyDiffDB)
+
+			// Note: We store the compressed public key here.
+			pkBytes := bls.PublicKeyToBytes(blsKey)
+			if err := nestedPKDiffDB.Put(nodeID[:], pkBytes); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+type subnetNodePair struct {
+	subnetID ids.ID
+	nodeID   ids.NodeID
+}
+
+func (s *state) processCurrentStakers() (
+	map[subnetNodePair]*ValidatorWeightDiff,
+	map[ids.NodeID]*bls.PublicKey,
+	error,
+) {
+	var (
+		outputWeights = make(map[subnetNodePair]*ValidatorWeightDiff)
+		outputBlsKey  = make(map[ids.NodeID]*bls.PublicKey)
+	)
+
+	for subnetID, subnetValidatorDiffs := range s.currentStakers.validatorDiffs {
+		// for now, let writeCurrentStakers consume s.currentStakers.validatorDiffs
+		for nodeID, validatorDiff := range subnetValidatorDiffs {
+			weightKey := subnetNodePair{
+				subnetID: subnetID,
+				nodeID:   nodeID,
+			}
+
+			// make sure there is an entry for delegators even in case
+			// there are no validators modified.
+			outputWeights[weightKey] = &ValidatorWeightDiff{
+				Decrease: validatorDiff.validatorStatus == deleted,
+			}
+
+			switch validatorDiff.validatorStatus {
+			case added:
+				var (
+					weight = validatorDiff.validator.Weight
+					blkKey = validatorDiff.validator.PublicKey
+				)
+
+				outputWeights[weightKey].Amount = weight
+
+				if blkKey != nil {
+					// Record that the public key for the validator is being
+					// added. This means the prior value for the public key was
+					// nil.
+					outputBlsKey[nodeID] = nil
+				}
+
+			case deleted:
+				var (
+					weight = validatorDiff.validator.Weight
+					blkKey = validatorDiff.validator.PublicKey
+				)
+				outputWeights[weightKey].Amount = weight
+				if blkKey != nil {
+					// Record that the public key for the validator is being
+					// removed. This means we must record the prior value of the
+					// public key.
+					outputBlsKey[nodeID] = blkKey
+				}
+
+			default:
+				// nothing to do
+			}
+
+			addedDelegatorIterator := NewTreeIterator(validatorDiff.addedDelegators)
+			defer addedDelegatorIterator.Release()
+			for addedDelegatorIterator.Next() {
+				staker := addedDelegatorIterator.Value()
+
+				if err := outputWeights[weightKey].Add(false, staker.Weight); err != nil {
+					return nil, nil, fmt.Errorf("failed to increase node weight diff: %w", err)
+				}
+			}
+
+			for _, staker := range validatorDiff.deletedDelegators {
+				if err := outputWeights[weightKey].Add(true, staker.Weight); err != nil {
+					return nil, nil, fmt.Errorf("failed to decrease node weight diff: %w", err)
+				}
+			}
+		}
+	}
+	return outputWeights, outputBlsKey, nil
+}
+
 func (s *state) Close() error {
 	errs := wrappers.Errs{}
 	errs.Add(
@@ -1905,11 +2061,7 @@ func (s *state) GetBlockIDAtHeight(height uint64) (ids.ID, error) {
 	return blkID, nil
 }
 
-func (s *state) writeCurrentStakers(updateValidators bool, height uint64) error {
-	heightBytes := database.PackUInt64(height)
-	rawNestedPublicKeyDiffDB := prefixdb.New(heightBytes, s.nestedValidatorPublicKeyDiffsDB)
-	nestedPKDiffDB := linkeddb.NewDefault(rawNestedPublicKeyDiffDB)
-
+func (s *state) writeCurrentStakers(updateValidators bool) error {
 	for subnetID, validatorDiffs := range s.currentStakers.validatorDiffs {
 		delete(s.currentStakers.validatorDiffs, subnetID)
 
@@ -1921,17 +2073,6 @@ func (s *state) writeCurrentStakers(updateValidators bool, height uint64) error
 			delegatorDB = s.currentDelegatorList
 		}
 
-		prefixStruct := heightWithSubnet{
-			Height:   height,
-			SubnetID: subnetID,
-		}
-		prefixBytes, err := block.GenesisCodec.Marshal(block.Version, prefixStruct)
-		if err != nil {
-			return fmt.Errorf("failed to create prefix bytes: %w", err)
-		}
-		rawNestedWeightDiffDB := prefixdb.New(prefixBytes, s.nestedValidatorWeightDiffsDB)
-		nestedWeightDiffDB := linkeddb.NewDefault(rawNestedWeightDiffDB)
-
 		// Record the change in weight and/or public key for each validator.
 		for nodeID, validatorDiff := range validatorDiffs {
 			// Copy [nodeID] so it doesn't get overwritten next iteration.
@@ -1945,21 +2086,6 @@ func (s *state) writeCurrentStakers(updateValidators bool, height uint64) error
 				staker := validatorDiff.validator
 				weightDiff.Amount = staker.Weight
 
-				// Invariant: Only the Primary Network contains non-nil public
-				// keys.
-				if staker.PublicKey != nil {
-					// Record that the public key for the validator is being
-					// added. This means the prior value for the public key was
-					// nil.
-					err := s.flatValidatorPublicKeyDiffsDB.Put(
-						marshalDiffKey(constants.PrimaryNetworkID, height, nodeID),
-						nil,
-					)
-					if err != nil {
-						return err
-					}
-				}
-
 				// The validator is being added.
 				//
 				// Invariant: It's impossible for a delegator to have been
@@ -1988,34 +2114,6 @@ func (s *state) writeCurrentStakers(updateValidators bool, height uint64) error
 				staker := validatorDiff.validator
 				weightDiff.Amount = staker.Weight
 
-				// Invariant: Only the Primary Network contains non-nil public
-				// keys.
-				if staker.PublicKey != nil {
-					// Record that the public key for the validator is being
-					// removed. This means we must record the prior value of the
-					// public key.
-					//
-					// Note: We store the uncompressed public key here as it is
-					// significantly more efficient to parse when applying
-					// diffs.
-					err := s.flatValidatorPublicKeyDiffsDB.Put(
-						marshalDiffKey(constants.PrimaryNetworkID, height, nodeID),
-						staker.PublicKey.Serialize(),
-					)
-					if err != nil {
-						return err
-					}
-
-					// TODO: Remove this once we no longer support version
-					// rollbacks.
-					//
-					// Note: We store the compressed public key here.
-					pkBytes := bls.PublicKeyToBytes(staker.PublicKey)
-					if err := nestedPKDiffDB.Put(nodeID[:], pkBytes); err != nil {
-						return err
-					}
-				}
-
 				if err := validatorDB.Delete(staker.TxID[:]); err != nil {
 					return fmt.Errorf("failed to delete current staker: %w", err)
 				}
@@ -2032,28 +2130,6 @@ func (s *state) writeCurrentStakers(updateValidators bool, height uint64) error
 				return err
 			}
 
-			if weightDiff.Amount == 0 {
-				// No weight change to record; go to next validator.
-				continue
-			}
-
-			err = s.flatValidatorWeightDiffsDB.Put(
-				marshalDiffKey(subnetID, height, nodeID),
-				marshalWeightDiff(weightDiff),
-			)
-			if err != nil {
-				return err
-			}
-
-			// TODO: Remove this once we no longer support version rollbacks.
-			weightDiffBytes, err := block.GenesisCodec.Marshal(block.Version, weightDiff)
-			if err != nil {
-				return fmt.Errorf("failed to serialize validator weight diff: %w", err)
-			}
-			if err := nestedWeightDiffDB.Put(nodeID[:], weightDiffBytes); err != nil {
-				return err
-			}
-
 			// TODO: Move the validator set management out of the state package
 			if !updateValidators {
 				continue