Move chainheads to blockchain pkg

beacon-chain/blockchain/chain_info.go

@@ -48,6 +48,7 @@ type HeadFetcher interface {
 	HeadGenesisValidatorRoot() [32]byte
 	HeadETH1Data() *ethpb.Eth1Data
 	ProtoArrayStore() *protoarray.Store
+	ChainHeads() ([][32]byte, []types.Slot)
 }
 
 // ForkFetcher retrieves the current fork information of the Ethereum beacon chain.
@@ -259,3 +260,25 @@ func (s *Service) IsCanonical(ctx context.Context, blockRoot [32]byte) (bool, er
 	// If the block has not been finalized, check fork choice store to see if the block is canonical
 	return s.cfg.ForkChoiceStore.IsCanonical(blockRoot), nil
 }
+
+// ChainHeads returns all possible chain heads (leaves of fork choice tree).
+// Heads roots and heads slots are returned.
+func (s *Service) ChainHeads() ([][32]byte, []types.Slot) {
+	nodes := s.ProtoArrayStore().Nodes()
+
+	// 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)
+
+	nonExistentNode := ^uint64(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
+}

beacon-chain/blockchain/chain_info_test.go

@@ -282,3 +282,17 @@ func TestService_ProtoArrayStore(t *testing.T) {
 	p := c.ProtoArrayStore()
 	require.Equal(t, 0, int(p.FinalizedEpoch()))
 }
+
+func TestService_ChainHeads(t *testing.T) {
+	ctx := context.Background()
+	c := &Service{cfg: &Config{ForkChoiceStore: protoarray.New(0, 0, [32]byte{})}}
+	require.NoError(t, c.cfg.ForkChoiceStore.ProcessBlock(ctx, 100, [32]byte{'a'}, [32]byte{}, [32]byte{}, 0, 0))
+	require.NoError(t, c.cfg.ForkChoiceStore.ProcessBlock(ctx, 101, [32]byte{'b'}, [32]byte{'a'}, [32]byte{}, 0, 0))
+	require.NoError(t, c.cfg.ForkChoiceStore.ProcessBlock(ctx, 102, [32]byte{'c'}, [32]byte{'b'}, [32]byte{}, 0, 0))
+	require.NoError(t, c.cfg.ForkChoiceStore.ProcessBlock(ctx, 103, [32]byte{'d'}, [32]byte{}, [32]byte{}, 0, 0))
+	require.NoError(t, c.cfg.ForkChoiceStore.ProcessBlock(ctx, 104, [32]byte{'e'}, [32]byte{'b'}, [32]byte{}, 0, 0))
+
+	roots, slots := c.ChainHeads()
+	require.DeepEqual(t, [][32]byte{{'c'}, {'d'}, {'e'}}, roots)
+	require.DeepEqual(t, []types.Slot{102, 103, 104}, slots)
+}

beacon-chain/blockchain/testing/mock.go

@@ -382,3 +382,8 @@ func (s *ChainService) VerifyFinalizedConsistency(_ context.Context, r []byte) e
 	}
 	return nil
 }
+
+// ChainHeads mocks ChainHeads and always return nil.
+func (s *ChainService) ChainHeads() ([][32]byte, []types.Slot) {
+	return [][32]byte{}, []types.Slot{}
+}

beacon-chain/forkchoice/protoarray/store.go

@@ -241,29 +241,6 @@ 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) {

beacon-chain/forkchoice/protoarray/store_test.go

@@ -589,20 +589,3 @@ 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)
-}