From b313b46f79d89674d662bee3c2d801ebb84f5abf Mon Sep 17 00:00:00 2001
From: terence tsao <terence@prysmaticlabs.com>
Date: Fri, 24 Jan 2020 10:15:01 -0800
Subject: [PATCH] Part 9 of proto array fork choice - get head (#4643)

---
 beacon-chain/forkchoice/protoarray/nodes.go   | 37 ++++++++++++
 .../forkchoice/protoarray/nodes_test.go       | 56 +++++++++++++++++++
 2 files changed, 93 insertions(+)

diff --git a/beacon-chain/forkchoice/protoarray/nodes.go b/beacon-chain/forkchoice/protoarray/nodes.go
index ca724a5a00..5231899798 100644
--- a/beacon-chain/forkchoice/protoarray/nodes.go
+++ b/beacon-chain/forkchoice/protoarray/nodes.go
@@ -4,12 +4,49 @@ import (
 	"bytes"
 	"context"
 	"errors"
+	"fmt"
 	"math"
 
 	"github.com/prysmaticlabs/prysm/shared/params"
 	"go.opencensus.io/trace"
 )
 
+// 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) {
+	ctx, span := trace.StartSpan(ctx, "protoArrayForkChoice.head")
+	defer span.End()
+
+	// Justified index has to be valid in node indices map, and can not be out of bound.
+	justifiedIndex, ok := s.nodeIndices[justifiedRoot]
+	if !ok {
+		return [32]byte{}, errUnknownJustifiedRoot
+	}
+	if justifiedIndex >= uint64(len(s.nodes)) {
+		return [32]byte{}, errInvalidJustifiedIndex
+	}
+
+	justifiedNode := s.nodes[justifiedIndex]
+	bestDescendantIndex := justifiedNode.bestDescendant
+	// If the justified node doesn't have a best descendent,
+	// the best node is itself.
+	if bestDescendantIndex == nonExistentNode {
+		bestDescendantIndex = justifiedIndex
+	}
+	if bestDescendantIndex >= uint64(len(s.nodes)) {
+		return [32]byte{}, errInvalidBestDescendantIndex
+	}
+
+	bestNode := s.nodes[bestDescendantIndex]
+	//
+	if !s.viableForHead(ctx, bestNode) {
+		return [32]byte{}, fmt.Errorf("after tree filter, best node can't be head, finalized epochs %d != %d, justified epoch %d != %d",
+			bestNode.finalizedEpoch, s.finalizedEpoch, bestNode.justifiedEpoch, s.justifiedEpoch)
+	}
+
+	return bestNode.root, nil
+}
+
 // insert registers a new block node to the fork choice store's node list.
 // It then updates the new node's parent with best child and descendant node.
 func (s *Store) insert(ctx context.Context,
diff --git a/beacon-chain/forkchoice/protoarray/nodes_test.go b/beacon-chain/forkchoice/protoarray/nodes_test.go
index dd2b535599..6a08d30946 100644
--- a/beacon-chain/forkchoice/protoarray/nodes_test.go
+++ b/beacon-chain/forkchoice/protoarray/nodes_test.go
@@ -5,6 +5,62 @@ import (
 	"testing"
 )
 
+func TestStore_Head_UnknownJustifiedRoot(t *testing.T) {
+	s := &Store{nodeIndices: make(map[[32]byte]uint64)}
+
+	if _, err := s.head(context.Background(), [32]byte{}); err.Error() != errUnknownJustifiedRoot.Error() {
+		t.Fatal("Did not get wanted error")
+	}
+}
+
+func TestStore_Head_UnknownJustifiedIndex(t *testing.T) {
+	r := [32]byte{'A'}
+	indices := make(map[[32]byte]uint64)
+	indices[r] = 1
+	s := &Store{nodeIndices: indices}
+
+	if _, err := s.head(context.Background(), r); err.Error() != errInvalidJustifiedIndex.Error() {
+		t.Fatal("Did not get wanted error")
+	}
+}
+
+func TestStore_Head_Itself(t *testing.T) {
+	r := [32]byte{'A'}
+	indices := make(map[[32]byte]uint64)
+	indices[r] = 0
+
+	// Since the justified node does not have a best descendant so the best node
+	// is itself.
+	s := &Store{nodeIndices: indices, nodes: []*Node{{root: r, bestDescendant: nonExistentNode}}}
+	h, err := s.head(context.Background(), r)
+	if err != nil {
+		t.Fatal("Did not get wanted error")
+	}
+
+	if h != r {
+		t.Error("Did not get wanted head")
+	}
+}
+
+func TestStore_Head_BestDescendant(t *testing.T) {
+	r := [32]byte{'A'}
+	best := [32]byte{'B'}
+	indices := make(map[[32]byte]uint64)
+	indices[r] = 0
+
+	// Since the justified node's best descendent is at index 1 and it's root is `best`,
+	// the head should be `best`.
+	s := &Store{nodeIndices: indices, nodes: []*Node{{root: r, bestDescendant: 1}, {root: best}}}
+	h, err := s.head(context.Background(), r)
+	if err != nil {
+		t.Fatal("Did not get wanted error")
+	}
+
+	if h != best {
+		t.Error("Did not get wanted head")
+	}
+}
+
 func TestStore_Insert_UnknownParent(t *testing.T) {
 	// The new node does not have a parent.
 	s := &Store{nodeIndices: make(map[[32]byte]uint64)}