diff --git a/beacon-chain/forkchoice/protoarray/store.go b/beacon-chain/forkchoice/protoarray/store.go index 505704bedb35..b7bf0989a73c 100644 --- a/beacon-chain/forkchoice/protoarray/store.go +++ b/beacon-chain/forkchoice/protoarray/store.go @@ -230,6 +230,29 @@ func (s *Store) NodesIndices() map[[32]byte]uint64 { return s.nodesIndices } +// ChainHeads returns all possible chain heads (leaves of fork choice tree). +// Heads roots and heads slots are returned. +func (s *Store) ChainHeads() ([][32]byte, []types.Slot) { + s.nodesLock.RLock() + nodes := s.Nodes() + s.nodesLock.RUnlock() + + // Deliberate choice to not preallocate space for below. + // Heads cant be more than 2-3 in the worst case where pre-allocation will be 64 to begin with. + headsRoots := make([][32]byte, 0) + headsSlots := make([]types.Slot, 0) + + for _, node := range nodes { + // Possible heads have no children. + if node.bestDescendant == NonExistentNode && node.bestChild == NonExistentNode { + headsRoots = append(headsRoots, node.root) + headsSlots = append(headsSlots, node.slot) + } + } + + return headsRoots, headsSlots +} + // head starts from justified root and then follows the best descendant links // to find the best block for head. func (s *Store) head(ctx context.Context, justifiedRoot [32]byte) ([32]byte, error) { diff --git a/beacon-chain/forkchoice/protoarray/store_test.go b/beacon-chain/forkchoice/protoarray/store_test.go index 4f50a4198b94..46d293aa3a52 100644 --- a/beacon-chain/forkchoice/protoarray/store_test.go +++ b/beacon-chain/forkchoice/protoarray/store_test.go @@ -589,3 +589,20 @@ func TestStore_UpdateCanonicalNodes_ContextCancelled(t *testing.T) { cancel() require.ErrorContains(t, "context canceled", f.store.updateCanonicalNodes(ctx, [32]byte{'c'})) } + +func TestStore_ChainHeads(t *testing.T) { + nodes := []*Node{ + {slot: 100, root: [32]byte{'a'}, bestChild: NonExistentNode, bestDescendant: NonExistentNode}, + {slot: 101, root: [32]byte{'b'}}, + {slot: 102, root: [32]byte{'c'}, bestDescendant: NonExistentNode}, + {slot: 103, root: [32]byte{'d'}, bestChild: NonExistentNode, bestDescendant: NonExistentNode}, + {slot: 104, root: [32]byte{'e'}, bestChild: NonExistentNode}, + } + + s := &Store{ + nodes: nodes, + } + roots, slots := s.ChainHeads() + require.DeepEqual(t, [][32]byte{{'a'}, {'d'}}, roots) + require.DeepEqual(t, []types.Slot{100, 103}, slots) +}