feat(node-api): Optimizations usage of math.U64 (#1818)

berachain · Jul 30, 2024 · a9255e3 · a9255e3
1 parent ee0d018
commit a9255e3
Show file tree

Hide file tree

Showing 17 changed files with 91 additions and 115 deletions.
diff --git a/mod/node-api/backend/backend.go b/mod/node-api/backend/backend.go
@@ -136,15 +136,15 @@ func (b *Backend[
 
 func (b *Backend[
 	_, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _,
-]) GetSlotByRoot(root [32]byte) (uint64, error) {
+]) GetSlotByRoot(root [32]byte) (math.Slot, error) {
 	return b.sb.BlockStore().GetSlotByRoot(root)
 }
 
 // stateFromSlot returns the state at the given slot, after also processing the
 // next slot to ensure the returned beacon state is up to date.
 func (b *Backend[
 	_, _, _, _, BeaconStateT, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _,
-]) stateFromSlot(slot uint64) (BeaconStateT, uint64, error) {
+]) stateFromSlot(slot math.Slot) (BeaconStateT, math.Slot, error) {
 	var (
 		st  BeaconStateT
 		err error
@@ -154,13 +154,13 @@ func (b *Backend[
 	}
 
 	// Process the slot to update the latest state and block roots.
-	if _, err = b.sp.ProcessSlots(st, math.U64(slot+1)); err != nil {
+	if _, err = b.sp.ProcessSlots(st, slot+1); err != nil {
 		return st, slot, err
 	}
 
 	// We need to set the slot on the state back since ProcessSlot will update
 	// it to slot + 1.
-	err = st.SetSlot(math.Slot(slot))
+	err = st.SetSlot(slot)
 	return st, slot, err
 }
 
@@ -169,7 +169,7 @@ func (b *Backend[
 // slot on the beacon state.
 func (b *Backend[
 	_, _, _, _, BeaconStateT, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _,
-]) stateFromSlotRaw(slot uint64) (BeaconStateT, uint64, error) {
+]) stateFromSlotRaw(slot math.Slot) (BeaconStateT, math.Slot, error) {
 	var st BeaconStateT
 	//#nosec:G701 // not an issue in practice.
 	queryCtx, err := b.node.CreateQueryContext(int64(slot), false)
@@ -181,12 +181,10 @@ func (b *Backend[
 	// If using height 0 for the query context, make sure to return the latest
 	// slot.
 	if slot == 0 {
-		var latestSlot math.U64
-		latestSlot, err = st.GetSlot()
+		slot, err = st.GetSlot()
 		if err != nil {
 			return st, slot, err
 		}
-		slot = latestSlot.Unwrap()
 	}
 	return st, slot, err
 }
diff --git a/mod/node-api/backend/block.go b/mod/node-api/backend/block.go
@@ -23,15 +23,14 @@ package backend
 import (
 	types "github.com/berachain/beacon-kit/mod/node-api/handlers/beacon/types"
 	"github.com/berachain/beacon-kit/mod/primitives/pkg/common"
+	"github.com/berachain/beacon-kit/mod/primitives/pkg/math"
 )
 
 // BlockHeader returns the block header at the given slot.
 func (b Backend[
-	_, _, _, BeaconBlockHeaderT,
-	_, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _,
-]) BlockHeaderAtSlot(
-	slot uint64,
-) (BeaconBlockHeaderT, error) {
+	_, _, _, BeaconBlockHeaderT, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _,
+	_,
+]) BlockHeaderAtSlot(slot math.Slot) (BeaconBlockHeaderT, error) {
 	var blockHeader BeaconBlockHeaderT
 
 	st, _, err := b.stateFromSlot(slot)
@@ -46,26 +45,21 @@ func (b Backend[
 // GetBlockRoot returns the root of the block at the given stateID.
 func (b Backend[
 	_, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _,
-]) BlockRootAtSlot(
-	slot uint64,
-) (common.Root, error) {
+]) BlockRootAtSlot(slot math.Slot) (common.Root, error) {
 	st, slot, err := b.stateFromSlot(slot)
 	if err != nil {
 		return common.Root{}, err
 	}
 
 	// As calculated by the beacon chain. Ideally, this logic
 	// should be abstracted by the beacon chain.
-	index := slot % b.cs.SlotsPerHistoricalRoot()
-	return st.GetBlockRootAtIndex(index)
+	return st.GetBlockRootAtIndex(slot.Unwrap() % b.cs.SlotsPerHistoricalRoot())
 }
 
 // TODO: Implement this.
 func (b Backend[
 	_, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _,
-]) BlockRewardsAtSlot(
-	_ uint64,
-) (*types.BlockRewardsData, error) {
+]) BlockRewardsAtSlot(math.Slot) (*types.BlockRewardsData, error) {
 	return &types.BlockRewardsData{
 		ProposerIndex:     1,
 		Total:             1,

diff --git a/mod/node-api/backend/genesis.go b/mod/node-api/backend/genesis.go
@@ -20,14 +20,15 @@
 
 package backend
 
-import "github.com/berachain/beacon-kit/mod/primitives/pkg/common"
+import (
+	"github.com/berachain/beacon-kit/mod/primitives/pkg/common"
+	"github.com/berachain/beacon-kit/mod/primitives/pkg/math"
+)
 
 // GetGenesis returns the genesis state of the beacon chain.
 func (b Backend[
 	_, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _,
-]) GenesisValidatorsRoot(
-	slot uint64,
-) (common.Root, error) {
+]) GenesisValidatorsRoot(slot math.Slot) (common.Root, error) {
 	// needs genesis_time and gensis_fork_version
 	st, _, err := b.stateFromSlot(slot)
 	if err != nil {

diff --git a/mod/node-api/backend/beacon_test.go → mod/node-api/backend/genesis_test.go b/mod/node-api/backend/beacon_test.go → mod/node-api/backend/genesis_test.go
diff --git a/mod/node-api/backend/mocks/block_store.mock.go b/mod/node-api/backend/mocks/block_store.mock.go
diff --git a/mod/node-api/backend/randao.go b/mod/node-api/backend/randao.go
@@ -27,18 +27,15 @@ import (
 
 func (b Backend[
 	_, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _,
-]) RandaoAtEpoch(
-	slot, epoch uint64,
-) (common.Bytes32, error) {
+]) RandaoAtEpoch(slot math.Slot, epoch math.Epoch) (common.Bytes32, error) {
 	st, slot, err := b.stateFromSlot(slot)
 	if err != nil {
 		return common.Bytes32{}, err
 	}
 	// Infer the epoch if not provided.
 	if epoch == 0 {
-		latestEpoch := b.cs.SlotToEpoch(math.U64(slot))
-		epoch = latestEpoch.Unwrap()
+		epoch = b.cs.SlotToEpoch(slot)
 	}
-	index := epoch % b.cs.EpochsPerHistoricalVector()
+	index := epoch.Unwrap() % b.cs.EpochsPerHistoricalVector()
 	return st.GetRandaoMixAtIndex(index)
 }
diff --git a/mod/node-api/backend/state.go b/mod/node-api/backend/state.go
@@ -22,40 +22,36 @@ package backend
 
 import (
 	"github.com/berachain/beacon-kit/mod/primitives/pkg/common"
+	"github.com/berachain/beacon-kit/mod/primitives/pkg/math"
 )
 
 // StateFromSlotForProof returns the beacon state of the version that was used
 // to calculate the parent beacon block root, which has the empty state root in
 // the latest block header. Hence we do not process the next slot.
 func (b *Backend[
 	_, _, _, _, BeaconStateT, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _,
-]) StateFromSlotForProof(slot uint64) (BeaconStateT, uint64, error) {
+]) StateFromSlotForProof(slot math.Slot) (BeaconStateT, math.Slot, error) {
 	return b.stateFromSlotRaw(slot)
 }
 
 // GetStateRoot returns the root of the state at the given slot.
 func (b Backend[
 	_, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _,
-]) StateRootAtSlot(
-	slot uint64,
-) (common.Root, error) {
+]) StateRootAtSlot(slot math.Slot) (common.Root, error) {
 	st, slot, err := b.stateFromSlot(slot)
 	if err != nil {
 		return common.Root{}, err
 	}
 
 	// As calculated by the beacon chain. Ideally, this logic
 	// should be abstracted by the beacon chain.
-	index := slot % b.cs.SlotsPerHistoricalRoot()
-	return st.StateRootAtIndex(index)
+	return st.StateRootAtIndex(slot.Unwrap() % b.cs.SlotsPerHistoricalRoot())
 }
 
 // GetStateFork returns the fork of the state at the given stateID.
 func (b Backend[
 	_, _, _, _, _, _, _, _, _, _, _, _, _, ForkT, _, _, _, _, _, _, _,
-]) StateForkAtSlot(
-	slot uint64,
-) (ForkT, error) {
+]) StateForkAtSlot(slot math.Slot) (ForkT, error) {
 	var fork ForkT
 	st, _, err := b.stateFromSlot(slot)
 	if err != nil {

diff --git a/mod/node-api/backend/types.go b/mod/node-api/backend/types.go
@@ -85,7 +85,7 @@ type BlockStore[BeaconBlockT any] interface {
 	// Get retrieves the block at the given slot.
 	Get(slot uint64) (BeaconBlockT, error)
 	// GetSlotByRoot retrieves the slot by a given root from the store.
-	GetSlotByRoot(root [32]byte) (uint64, error)
+	GetSlotByRoot(root [32]byte) (math.Slot, error)
 	// Set sets the block at the given slot.
 	Set(slot uint64, block BeaconBlockT) error
 	// Prune prunes the block store of [start, end).

diff --git a/mod/node-api/backend/validator.go b/mod/node-api/backend/validator.go
@@ -29,8 +29,7 @@ import (
 func (b Backend[
 	_, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, ValidatorT, _, _, _,
 ]) ValidatorByID(
-	slot uint64,
-	id string,
+	slot math.Slot, id string,
 ) (*beacontypes.ValidatorData[ValidatorT], error) {
 	// TODO: to adhere to the spec, this shouldn't error if the error
 	// is not found, but i can't think of a way to do that without coupling
@@ -61,12 +60,11 @@ func (b Backend[
 	}, nil
 }
 
+// TODO: filter by status
 func (b Backend[
 	_, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, ValidatorT, _, _, _,
 ]) ValidatorsByIDs(
-	slot uint64,
-	ids []string,
-	_ []string, // TODO: filter by status
+	slot math.Slot, ids []string, _ []string,
 ) ([]*beacontypes.ValidatorData[ValidatorT], error) {
 	validatorsData := make([]*beacontypes.ValidatorData[ValidatorT], 0)
 	for _, id := range ids {
@@ -85,8 +83,7 @@ func (b Backend[
 func (b Backend[
 	_, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _,
 ]) ValidatorBalancesByIDs(
-	slot uint64,
-	ids []string,
+	slot math.Slot, ids []string,
 ) ([]*beacontypes.ValidatorBalanceData, error) {
 	var index math.U64
 	st, _, err := b.stateFromSlot(slot)

diff --git a/mod/node-api/handlers/beacon/backend.go b/mod/node-api/handlers/beacon/backend.go
@@ -23,6 +23,7 @@ package beacon
 import (
 	"github.com/berachain/beacon-kit/mod/node-api/handlers/beacon/types"
 	"github.com/berachain/beacon-kit/mod/primitives/pkg/common"
+	"github.com/berachain/beacon-kit/mod/primitives/pkg/math"
 )
 
 // Backend is the interface for backend of the beacon API.
@@ -33,63 +34,44 @@ type Backend[BlockHeaderT, ForkT, ValidatorT any] interface {
 	StateBackend[ForkT]
 	ValidatorBackend[ValidatorT]
 	HistoricalBackend[ForkT]
-	GetSlotByRoot(root [32]byte) (uint64, error)
+	GetSlotByRoot(root [32]byte) (math.Slot, error)
 }
 
 type GenesisBackend interface {
-	GenesisValidatorsRoot(
-		slot uint64,
-	) (common.Root, error)
+	GenesisValidatorsRoot(slot math.Slot) (common.Root, error)
 }
 
 type HistoricalBackend[ForkT any] interface {
-	StateRootAtSlot(
-		slot uint64,
-	) (common.Root, error)
-	StateForkAtSlot(
-		slot uint64,
-	) (ForkT, error)
+	StateRootAtSlot(slot math.Slot) (common.Root, error)
+	StateForkAtSlot(slot math.Slot) (ForkT, error)
 }
 
 type RandaoBackend interface {
-	RandaoAtEpoch(
-		slot, epoch uint64,
-	) (common.Bytes32, error)
+	RandaoAtEpoch(slot math.Slot, epoch math.Epoch) (common.Bytes32, error)
 }
 
 type BlockBackend[BeaconBlockHeaderT any] interface {
-	BlockRootAtSlot(
-		slot uint64,
-	) (common.Root, error)
-	BlockRewardsAtSlot(
-		slot uint64,
-	) (*types.BlockRewardsData, error)
-	BlockHeaderAtSlot(
-		slot uint64,
-	) (BeaconBlockHeaderT, error)
+	BlockRootAtSlot(slot math.Slot) (common.Root, error)
+	BlockRewardsAtSlot(slot math.Slot) (*types.BlockRewardsData, error)
+	BlockHeaderAtSlot(slot math.Slot) (BeaconBlockHeaderT, error)
 }
 
 type StateBackend[ForkT any] interface {
-	StateRootAtSlot(
-		slot uint64,
-	) (common.Root, error)
-	StateForkAtSlot(
-		slot uint64,
-	) (ForkT, error)
+	StateRootAtSlot(slot math.Slot) (common.Root, error)
+	StateForkAtSlot(slot math.Slot) (ForkT, error)
 }
 
 type ValidatorBackend[ValidatorT any] interface {
 	ValidatorByID(
-		slot uint64,
-		id string,
+		slot math.Slot, id string,
 	) (*types.ValidatorData[ValidatorT], error)
 	ValidatorsByIDs(
-		slot uint64,
+		slot math.Slot,
 		ids []string,
 		statuses []string,
 	) ([]*types.ValidatorData[ValidatorT], error)
 	ValidatorBalancesByIDs(
-		slot uint64,
+		slot math.Slot,
 		ids []string,
 	) ([]*types.ValidatorBalanceData, error)
 }
diff --git a/mod/node-api/handlers/beacon/header.go b/mod/node-api/handlers/beacon/header.go
@@ -21,8 +21,6 @@
 package beacon
 
 import (
-	"strconv"
-
 	beacontypes "github.com/berachain/beacon-kit/mod/node-api/handlers/beacon/types"
 	"github.com/berachain/beacon-kit/mod/node-api/handlers/utils"
 	"github.com/berachain/beacon-kit/mod/primitives/pkg/bytes"
@@ -37,7 +35,7 @@ func (h *Handler[
 	if err != nil {
 		return nil, err
 	}
-	slot, err := strconv.ParseUint(req.Slot, 10, 64)
+	slot, err := utils.U64FromString(req.Slot)
 	if err != nil {
 		return nil, err
 	}