perf(state-transition): reduce amount of withdrawals processed

.mockery.yaml

@@ -68,3 +68,8 @@ packages:
       recursive: False
       with-expecter: true
       all: True
+  github.com/berachain/beacon-kit/mod/state-transition/pkg/core:
+    config:
+      recursive: False
+      with-expecter: true
+      include-regex: ExecutionEngine

mod/state-transition/pkg/core/helpers_test.go

@@ -0,0 +1,138 @@
+// SPDX-License-Identifier: BUSL-1.1
+//
+// Copyright (C) 2024, Berachain Foundation. All rights reserved.
+// Use of this software is governed by the Business Source License included
+// in the LICENSE file of this repository and at www.mariadb.com/bsl11.
+//
+// ANY USE OF THE LICENSED WORK IN VIOLATION OF THIS LICENSE WILL AUTOMATICALLY
+// TERMINATE YOUR RIGHTS UNDER THIS LICENSE FOR THE CURRENT AND ALL OTHER
+// VERSIONS OF THE LICENSED WORK.
+//
+// THIS LICENSE DOES NOT GRANT YOU ANY RIGHT IN ANY TRADEMARK OR LOGO OF
+// LICENSOR OR ITS AFFILIATES (PROVIDED THAT YOU MAY USE A TRADEMARK OR LOGO OF
+// LICENSOR AS EXPRESSLY REQUIRED BY THIS LICENSE).
+//
+// TO THE EXTENT PERMITTED BY APPLICABLE LAW, THE LICENSED WORK IS PROVIDED ON
+// AN “AS IS” BASIS. LICENSOR HEREBY DISCLAIMS ALL WARRANTIES AND CONDITIONS,
+// EXPRESS OR IMPLIED, INCLUDING (WITHOUT LIMITATION) WARRANTIES OF
+// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, NON-INFRINGEMENT, AND
+// TITLE.
+
+package core_test
+
+import (
+	"context"
+	"fmt"
+
+	corestore "cosmossdk.io/core/store"
+	"cosmossdk.io/log"
+	"cosmossdk.io/store"
+	"cosmossdk.io/store/metrics"
+	storetypes "cosmossdk.io/store/types"
+	"github.com/berachain/beacon-kit/mod/consensus-types/pkg/types"
+	engineprimitives "github.com/berachain/beacon-kit/mod/engine-primitives/pkg/engine-primitives"
+	"github.com/berachain/beacon-kit/mod/node-core/pkg/components"
+	statedb "github.com/berachain/beacon-kit/mod/state-transition/pkg/core/state"
+	"github.com/berachain/beacon-kit/mod/storage/pkg/beacondb"
+	"github.com/berachain/beacon-kit/mod/storage/pkg/db"
+	"github.com/berachain/beacon-kit/mod/storage/pkg/encoding"
+	dbm "github.com/cosmos/cosmos-db"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
+
+type (
+	TestBeaconStateMarshallableT = types.BeaconState[
+		*types.BeaconBlockHeader,
+		*types.Eth1Data,
+		*types.ExecutionPayloadHeader,
+		*types.Fork,
+		*types.Validator,
+		types.BeaconBlockHeader,
+		types.Eth1Data,
+		types.ExecutionPayloadHeader,
+		types.Fork,
+		types.Validator,
+	]
+
+	TestKVStoreT = beacondb.KVStore[
+		*types.BeaconBlockHeader,
+		*types.Eth1Data,
+		*types.ExecutionPayloadHeader,
+		*types.Fork,
+		*types.Validator,
+		types.Validators,
+	]
+
+	TestBeaconStateT = statedb.StateDB[
+		*types.BeaconBlockHeader,
+		*TestBeaconStateMarshallableT,
+		*types.Eth1Data,
+		*types.ExecutionPayloadHeader,
+		*types.Fork,
+		*TestKVStoreT,
+		*types.Validator,
+		types.Validators,
+		*engineprimitives.Withdrawal,
+		types.WithdrawalCredentials,
+	]
+)
+
+type testKVStoreService struct {
+	ctx sdk.Context
+}
+
+func (kvs *testKVStoreService) OpenKVStore(context.Context) corestore.KVStore {
+	//nolint:contextcheck // fine with tests
+	return components.NewKVStore(
+		sdk.UnwrapSDKContext(kvs.ctx).KVStore(testStoreKey),
+	)
+}
+
+var (
+	testStoreKey = storetypes.NewKVStoreKey("state-transition-tests")
+	testCodec    = &encoding.SSZInterfaceCodec[*types.ExecutionPayloadHeader]{}
+)
+
+func initTestStore() (
+	*beacondb.KVStore[
+		*types.BeaconBlockHeader,
+		*types.Eth1Data,
+		*types.ExecutionPayloadHeader,
+		*types.Fork,
+		*types.Validator,
+		types.Validators,
+	], error) {
+	db, err := db.OpenDB("", dbm.MemDBBackend)
+	if err != nil {
+		return nil, fmt.Errorf("failed opening mem db: %w", err)
+	}
+	var (
+		nopLog     = log.NewNopLogger()
+		nopMetrics = metrics.NewNoOpMetrics()
+	)
+
+	cms := store.NewCommitMultiStore(
+		db,
+		nopLog,
+		nopMetrics,
+	)
+
+	ctx := sdk.NewContext(cms, true, nopLog)
+	cms.MountStoreWithDB(testStoreKey, storetypes.StoreTypeIAVL, nil)
+	if err = cms.LoadLatestVersion(); err != nil {
+		return nil, fmt.Errorf("failed to load latest version: %w", err)
+	}
+	testStoreService := &testKVStoreService{ctx: ctx}
+
+	return beacondb.New[
+		*types.BeaconBlockHeader,
+		*types.Eth1Data,
+		*types.ExecutionPayloadHeader,
+		*types.Fork,
+		*types.Validator,
+		types.Validators,
+	](
+		testStoreService,
+		testCodec,
+	), nil
+}

mod/state-transition/pkg/core/state/statedb.go

@@ -225,10 +225,6 @@ func (s *StateDB[
 
 	// Iterate through indices to find the next validators to withdraw.
 	for range bound {
-		var (
-			withdrawal WithdrawalT
-			amount     math.Gwei
-		)
 		validator, err = s.ValidatorByIndex(validatorIndex)
 		if err != nil {
 			return nil, err
@@ -247,21 +243,24 @@ func (s *StateDB[
 
 		// Set the amount of the withdrawal depending on the balance of the
 		// validator.
+		var withdrawal WithdrawalT
 		if validator.IsFullyWithdrawable(balance, epoch) {
-			amount = balance
+			withdrawals = append(withdrawals, withdrawal.New(
+				math.U64(withdrawalIndex),
+				validatorIndex,
+				withdrawalAddress,
+				balance,
+			))
 		} else if validator.IsPartiallyWithdrawable(
 			balance, math.Gwei(s.cs.MaxEffectiveBalance()),
 		) {
-			amount = balance - math.Gwei(s.cs.MaxEffectiveBalance())
+			withdrawals = append(withdrawals, withdrawal.New(
+				math.U64(withdrawalIndex),
+				validatorIndex,
+				withdrawalAddress,
+				balance-math.Gwei(s.cs.MaxEffectiveBalance()),
+			))
 		}
-		withdrawal = withdrawal.New(
-			math.U64(withdrawalIndex),
-			validatorIndex,
-			withdrawalAddress,
-			amount,
-		)
-
-		withdrawals = append(withdrawals, withdrawal)
 
 		// Increment the withdrawal index to process the next withdrawal.
 		withdrawalIndex++

mod/state-transition/pkg/core/state_processor.go

@@ -85,6 +85,10 @@ type StateProcessor[
 	executionEngine ExecutionEngine[
 		ExecutionPayloadT, ExecutionPayloadHeaderT, WithdrawalsT,
 	]
+
+	// processingGenesis allows initializing correctly
+	// eth1 deposit index upon genesis
+	processingGenesis bool
 }
 
 // NewStateProcessor creates a new state processor.

mod/state-transition/pkg/core/state_processor_genesis.go

@@ -47,6 +47,11 @@ func (sp *StateProcessor[
 	executionPayloadHeader ExecutionPayloadHeaderT,
 	genesisVersion common.Version,
 ) (transition.ValidatorUpdates, error) {
+	sp.processingGenesis = true
+	defer func() {
+		sp.processingGenesis = false
+	}()
+
 	var (
 		blkHeader BeaconBlockHeaderT
 		blkBody   BeaconBlockBodyT
@@ -67,25 +72,28 @@ func (sp *StateProcessor[
 		return nil, err
 	}
 
-	if err := st.SetEth1DepositIndex(0); err != nil {
-		return nil, err
-	}
+	// Eth1DepositIndex will be set in processDeposit
 
-	if err := st.SetEth1Data(eth1Data.New(
-		common.Root{},
-		0,
-		executionPayloadHeader.GetBlockHash(),
-	)); err != nil {
+	if err := st.SetEth1Data(
+		eth1Data.New(
+			common.Root{},
+			0,
+			executionPayloadHeader.GetBlockHash(),
+		)); err != nil {
 		return nil, err
 	}
 
 	// TODO: we need to handle common.Version vs
 	// uint32 better.
-	bodyRoot := blkBody.Empty(
-		version.ToUint32(genesisVersion)).HashTreeRoot()
-	if err := st.SetLatestBlockHeader(blkHeader.New(
-		0, 0, common.Root{}, common.Root{}, bodyRoot,
-	)); err != nil {
+	bodyRoot := blkBody.Empty(version.ToUint32(genesisVersion)).HashTreeRoot()
+	if err := st.SetLatestBlockHeader(
+		blkHeader.New(
+			0,             // slot
+			0,             // proposer index
+			common.Root{}, // parent block root
+			common.Root{}, // state root
+			bodyRoot,
+		)); err != nil {
 		return nil, err
 	}