diff --git a/beacon-chain/forkchoice/doubly-linked-tree/forkchoice.go b/beacon-chain/forkchoice/doubly-linked-tree/forkchoice.go index ce42b2fae..a46d1bca0 100644 --- a/beacon-chain/forkchoice/doubly-linked-tree/forkchoice.go +++ b/beacon-chain/forkchoice/doubly-linked-tree/forkchoice.go @@ -637,3 +637,15 @@ func (f *ForkChoice) Slot(root [32]byte) (primitives.Slot, error) { } return n.slot, nil } + +// TargetRoot returns the root of the checkpoint block for the given head root +func (f *ForkChoice) TargetRoot(root [32]byte) ([32]byte, error) { + n, ok := f.store.nodeByRoot[root] + if !ok || n == nil { + return [32]byte{}, ErrNilNode + } + if n.target == nil { + return [32]byte{}, nil + } + return n.target.root, nil +} diff --git a/beacon-chain/forkchoice/doubly-linked-tree/store.go b/beacon-chain/forkchoice/doubly-linked-tree/store.go index f50da85c9..79eb00f9a 100644 --- a/beacon-chain/forkchoice/doubly-linked-tree/store.go +++ b/beacon-chain/forkchoice/doubly-linked-tree/store.go @@ -75,7 +75,6 @@ func (s *Store) insert(ctx context.Context, } parent := s.nodeByRoot[parentRoot] - n := &Node{ slot: slot, root: root, @@ -89,6 +88,17 @@ func (s *Store) insert(ctx context.Context, timestamp: uint64(time.Now().Unix()), } + // Set the node's target checkpoint + if slot%params.BeaconConfig().SlotsPerEpoch == 0 { + n.target = n + } else if parent != nil { + if slots.ToEpoch(slot) == slots.ToEpoch(parent.slot) { + n.target = parent.target + } else { + n.target = parent + } + } + s.nodeByPayload[payloadHash] = n s.nodeByRoot[root] = n if parent == nil { @@ -145,6 +155,9 @@ func (s *Store) pruneFinalizedNodeByRootMap(ctx context.Context, node, finalized return ctx.Err() } if node == finalizedNode { + if node.target != node { + node.target = nil + } return nil } for _, child := range node.children { diff --git a/beacon-chain/forkchoice/doubly-linked-tree/store_test.go b/beacon-chain/forkchoice/doubly-linked-tree/store_test.go index 48b2f71f3..8f05a197d 100644 --- a/beacon-chain/forkchoice/doubly-linked-tree/store_test.go +++ b/beacon-chain/forkchoice/doubly-linked-tree/store_test.go @@ -435,3 +435,69 @@ func TestForkChoice_ReceivedBlocksLastEpoch(t *testing.T) { require.NoError(t, err) require.Equal(t, uint64(0), count) } + +func TestStore_Target(t *testing.T) { + ctx := context.Background() + f := setup(1, 1) + + state, blkRoot, err := prepareForkchoiceState(ctx, params.BeaconConfig().SlotsPerEpoch, [32]byte{'a'}, params.BeaconConfig().ZeroHash, params.BeaconConfig().ZeroHash, 1, 1) + require.NoError(t, err) + require.NoError(t, f.InsertNode(ctx, state, blkRoot)) + target, err := f.TargetRoot(blkRoot) + require.NoError(t, err) + require.Equal(t, target, blkRoot) + + state, root1, err := prepareForkchoiceState(ctx, params.BeaconConfig().SlotsPerEpoch+1, [32]byte{'b'}, blkRoot, params.BeaconConfig().ZeroHash, 1, 1) + require.NoError(t, err) + require.NoError(t, f.InsertNode(ctx, state, root1)) + target, err = f.TargetRoot(root1) + require.NoError(t, err) + require.Equal(t, target, blkRoot) + + // Insert a block for the next epoch (missed slot 0) + + state, root2, err := prepareForkchoiceState(ctx, 2*params.BeaconConfig().SlotsPerEpoch+1, [32]byte{'c'}, root1, params.BeaconConfig().ZeroHash, 1, 1) + require.NoError(t, err) + require.NoError(t, f.InsertNode(ctx, state, root2)) + target, err = f.TargetRoot(root2) + require.NoError(t, err) + require.Equal(t, target, root1) + + state, root3, err := prepareForkchoiceState(ctx, 2*params.BeaconConfig().SlotsPerEpoch+2, [32]byte{'d'}, root2, params.BeaconConfig().ZeroHash, 1, 1) + require.NoError(t, err) + require.NoError(t, f.InsertNode(ctx, state, root3)) + target, err = f.TargetRoot(root2) + require.NoError(t, err) + require.Equal(t, target, root1) + + // Prune finalization + s := f.store + s.finalizedCheckpoint.Root = root1 + require.NoError(t, s.prune(ctx)) + target, err = f.TargetRoot(root1) + require.NoError(t, err) + require.Equal(t, [32]byte{}, target) + + // Insert a block for next epoch (slot 0 present) + + state, root4, err := prepareForkchoiceState(ctx, 3*params.BeaconConfig().SlotsPerEpoch, [32]byte{'e'}, root1, params.BeaconConfig().ZeroHash, 1, 1) + require.NoError(t, err) + require.NoError(t, f.InsertNode(ctx, state, root4)) + target, err = f.TargetRoot(root4) + require.NoError(t, err) + require.Equal(t, target, root4) + + state, root5, err := prepareForkchoiceState(ctx, 3*params.BeaconConfig().SlotsPerEpoch+1, [32]byte{'f'}, root4, params.BeaconConfig().ZeroHash, 1, 1) + require.NoError(t, err) + require.NoError(t, f.InsertNode(ctx, state, root5)) + target, err = f.TargetRoot(root5) + require.NoError(t, err) + require.Equal(t, target, root4) + + // Prune finalization + s.finalizedCheckpoint.Root = root4 + require.NoError(t, s.prune(ctx)) + target, err = f.TargetRoot(root4) + require.NoError(t, err) + require.Equal(t, root4, target) +} diff --git a/beacon-chain/forkchoice/doubly-linked-tree/types.go b/beacon-chain/forkchoice/doubly-linked-tree/types.go index e33f3a812..416884c33 100644 --- a/beacon-chain/forkchoice/doubly-linked-tree/types.go +++ b/beacon-chain/forkchoice/doubly-linked-tree/types.go @@ -50,6 +50,7 @@ type Node struct { root [fieldparams.RootLength]byte // root of the block converted to the node. payloadHash [fieldparams.RootLength]byte // payloadHash of the block converted to the node. parent *Node // parent index of this node. + target *Node // target checkpoint for children []*Node // the list of direct children of this Node justifiedEpoch primitives.Epoch // justifiedEpoch of this node. unrealizedJustifiedEpoch primitives.Epoch // the epoch that would be justified if the block would be advanced to the next epoch. diff --git a/beacon-chain/forkchoice/interfaces.go b/beacon-chain/forkchoice/interfaces.go index 21c79dd85..1eb6dd013 100644 --- a/beacon-chain/forkchoice/interfaces.go +++ b/beacon-chain/forkchoice/interfaces.go @@ -69,6 +69,7 @@ type Getter interface { ShouldOverrideFCU() bool Slot([32]byte) (primitives.Slot, error) LastRoot(primitives.Epoch) [32]byte + TargetRoot([32]byte) ([32]byte, error) } // Setter allows to set forkchoice information