Compute state up to slot (#5035)

beacon-chain/state/stategen/replay.go

@@ -16,6 +16,53 @@ import (
 	"go.opencensus.io/trace"
 )
 
+// ComputeStateUpToSlot returns a processed state up to input target slot.
+// If the last processed block is at slot 32, given input target slot at 40, this
+// returns processed state up to slot 40 via empty slots.
+// If there's duplicated blocks in a single slot, the canonical block will be returned.
+func (s *State) ComputeStateUpToSlot(ctx context.Context, targetSlot uint64) (*state.BeaconState, error) {
+	ctx, span := trace.StartSpan(ctx, "stateGen.ComputeStateUpToSlot")
+	defer span.End()
+
+	// Return genesis state if target slot is 0.
+	if targetSlot == 0 {
+		return s.beaconDB.GenesisState(ctx)
+	}
+
+	lastBlockRoot, lastBlockSlot, err := s.lastSavedBlock(ctx, targetSlot)
+	if err != nil {
+		return nil, errors.Wrap(err, "could not get last saved block")
+	}
+
+	lastBlockRootForState, err := s.lastSavedState(ctx, targetSlot)
+	if err != nil {
+		return nil, errors.Wrap(err, "could not get last valid state")
+	}
+	lastState, err := s.beaconDB.State(ctx, lastBlockRootForState)
+	if err != nil {
+		return nil, err
+	}
+	if lastState == nil {
+		return nil, errUnknownState
+	}
+
+	// Return if the last valid state's slot is higher than the target slot.
+	if lastState.Slot() >= targetSlot {
+		return lastState, nil
+	}
+
+	blks, err := s.LoadBlocks(ctx, lastState.Slot()+1, lastBlockSlot, lastBlockRoot)
+	if err != nil {
+		return nil, errors.Wrap(err, "could not load blocks")
+	}
+	lastState, err = s.ReplayBlocks(ctx, lastState, blks, targetSlot)
+	if err != nil {
+		return nil, errors.Wrap(err, "could not replay blocks")
+	}
+
+	return lastState, nil
+}
+
 // ReplayBlocks replays the input blocks on the input state until the target slot is reached.
 func (s *State) ReplayBlocks(ctx context.Context, state *state.BeaconState, signed []*ethpb.SignedBeaconBlock, targetSlot uint64) (*state.BeaconState, error) {
 	var err error

beacon-chain/state/stategen/replay_test.go

@@ -5,6 +5,7 @@ import (
 	"reflect"
 	"testing"
 
+	"github.com/gogo/protobuf/proto"
 	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
 	"github.com/prysmaticlabs/go-ssz"
 	"github.com/prysmaticlabs/prysm/beacon-chain/core/blocks"
@@ -17,6 +18,70 @@ import (
 	"github.com/prysmaticlabs/prysm/shared/testutil"
 )
 
+func TestComputeStateUpToSlot_GenesisState(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	service := New(db)
+
+	gBlk := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{}}
+	gRoot, err := ssz.HashTreeRoot(gBlk.Block)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := service.beaconDB.SaveBlock(ctx, gBlk); err != nil {
+		t.Fatal(err)
+	}
+	if err := service.beaconDB.SaveGenesisBlockRoot(ctx, gRoot); err != nil {
+		t.Fatal(err)
+	}
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	if err := service.beaconDB.SaveState(ctx, beaconState, gRoot); err != nil {
+		t.Fatal(err)
+	}
+
+	s, err := service.ComputeStateUpToSlot(ctx, 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !proto.Equal(s.InnerStateUnsafe(), beaconState.InnerStateUnsafe()) {
+		t.Error("Did not receive correct genesis state")
+	}
+}
+
+func TestComputeStateUpToSlot_CanProcessUpTo(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	service := New(db)
+
+	gBlk := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{}}
+	gRoot, err := ssz.HashTreeRoot(gBlk.Block)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := service.beaconDB.SaveBlock(ctx, gBlk); err != nil {
+		t.Fatal(err)
+	}
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	if err := service.beaconDB.SaveState(ctx, beaconState, gRoot); err != nil {
+		t.Fatal(err)
+	}
+
+	s, err := service.ComputeStateUpToSlot(ctx, params.BeaconConfig().SlotsPerEpoch+1)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if s.Slot() != params.BeaconConfig().SlotsPerEpoch+1 {
+		t.Log(s.Slot())
+		t.Error("Did not receive correct processed state")
+	}
+}
+
 func TestReplayBlocks_AllSkipSlots(t *testing.T) {
 	db := testDB.SetupDB(t)
 	defer testDB.TeardownDB(t, db)