feat(beradb): storage rw (#1992)

Co-authored-by: berabera <[email protected]>
berachain · Sep 4, 2024 · 99aeabe · 99aeabe
1 parent ea8e773
commit 99aeabe
Show file tree

Hide file tree

Showing 6 changed files with 59 additions and 38 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) GetLatestExecutionPayloadHeader() (
 	ExecutionPayloadHeaderT, error,
 ) {
@@ -113,7 +113,7 @@ func (kv *KVStore[
 // the BeaconStore.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) SetLatestExecutionPayloadHeader(
 	payloadHeader ExecutionPayloadHeaderT,
 ) error {

diff --git a/examples/berad/pkg/storage/kvstore.go b/examples/berad/pkg/storage/kvstore.go
@@ -51,6 +51,7 @@ type KVStore[
 	ValidatorT Validator[ValidatorT],
 	ValidatorsT ~[]ValidatorT,
 	WithdrawalT any,
+	WithdrawalsT ~[]WithdrawalT,
 ] struct {
 	ctx context.Context
 	// Versioning
@@ -90,7 +91,7 @@ type KVStore[
 	randaoMix sdkcollections.Map[uint64, []byte]
 	// Staking
 	// withdrawals stores a list of initiated withdrawals.
-	withdrawals sdkcollections.Map[uint64, WithdrawalT]
+	withdrawals sdkcollections.Item[WithdrawalsT]
 }
 
 // New creates a new instance of Store.
@@ -120,17 +121,22 @@ func New[
 		constraints.Empty[WithdrawalT]
 		constraints.SSZMarshallable
 	},
+	WithdrawalsT interface {
+		~[]WithdrawalT
+		constraints.Empty[WithdrawalsT]
+		constraints.SSZMarshallable
+	},
 ](
 	kss store.KVStoreService,
 	payloadCodec *encoding.SSZInterfaceCodec[ExecutionPayloadHeaderT],
 ) *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ] {
 	schemaBuilder := sdkcollections.NewSchemaBuilder(kss)
 	return &KVStore[
 		BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-		ForkT, ValidatorT, ValidatorsT, WithdrawalT,
+		ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 	]{
 		ctx: nil,
 		genesisValidatorsRoot: sdkcollections.NewItem(
@@ -216,23 +222,22 @@ func New[
 			keys.LatestBeaconBlockHeaderPrefixHumanReadable,
 			encoding.SSZValueCodec[BeaconBlockHeaderT]{},
 		),
-		withdrawals: sdkcollections.NewMap(
+		withdrawals: sdkcollections.NewItem(
 			schemaBuilder,
 			sdkcollections.NewPrefix([]byte{keys.WithdrawalsPrefix}),
 			keys.WithdrawalsPrefixHumanReadable,
-			sdkcollections.Uint64Key,
-			encoding.SSZValueCodec[WithdrawalT]{},
+			encoding.SSZValueCodec[WithdrawalsT]{},
 		),
 	}
 }
 
 // Copy returns a copy of the Store.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) Copy() *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ] {
 	// TODO: Decouple the KVStore type from the Cosmos-SDK.
 	cctx, _ := sdk.UnwrapSDKContext(kv.ctx).CacheContext()
@@ -243,20 +248,20 @@ func (kv *KVStore[
 // Context returns the context of the Store.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) WithContext(
 	ctx context.Context,
 ) *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ] {
 	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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) GetTotalActiveBalances(
 	slotsPerEpoch uint64,
 ) (math.Gwei, error) {

diff --git a/examples/berad/pkg/storage/staking.go b/examples/berad/pkg/storage/staking.go
@@ -23,17 +23,33 @@ package beacondb
 // GetEth1DepositIndex retrieves the eth1 deposit index from the beacon state.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) SetEth1DepositIndex(
 	index uint64,
 ) error {
 	return kv.eth1DepositIndex.Set(kv.ctx, index)
 }
+
+// GetWithdrawals retrieves the pending withdrawals from the beacon state.
+func (kv *KVStore[
+	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
+]) GetWithdrawals() (WithdrawalsT, error) {
+	return kv.withdrawals.Get(kv.ctx)
+}
+
+// SetWithdrawals overwrites the pending withdrawals from the beacon state.
+func (kv *KVStore[
+	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
+]) SetWithdrawals(withdrawals WithdrawalsT) error {
+	return kv.withdrawals.Set(kv.ctx, withdrawals)
+}
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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) SetGenesisValidatorsRoot(
 	root common.Root,
 ) error {
@@ -40,7 +40,7 @@ func (kv *KVStore[
 // beacon state.
 func (kv *KVStore[
 	BeaconBlockHeaderT, ExecutionPayloadHeaderT,
-	ForkT, ValidatorT, ValidatorsT, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) 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, WithdrawalT,
+	ForkT, ValidatorT, ValidatorsT, WithdrawalT, WithdrawalsT,
 ]) GetFork() (ForkT, error) {
 	return kv.fork.Get(kv.ctx)
 }