diff --git a/beacon-chain/stategen/replay.go b/beacon-chain/stategen/replay.go index 85d36135d18..c0fc83af916 100644 --- a/beacon-chain/stategen/replay.go +++ b/beacon-chain/stategen/replay.go @@ -51,7 +51,17 @@ func (s *State) loadBlocks(ctx context.Context, startSlot uint64, endSlot uint64 } // The last retrieved block root has to match input end block root. + // Covers the edge case if there's multiple blocks on the same end slot, + // the end root may not be the last index in `blockRoots`. length := len(blocks) + for length >= 3 && blocks[length-1].Block.Slot == blocks[length-2].Block.Slot && blockRoots[length-1] != endBlockRoot { + length-- + if blockRoots[length-2] == endBlockRoot { + length-- + break + } + } + if blockRoots[length-1] != endBlockRoot { return nil, errors.New("end block roots don't match") } diff --git a/beacon-chain/stategen/replay_test.go b/beacon-chain/stategen/replay_test.go index 6c1028d80cc..692d807b461 100644 --- a/beacon-chain/stategen/replay_test.go +++ b/beacon-chain/stategen/replay_test.go @@ -209,6 +209,61 @@ func TestLoadBlocks_SameSlots(t *testing.T) { } } +func TestLoadBlocks_SameEndSlots(t *testing.T) { + db := testDB.SetupDB(t) + defer testDB.TeardownDB(t, db) + ctx := context.Background() + s := &State{ + beaconDB: db, + } + + roots, savedBlocks, err := tree3(db, []byte{'A'}) + if err != nil { + t.Fatal(err) + } + + filteredBlocks, err := s.loadBlocks(ctx, 0, 2, roots[2]) + if err != nil { + t.Fatal(err) + } + + wanted := []*ethpb.SignedBeaconBlock{ + {Block: savedBlocks[2]}, + {Block: savedBlocks[1]}, + {Block: savedBlocks[0]}, + } + if !reflect.DeepEqual(filteredBlocks, wanted) { + t.Error("Did not get wanted blocks") + } +} + +func TestLoadBlocks_SameEndSlotsWith2blocks(t *testing.T) { + db := testDB.SetupDB(t) + defer testDB.TeardownDB(t, db) + ctx := context.Background() + s := &State{ + beaconDB: db, + } + + roots, savedBlocks, err := tree4(db, []byte{'A'}) + if err != nil { + t.Fatal(err) + } + + filteredBlocks, err := s.loadBlocks(ctx, 0, 2, roots[1]) + if err != nil { + t.Fatal(err) + } + + wanted := []*ethpb.SignedBeaconBlock{ + {Block: savedBlocks[1]}, + {Block: savedBlocks[0]}, + } + if !reflect.DeepEqual(filteredBlocks, wanted) { + t.Error("Did not get wanted blocks") + } +} + func TestLoadBlocks_BadStart(t *testing.T) { db := testDB.SetupDB(t) defer testDB.TeardownDB(t, db) @@ -301,3 +356,73 @@ func tree2(db db.Database, genesisRoot []byte) ([][32]byte, []*ethpb.BeaconBlock } return [][32]byte{r0, r1, r21, r22, r23, r24, r3}, []*ethpb.BeaconBlock{b0, b1, b21, b22, b23, b24, b3}, nil } + +// tree3 constructs the following tree: +// B0 - B1 +// \- B2 +// \- B2 +// \- B2 +// \- B2 +func tree3(db db.Database, genesisRoot []byte) ([][32]byte, []*ethpb.BeaconBlock, error) { + b0 := ðpb.BeaconBlock{Slot: 0, ParentRoot: genesisRoot} + r0, _ := ssz.HashTreeRoot(b0) + b1 := ðpb.BeaconBlock{Slot: 1, ParentRoot: r0[:]} + r1, _ := ssz.HashTreeRoot(b1) + b21 := ðpb.BeaconBlock{Slot: 2, ParentRoot: r1[:], StateRoot: []byte{'A'}} + r21, _ := ssz.HashTreeRoot(b21) + b22 := ðpb.BeaconBlock{Slot: 2, ParentRoot: r1[:], StateRoot: []byte{'B'}} + r22, _ := ssz.HashTreeRoot(b22) + b23 := ðpb.BeaconBlock{Slot: 2, ParentRoot: r1[:], StateRoot: []byte{'C'}} + r23, _ := ssz.HashTreeRoot(b23) + b24 := ðpb.BeaconBlock{Slot: 2, ParentRoot: r1[:], StateRoot: []byte{'D'}} + r24, _ := ssz.HashTreeRoot(b24) + st, err := stateTrie.InitializeFromProtoUnsafe(&pb.BeaconState{}) + if err != nil { + return nil, nil, err + } + + for _, b := range []*ethpb.BeaconBlock{b0, b1, b21, b22, b23, b24} { + if err := db.SaveBlock(context.Background(), ðpb.SignedBeaconBlock{Block: b}); err != nil { + return nil, nil, err + } + if err := db.SaveState(context.Background(), st.Copy(), bytesutil.ToBytes32(b.ParentRoot)); err != nil { + return nil, nil, err + } + } + + return [][32]byte{r0, r1, r21, r22, r23, r24}, []*ethpb.BeaconBlock{b0, b1, b21, b22, b23, b24}, nil +} + +// tree4 constructs the following tree: +// B0 +// \- B2 +// \- B2 +// \- B2 +// \- B2 +func tree4(db db.Database, genesisRoot []byte) ([][32]byte, []*ethpb.BeaconBlock, error) { + b0 := ðpb.BeaconBlock{Slot: 0, ParentRoot: genesisRoot} + r0, _ := ssz.HashTreeRoot(b0) + b21 := ðpb.BeaconBlock{Slot: 2, ParentRoot: r0[:], StateRoot: []byte{'A'}} + r21, _ := ssz.HashTreeRoot(b21) + b22 := ðpb.BeaconBlock{Slot: 2, ParentRoot: r0[:], StateRoot: []byte{'B'}} + r22, _ := ssz.HashTreeRoot(b22) + b23 := ðpb.BeaconBlock{Slot: 2, ParentRoot: r0[:], StateRoot: []byte{'C'}} + r23, _ := ssz.HashTreeRoot(b23) + b24 := ðpb.BeaconBlock{Slot: 2, ParentRoot: r0[:], StateRoot: []byte{'D'}} + r24, _ := ssz.HashTreeRoot(b24) + st, err := stateTrie.InitializeFromProtoUnsafe(&pb.BeaconState{}) + if err != nil { + return nil, nil, err + } + + for _, b := range []*ethpb.BeaconBlock{b0, b21, b22, b23, b24} { + if err := db.SaveBlock(context.Background(), ðpb.SignedBeaconBlock{Block: b}); err != nil { + return nil, nil, err + } + if err := db.SaveState(context.Background(), st.Copy(), bytesutil.ToBytes32(b.ParentRoot)); err != nil { + return nil, nil, err + } + } + + return [][32]byte{r0, r21, r22, r23, r24}, []*ethpb.BeaconBlock{b0, b21, b22, b23, b24}, nil +}