Load epoch boundary root (#5079)

* loadEpochBoundaryRoot

* Tests

* Span

Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>

beacon-chain/state/stategen/hot.go

@@ -7,6 +7,7 @@ import (
 	"github.com/prysmaticlabs/prysm/beacon-chain/core/helpers"
 	"github.com/prysmaticlabs/prysm/beacon-chain/state"
 	"github.com/prysmaticlabs/prysm/shared/bytesutil"
+	"github.com/prysmaticlabs/prysm/shared/params"
 	"go.opencensus.io/trace"
 )
 
@@ -102,3 +103,45 @@ func (s *State) loadHotStateBySlot(ctx context.Context, slot uint64) (*state.Bea
 
 	return s.ReplayBlocks(ctx, boundaryState, replayBlks, slot)
 }
+
+// This loads the epoch boundary root of a given state based on the state slot.
+// If the epoch boundary does not have a valid root, it then recovers by going
+// back to find the last slot before boundary which has a valid block.
+func (s *State) loadEpochBoundaryRoot(ctx context.Context, blockRoot [32]byte, state *state.BeaconState) ([32]byte, error) {
+	ctx, span := trace.StartSpan(ctx, "stateGen.loadEpochBoundaryRoot")
+	defer span.End()
+
+	boundarySlot := helpers.CurrentEpoch(state) * params.BeaconConfig().SlotsPerEpoch
+
+	// First checks if epoch boundary root already exists in cache.
+	r, ok := s.epochBoundarySlotToRoot[boundarySlot]
+	if ok {
+		return r, nil
+	}
+
+	// At epoch boundary, return the root which is just itself.
+	if state.Slot() == boundarySlot {
+		return blockRoot, nil
+	}
+
+	// Node uses genesis getters if the epoch boundary slot is genesis slot.
+	if boundarySlot == 0 {
+		r, err := s.genesisRoot(ctx)
+		if err != nil {
+			return [32]byte{}, nil
+		}
+		s.setEpochBoundaryRoot(boundarySlot, r)
+		return r, nil
+	}
+
+	// Now to find the epoch boundary root via DB.
+	r, _, err := s.lastSavedBlock(ctx, boundarySlot)
+	if err != nil {
+		return [32]byte{}, errors.Wrap(err, "could not get last saved block for epoch boundary root")
+	}
+
+	// Set the epoch boundary root cache.
+	s.setEpochBoundaryRoot(boundarySlot, r)
+
+	return r, nil
+}

beacon-chain/state/stategen/hot_test.go

@@ -5,8 +5,11 @@ import (
 	"testing"
 
 	"github.com/gogo/protobuf/proto"
+	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
+	"github.com/prysmaticlabs/go-ssz"
 	testDB "github.com/prysmaticlabs/prysm/beacon-chain/db/testing"
 	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
+	"github.com/prysmaticlabs/prysm/shared/params"
 	"github.com/prysmaticlabs/prysm/shared/testutil"
 )
 
@@ -136,3 +139,92 @@ func TestLoadHoteStateBySlot_CanAdvanceSlotUsingDB(t *testing.T) {
 		t.Error("Did not correctly load state")
 	}
 }
+
+func TestLoadEpochBoundaryRoot_Exists(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+	service := New(db)
+	r := [32]byte{'a'}
+	service.setEpochBoundaryRoot(params.BeaconConfig().SlotsPerEpoch, r)
+
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	boundaryRoot, err := service.loadEpochBoundaryRoot(ctx, r, beaconState)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if r != boundaryRoot {
+		t.Error("Did not correctly load boundary root")
+	}
+}
+
+func TestLoadEpochBoundaryRoot_SameSlot(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+	service := New(db)
+	r := [32]byte{'a'}
+
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	beaconState.SetSlot(params.BeaconConfig().SlotsPerEpoch)
+	boundaryRoot, err := service.loadEpochBoundaryRoot(ctx, r, beaconState)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if r != boundaryRoot {
+		t.Error("Did not correctly load boundary root")
+	}
+}
+
+func TestLoadEpochBoundaryRoot_Genesis(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+	service := New(db)
+	r := [32]byte{'a'}
+
+	b := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{}}
+	if err := db.SaveBlock(ctx, b); err != nil {
+		t.Fatal(err)
+	}
+	gRoot, _ := ssz.HashTreeRoot(b.Block)
+	if err := db.SaveGenesisBlockRoot(ctx, gRoot); err != nil {
+		t.Fatal(err)
+	}
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	beaconState.SetSlot(1)
+	boundaryRoot, err := service.loadEpochBoundaryRoot(ctx, r, beaconState)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if boundaryRoot != gRoot {
+		t.Error("Did not correctly load boundary root")
+	}
+}
+
+func TestLoadEpochBoundaryRoot_LastSavedBlock(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+	service := New(db)
+
+	b1 := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{Slot: service.lastArchivedSlot + 5}}
+	if err := service.beaconDB.SaveBlock(ctx, b1); err != nil {
+		t.Fatal(err)
+	}
+	b1Root, _ := ssz.HashTreeRoot(b1.Block)
+
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	beaconState.SetSlot(params.BeaconConfig().SlotsPerEpoch + 10)
+	boundaryRoot, err := service.loadEpochBoundaryRoot(ctx, [32]byte{}, beaconState)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if boundaryRoot != b1Root {
+		t.Error("Did not correctly load boundary root")
+	}
+}