Deneb DB methods (#12379)

beacon-chain/db/kv/state.go

@@ -229,34 +229,28 @@ func (s *Store) saveStatesEfficientInternal(ctx context.Context, tx *bolt.Tx, bl
 		// look at issue https://github.com/prysmaticlabs/prysm/issues/9262.
 		switch rawType := states[i].ToProtoUnsafe().(type) {
 		case *ethpb.BeaconState:
-			pbState, err := statenative.ProtobufBeaconStatePhase0(rawType)
+			pbState, err := getPhase0PbState(rawType)
 			if err != nil {
 				return err
 			}
-			if pbState == nil {
-				return errors.New("nil state")
-			}
 			valEntries := pbState.Validators
 			pbState.Validators = make([]*ethpb.Validator, 0)
 			encodedState, err := encode(ctx, pbState)
 			if err != nil {
 				return err
 			}
+			pbState.Validators = valEntries
 			if err := bucket.Put(rt[:], encodedState); err != nil {
 				return err
 			}
-			pbState.Validators = valEntries
 			if err := valIdxBkt.Put(rt[:], validatorKeys[i]); err != nil {
 				return err
 			}
 		case *ethpb.BeaconStateAltair:
-			pbState, err := statenative.ProtobufBeaconStateAltair(rawType)
+			pbState, err := getAltairPbState(rawType)
 			if err != nil {
 				return err
 			}
-			if pbState == nil {
-				return errors.New("nil state")
-			}
 			valEntries := pbState.Validators
 			pbState.Validators = make([]*ethpb.Validator, 0)
 			rawObj, err := pbState.MarshalSSZ()
@@ -272,13 +266,10 @@ func (s *Store) saveStatesEfficientInternal(ctx context.Context, tx *bolt.Tx, bl
 				return err
 			}
 		case *ethpb.BeaconStateBellatrix:
-			pbState, err := statenative.ProtobufBeaconStateBellatrix(rawType)
+			pbState, err := getBellatrixPbState(rawType)
 			if err != nil {
 				return err
 			}
-			if pbState == nil {
-				return errors.New("nil state")
-			}
 			valEntries := pbState.Validators
 			pbState.Validators = make([]*ethpb.Validator, 0)
 			rawObj, err := pbState.MarshalSSZ()
@@ -294,12 +285,28 @@ func (s *Store) saveStatesEfficientInternal(ctx context.Context, tx *bolt.Tx, bl
 				return err
 			}
 		case *ethpb.BeaconStateCapella:
-			pbState, err := statenative.ProtobufBeaconStateCapella(rawType)
+			pbState, err := getCapellaPbState(rawType)
 			if err != nil {
 				return err
 			}
-			if pbState == nil {
-				return errors.New("nil state")
+			valEntries := pbState.Validators
+			pbState.Validators = make([]*ethpb.Validator, 0)
+			rawObj, err := pbState.MarshalSSZ()
+			if err != nil {
+				return err
+			}
+			encodedState := snappy.Encode(nil, append(capellaKey, rawObj...))
+			if err := bucket.Put(rt[:], encodedState); err != nil {
+				return err
+			}
+			pbState.Validators = valEntries
+			if err := valIdxBkt.Put(rt[:], validatorKeys[i]); err != nil {
+				return err
+			}
+		case *ethpb.BeaconStateDeneb:
+			pbState, err := getDenebPbState(rawType)
+			if err != nil {
+				return err
 			}
 			valEntries := pbState.Validators
 			pbState.Validators = make([]*ethpb.Validator, 0)
@@ -323,6 +330,61 @@ func (s *Store) saveStatesEfficientInternal(ctx context.Context, tx *bolt.Tx, bl
 	return s.storeValidatorEntriesSeparately(ctx, tx, validatorsEntries)
 }
 
+func getPhase0PbState(rawState interface{}) (*ethpb.BeaconState, error) {
+	pbState, err := statenative.ProtobufBeaconStatePhase0(rawState)
+	if err != nil {
+		return nil, err
+	}
+	if pbState == nil {
+		return nil, errors.New("nil state")
+	}
+	return pbState, nil
+}
+
+func getAltairPbState(rawState interface{}) (*ethpb.BeaconStateAltair, error) {
+	pbState, err := statenative.ProtobufBeaconStateAltair(rawState)
+	if err != nil {
+		return nil, err
+	}
+	if pbState == nil {
+		return nil, errors.New("nil state")
+	}
+	return pbState, nil
+}
+
+func getBellatrixPbState(rawState interface{}) (*ethpb.BeaconStateBellatrix, error) {
+	pbState, err := statenative.ProtobufBeaconStateBellatrix(rawState)
+	if err != nil {
+		return nil, err
+	}
+	if pbState == nil {
+		return nil, errors.New("nil state")
+	}
+	return pbState, nil
+}
+
+func getCapellaPbState(rawState interface{}) (*ethpb.BeaconStateCapella, error) {
+	pbState, err := statenative.ProtobufBeaconStateCapella(rawState)
+	if err != nil {
+		return nil, err
+	}
+	if pbState == nil {
+		return nil, errors.New("nil state")
+	}
+	return pbState, nil
+}
+
+func getDenebPbState(rawState interface{}) (*ethpb.BeaconStateDeneb, error) {
+	pbState, err := statenative.ProtobufBeaconStateDeneb(rawState)
+	if err != nil {
+		return nil, err
+	}
+	if pbState == nil {
+		return nil, errors.New("nil state")
+	}
+	return pbState, nil
+}
+
 func (s *Store) storeValidatorEntriesSeparately(ctx context.Context, tx *bolt.Tx, validatorsEntries map[string]*ethpb.Validator) error {
 	valBkt := tx.Bucket(stateValidatorsBucket)
 	for hashStr, validatorEntry := range validatorsEntries {
@@ -472,6 +534,19 @@ func (s *Store) unmarshalState(_ context.Context, enc []byte, validatorEntries [
 	}
 
 	switch {
+	case hasDenebKey(enc):
+		protoState := &ethpb.BeaconStateDeneb{}
+		if err := protoState.UnmarshalSSZ(enc[len(denebKey):]); err != nil {
+			return nil, errors.Wrap(err, "failed to unmarshal encoding for Deneb")
+		}
+		ok, err := s.isStateValidatorMigrationOver()
+		if err != nil {
+			return nil, err
+		}
+		if ok {
+			protoState.Validators = validatorEntries
+		}
+		return statenative.InitializeFromProtoUnsafeDeneb(protoState)
 	case hasCapellaKey(enc):
 		// Marshal state bytes to capella beacon state.
 		protoState := &ethpb.BeaconStateCapella{}
@@ -579,6 +654,19 @@ func marshalState(ctx context.Context, st state.ReadOnlyBeaconState) ([]byte, er
 			return nil, err
 		}
 		return snappy.Encode(nil, append(capellaKey, rawObj...)), nil
+	case *ethpb.BeaconStateDeneb:
+		rState, ok := st.ToProtoUnsafe().(*ethpb.BeaconStateDeneb)
+		if !ok {
+			return nil, errors.New("non valid inner state")
+		}
+		if rState == nil {
+			return nil, errors.New("nil state")
+		}
+		rawObj, err := rState.MarshalSSZ()
+		if err != nil {
+			return nil, err
+		}
+		return snappy.Encode(nil, append(denebKey, rawObj...)), nil
 	default:
 		return nil, errors.New("invalid inner state")
 	}