feat(beradb): withdrawals (#1991)

berachain · Sep 3, 2024 · ea8e773 · ea8e773
1 parent f598621
commit ea8e773
Show file tree

Hide file tree

Showing 7 changed files with 54 additions and 39 deletions.
diff --git a/examples/berad/pkg/storage/history.go b/examples/berad/pkg/storage/history.go
@@ -25,7 +25,7 @@ import "github.com/berachain/beacon-kit/mod/primitives/pkg/common"
 // UpdateBlockRootAtIndex sets a block root in the BeaconStore.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) UpdateBlockRootAtIndex(
 	index uint64,
 	root common.Root,
@@ -36,7 +36,7 @@ func (kv *KVStore[
 // GetBlockRootAtIndex retrieves the block root from the BeaconStore.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) GetBlockRootAtIndex(
 	index uint64,
 ) (common.Root, error) {
@@ -50,7 +50,7 @@ func (kv *KVStore[
 // SetLatestBlockHeader sets the latest block header in the BeaconStore.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) SetLatestBlockHeader(
 	header BeaconBlockHeaderT,
 ) error {
@@ -60,7 +60,7 @@ func (kv *KVStore[
 // GetLatestBlockHeader retrieves the latest block header from the BeaconStore.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) GetLatestBlockHeader() (
 	BeaconBlockHeaderT, error,
 ) {
@@ -70,7 +70,7 @@ func (kv *KVStore[
 // UpdateStateRootAtIndex updates the state root at the given slot.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) UpdateStateRootAtIndex(
 	idx uint64,
 	stateRoot common.Root,
@@ -81,7 +81,7 @@ func (kv *KVStore[
 // StateRootAtIndex returns the state root at the given slot.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) StateRootAtIndex(
 	idx uint64,
 ) (common.Root, error) {
@@ -96,7 +96,7 @@ func (kv *KVStore[
 // header from the BeaconStore.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) GetLatestExecutionPayloadHeader() (
 	ExecutionPayloadHeaderT, error,
 ) {
@@ -113,7 +113,7 @@ func (kv *KVStore[
 // the BeaconStore.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) SetLatestExecutionPayloadHeader(
 	payloadHeader ExecutionPayloadHeaderT,
 ) error {

diff --git a/examples/berad/pkg/storage/keys/keys.go b/examples/berad/pkg/storage/keys/keys.go
@@ -32,8 +32,9 @@ const (
 	LatestExecutionPayloadHeaderPrefix
 	LatestExecutionPayloadVersionPrefix
 	GenesisValidatorsRootPrefix
-	NextWithdrawalIndexPrefix
-	NextWithdrawalValidatorIndexPrefix
+	WithdrawalsPrefix
+	// NextWithdrawalIndexPrefix
+	// NextWithdrawalValidatorIndexPrefix.
 	ForkPrefix
 )
 
@@ -50,7 +51,6 @@ const (
 	LatestExecutionPayloadHeaderPrefixHumanReadable  = "LatestExecutionPayloadHeaderPrefix"
 	LatestExecutionPayloadVersionPrefixHumanReadable = "LatestExecutionPayloadVersionPrefix"
 	GenesisValidatorsRootPrefixHumanReadable         = "GenesisValidatorsRootPrefix"
-	NextWithdrawalIndexPrefixHumanReadable           = "NextWithdrawalIndexPrefix"
-	NextWithdrawalValidatorIndexPrefixHumanReadable  = "NextWithdrawalValidatorIndexPrefix"
+	WithdrawalsPrefixHumanReadable                   = "WithdrawalsPrefix"
 	ForkPrefixHumanReadable                          = "ForkPrefix"
 )
diff --git a/examples/berad/pkg/storage/kvstore.go b/examples/berad/pkg/storage/kvstore.go
@@ -25,9 +25,9 @@ import (
 
 	sdkcollections "cosmossdk.io/collections"
 	"cosmossdk.io/core/store"
+	"github.com/berachain/beacon-kit/examples/berad/pkg/storage/keys"
 	"github.com/berachain/beacon-kit/mod/primitives/pkg/constraints"
 	"github.com/berachain/beacon-kit/mod/storage/pkg/beacondb/index"
-	"github.com/berachain/beacon-kit/mod/storage/pkg/beacondb/keys"
 	"github.com/berachain/beacon-kit/mod/storage/pkg/encoding"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
@@ -50,6 +50,7 @@ type KVStore[
 	},
 	ValidatorT Validator[ValidatorT],
 	ValidatorsT ~[]ValidatorT,
+	WithdrawalT any,
 ] struct {
 	ctx context.Context
 	// Versioning
@@ -87,6 +88,9 @@ type KVStore[
 	// Randomness
 	// randaoMix stores the randao mix for the current epoch.
 	randaoMix sdkcollections.Map[uint64, []byte]
+	// Staking
+	// withdrawals stores a list of initiated withdrawals.
+	withdrawals sdkcollections.Map[uint64, WithdrawalT]
 }
 
 // New creates a new instance of Store.
@@ -112,17 +116,21 @@ func New[
 	},
 	ValidatorT Validator[ValidatorT],
 	ValidatorsT ~[]ValidatorT,
+	WithdrawalT interface {
+		constraints.Empty[WithdrawalT]
+		constraints.SSZMarshallable
+	},
 ](
 	kss store.KVStoreService,
 	payloadCodec *encoding.SSZInterfaceCodec[ExecutionPayloadHeaderT],
 ) *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ] {
 	schemaBuilder := sdkcollections.NewSchemaBuilder(kss)
 	return &KVStore[
 		BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-		ForkT, ValidatorT, ValidatorsT,
+		ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 	]{
 		ctx: nil,
 		genesisValidatorsRoot: sdkcollections.NewItem(
@@ -208,16 +216,23 @@ func New[
 			keys.LatestBeaconBlockHeaderPrefixHumanReadable,
 			encoding.SSZValueCodec[BeaconBlockHeaderT]{},
 		),
+		withdrawals: sdkcollections.NewMap(
+			schemaBuilder,
+			sdkcollections.NewPrefix([]byte{keys.WithdrawalsPrefix}),
+			keys.WithdrawalsPrefixHumanReadable,
+			sdkcollections.Uint64Key,
+			encoding.SSZValueCodec[WithdrawalT]{},
+		),
 	}
 }
 
 // Copy returns a copy of the Store.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) Copy() *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ] {
 	// TODO: Decouple the KVStore type from the Cosmos-SDK.
 	cctx, _ := sdk.UnwrapSDKContext(kv.ctx).CacheContext()
@@ -228,20 +243,20 @@ func (kv *KVStore[
 // Context returns the context of the Store.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) Context() context.Context {
 	return kv.ctx
 }
 
 // WithContext returns a copy of the Store with the given context.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) WithContext(
 	ctx context.Context,
 ) *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ] {
 	cpy := *kv
 	cpy.ctx = ctx

diff --git a/examples/berad/pkg/storage/randao.go b/examples/berad/pkg/storage/randao.go
@@ -25,7 +25,7 @@ import "github.com/berachain/beacon-kit/mod/primitives/pkg/common"
 // UpdateRandaoMixAtIndex sets the current RANDAO mix in the store.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) UpdateRandaoMixAtIndex(
 	index uint64,
 	mix common.Bytes32,
@@ -36,7 +36,7 @@ func (kv *KVStore[
 // GetRandaoMixAtIndex retrieves the current RANDAO mix from the store.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) GetRandaoMixAtIndex(
 	index uint64,
 ) (common.Bytes32, error) {

diff --git a/examples/berad/pkg/storage/registry.go b/examples/berad/pkg/storage/registry.go
@@ -29,7 +29,7 @@ import (
 // AddValidator registers a new validator in the beacon state.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) AddValidator(val ValidatorT) error {
 	// Get the next validator index from the sequence.
 	idx, err := kv.validatorIndex.Next(kv.ctx)
@@ -44,7 +44,7 @@ func (kv *KVStore[
 // UpdateValidatorAtIndex updates a validator at a specific index.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) UpdateValidatorAtIndex(
 	index math.ValidatorIndex,
 	val ValidatorT,
@@ -55,7 +55,7 @@ func (kv *KVStore[
 // ValidatorIndexByPubkey returns the validator address by index.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) ValidatorIndexByPubkey(
 	pubkey crypto.BLSPubkey,
 ) (math.ValidatorIndex, error) {
@@ -72,7 +72,7 @@ func (kv *KVStore[
 // ValidatorIndexByCometBFTAddress returns the validator address by index.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) ValidatorIndexByCometBFTAddress(
 	cometBFTAddress []byte,
 ) (math.ValidatorIndex, error) {
@@ -89,7 +89,7 @@ func (kv *KVStore[
 // ValidatorByIndex returns the validator address by index.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) ValidatorByIndex(
 	index math.ValidatorIndex,
 ) (ValidatorT, error) {
@@ -104,7 +104,7 @@ func (kv *KVStore[
 // GetValidators retrieves all validators from the beacon state.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) GetValidators() (
 	ValidatorsT, error,
 ) {
@@ -140,7 +140,7 @@ func (kv *KVStore[
 // GetTotalValidators returns the total number of validators.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) GetTotalValidators() (uint64, error) {
 	validators, err := kv.GetValidators()
 	if err != nil {
@@ -153,7 +153,7 @@ func (kv *KVStore[
 // effective balance from the beacon state.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) GetValidatorsByEffectiveBalance() (
 	[]ValidatorT, error,
 ) {
@@ -190,7 +190,7 @@ func (kv *KVStore[
 // TODO: this shouldn't live in KVStore
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) GetTotalActiveBalances(
 	slotsPerEpoch uint64,
 ) (math.Gwei, error) {

diff --git a/examples/berad/pkg/storage/staking.go b/examples/berad/pkg/storage/staking.go
@@ -23,15 +23,15 @@ package beacondb
 // GetEth1DepositIndex retrieves the eth1 deposit index from the beacon state.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) GetEth1DepositIndex() (uint64, error) {
 	return kv.eth1DepositIndex.Get(kv.ctx)
 }
 
 // SetEth1DepositIndex sets the eth1 deposit index in the beacon state.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) SetEth1DepositIndex(
 	index uint64,
 ) error {

diff --git a/examples/berad/pkg/storage/versioning.go b/examples/berad/pkg/storage/versioning.go
@@ -29,7 +29,7 @@ import (
 // state.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) SetGenesisValidatorsRoot(
 	root common.Root,
 ) error {
@@ -40,7 +40,7 @@ func (kv *KVStore[
 // beacon state.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) GetGenesisValidatorsRoot() (common.Root, error) {
 	bz, err := kv.genesisValidatorsRoot.Get(kv.ctx)
 	if err != nil {
@@ -52,7 +52,7 @@ func (kv *KVStore[
 // GetSlot returns the current slot.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) GetSlot() (math.Slot, error) {
 	slot, err := kv.slot.Get(kv.ctx)
 	return math.Slot(slot), err
@@ -61,7 +61,7 @@ func (kv *KVStore[
 // SetSlot sets the current slot.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) SetSlot(
 	slot math.Slot,
 ) error {
@@ -71,7 +71,7 @@ func (kv *KVStore[
 // SetFork sets the fork version for the given epoch.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) SetFork(
 	fork ForkT,
 ) error {
@@ -81,7 +81,7 @@ func (kv *KVStore[
 // GetFork gets the fork version for the given epoch.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
 ]) GetFork() (ForkT, error) {
 	return kv.fork.Get(kv.ctx)
 }