Part 5 of proto array fork choice - update best child and descendant (#…

…4629)

beacon-chain/forkchoice/protoarray/nodes.go

@@ -1,11 +1,116 @@
 package protoarray
 
 import (
+	"bytes"
 	"context"
 
 	"go.opencensus.io/trace"
 )
 
+// updateBestChildAndDescendant updates parent node's best child and descendent.
+// It looks at input parent node and input child node and potentially modifies parent's best
+// child and best descendent indices.
+// There are four outcomes:
+// 1.)  The child is already the best child but it's now invalid due to a FFG change and should be removed.
+// 2.)  The child is already the best child and the parent is updated with the new best descendant.
+// 3.)  The child is not the best child but becomes the best child.
+// 4.)  The child is not the best child and does not become best child.
+func (s *Store) updateBestChildAndDescendant(ctx context.Context, parentIndex uint64, childIndex uint64) error {
+	ctx, span := trace.StartSpan(ctx, "protoArrayForkChoice.updateBestChildAndDescendant")
+	defer span.End()
+
+	// Protection against parent index out of bound, this should not happen.
+	if parentIndex >= uint64(len(s.nodes)) {
+		return errInvalidNodeIndex
+	}
+	parent := s.nodes[parentIndex]
+
+	// Protection against child index out of bound, again this should not happen.
+	if childIndex >= uint64(len(s.nodes)) {
+		return errInvalidNodeIndex
+	}
+	child := s.nodes[childIndex]
+
+	// Is the child viable to become head? Based on justification and finalization rules.
+	childLeadsToViableHead, err := s.leadsToViableHead(ctx, child)
+	if err != nil {
+		return err
+	}
+
+	// Define 3 variables for the 3 outcomes mentioned above. This is to
+	// set `parent.bestChild` and `parent.bestDescendent` to. These
+	// aliases are to assist readability.
+	changeToNone := []uint64{nonExistentNode, nonExistentNode}
+	bestDescendant := child.bestDescendant
+	if bestDescendant == nonExistentNode {
+		bestDescendant = childIndex
+	}
+	changeToChild := []uint64{childIndex, bestDescendant}
+	noChange := []uint64{parent.bestChild, parent.bestDescendant}
+	newParentChild := make([]uint64, 0)
+
+	if parent.bestChild != nonExistentNode {
+		if parent.bestChild == childIndex && !childLeadsToViableHead {
+			// If the child is already the best child of the parent but it's not viable for head,
+			// we should remove it. (Outcome 1)
+			newParentChild = changeToNone
+		} else if parent.bestChild == childIndex {
+			// If the child is already the best child of the parent, set it again to ensure best
+			// descendent of the parent is updated. (Outcome 2)
+			newParentChild = changeToChild
+		} else {
+			// Protection against parent's best child going out of bound.
+			if parent.bestChild > uint64(len(s.nodes)) {
+				return errInvalidBestDescendantIndex
+			}
+			bestChild := s.nodes[parent.bestChild]
+			// Is current parent's best child viable to be head? Based on justification and finalization rules.
+			bestChildLeadsToViableHead, err := s.leadsToViableHead(ctx, bestChild)
+			if err != nil {
+				return err
+			}
+
+			if childLeadsToViableHead && !bestChildLeadsToViableHead {
+				// The child leads to a viable head, but the current parent's best child doesnt.
+				newParentChild = changeToChild
+			} else if !childLeadsToViableHead && bestChildLeadsToViableHead {
+				// The child doesn't lead to a viable head, the current parent's best child does.
+				newParentChild = noChange
+			} else if child.weight == bestChild.weight {
+				// If both are viable, compare their weights.
+				// Tie-breaker of equal weights by root.
+				if bytes.Compare(child.root[:], bestChild.root[:]) > 0 {
+					newParentChild = changeToChild
+				} else {
+					newParentChild = noChange
+				}
+			} else {
+				// Choose winner by weight.
+				if child.weight > bestChild.weight {
+					newParentChild = changeToChild
+				} else {
+					newParentChild = noChange
+				}
+			}
+		}
+	} else {
+		if childLeadsToViableHead {
+			// If parent doesn't have a best child and the child is viable.
+			newParentChild = changeToChild
+		} else {
+			// If parent doesn't have a best child and the child is not viable.
+			newParentChild = noChange
+		}
+	}
+
+	// Update parent with the outcome.
+	parent.bestChild = newParentChild[0]
+	parent.bestDescendant = newParentChild[1]
+	s.nodes[parentIndex] = parent
+
+	return nil
+}
+
 // leadsToViableHead returns true if the node or the best descendent of the node is viable for head.
 // Any node with diff finalized or justified epoch than the ones in fork choice store
 // should not be viable to head.

beacon-chain/forkchoice/protoarray/nodes_test.go

@@ -5,7 +5,177 @@ import (
 	"testing"
 )
 
-func TestStore_leadsToViableHead(t *testing.T) {
+func TestStore_UpdateBestChildAndDescendant_RemoveChild(t *testing.T) {
+	// Make parent's best child equal's to input child index and child is not viable.
+	s := &Store{nodes: []*Node{{bestChild: 1}, {}}, justifiedEpoch: 1, finalizedEpoch: 1}
+
+	if err := s.updateBestChildAndDescendant(context.Background(), 0, 1); err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify parent's best child and best descendant are `none`.
+	if s.nodes[0].bestChild != nonExistentNode {
+		t.Error("Did not get correct best child index")
+	}
+	if s.nodes[0].bestDescendant != nonExistentNode {
+		t.Error("Did not get correct best descendant index")
+	}
+}
+
+func TestStore_UpdateBestChildAndDescendant_UpdateDescendant(t *testing.T) {
+	// Make parent's best child equal to child index and child is viable.
+	s := &Store{nodes: []*Node{{bestChild: 1}, {bestDescendant: nonExistentNode}}}
+
+	if err := s.updateBestChildAndDescendant(context.Background(), 0, 1); err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify parent's best child is the same and best descendant is not set to child index.
+	if s.nodes[0].bestChild != 1 {
+		t.Error("Did not get correct best child index")
+	}
+	if s.nodes[0].bestDescendant != 1 {
+		t.Error("Did not get correct best descendant index")
+	}
+}
+
+func TestStore_UpdateBestChildAndDescendant_ChangeChildByViability(t *testing.T) {
+	// Make parent's best child not equal to child index, child leads to viable index and
+	// parents best child doesnt lead to viable index.
+	s := &Store{
+		justifiedEpoch: 1,
+		finalizedEpoch: 1,
+		nodes: []*Node{{bestChild: 1, justifiedEpoch: 1, finalizedEpoch: 1},
+			{bestDescendant: nonExistentNode},
+			{bestDescendant: nonExistentNode, justifiedEpoch: 1, finalizedEpoch: 1}}}
+
+	if err := s.updateBestChildAndDescendant(context.Background(), 0, 2); err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify parent's best child and best descendant are set to child index.
+	if s.nodes[0].bestChild != 2 {
+		t.Error("Did not get correct best child index")
+	}
+	if s.nodes[0].bestDescendant != 2 {
+		t.Error("Did not get correct best descendant index")
+	}
+}
+
+func TestStore_UpdateBestChildAndDescendant_ChangeChildByWeight(t *testing.T) {
+	// Make parent's best child not equal to child index, child leads to viable index and
+	// parents best child leads to viable index but child has more weight than parent's best child.
+	s := &Store{
+		justifiedEpoch: 1,
+		finalizedEpoch: 1,
+		nodes: []*Node{{bestChild: 1, justifiedEpoch: 1, finalizedEpoch: 1},
+			{bestDescendant: nonExistentNode, justifiedEpoch: 1, finalizedEpoch: 1},
+			{bestDescendant: nonExistentNode, justifiedEpoch: 1, finalizedEpoch: 1, weight: 1}}}
+
+	if err := s.updateBestChildAndDescendant(context.Background(), 0, 2); err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify parent's best child and best descendant are set to child index.
+	if s.nodes[0].bestChild != 2 {
+		t.Error("Did not get correct best child index")
+	}
+	if s.nodes[0].bestDescendant != 2 {
+		t.Error("Did not get correct best descendant index")
+	}
+}
+
+func TestStore_UpdateBestChildAndDescendant_ChangeChildAtLeaf(t *testing.T) {
+	// Make parent's best child to none and input child leads to viable index.
+	s := &Store{
+		justifiedEpoch: 1,
+		finalizedEpoch: 1,
+		nodes: []*Node{{bestChild: nonExistentNode, justifiedEpoch: 1, finalizedEpoch: 1},
+			{bestDescendant: nonExistentNode, justifiedEpoch: 1, finalizedEpoch: 1},
+			{bestDescendant: nonExistentNode, justifiedEpoch: 1, finalizedEpoch: 1}}}
+
+	if err := s.updateBestChildAndDescendant(context.Background(), 0, 2); err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify parent's best child and best descendant are set to child index.
+	if s.nodes[0].bestChild != 2 {
+		t.Error("Did not get correct best child index")
+	}
+	if s.nodes[0].bestDescendant != 2 {
+		t.Error("Did not get correct best descendant index")
+	}
+}
+
+func TestStore_UpdateBestChildAndDescendant_NoChangeByViability(t *testing.T) {
+	// Make parent's best child not equal to child index, child leads to not viable index and
+	// parents best child leads to viable index.
+	s := &Store{
+		justifiedEpoch: 1,
+		finalizedEpoch: 1,
+		nodes: []*Node{{bestChild: 1, justifiedEpoch: 1, finalizedEpoch: 1},
+			{bestDescendant: nonExistentNode, justifiedEpoch: 1, finalizedEpoch: 1},
+			{bestDescendant: nonExistentNode}}}
+
+	if err := s.updateBestChildAndDescendant(context.Background(), 0, 2); err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify parent's best child and best descendant are not changed.
+	if s.nodes[0].bestChild != 1 {
+		t.Error("Did not get correct best child index")
+	}
+	if s.nodes[0].bestDescendant != 0 {
+		t.Error("Did not get correct best descendant index")
+	}
+}
+
+func TestStore_UpdateBestChildAndDescendant_NoChangeByWeight(t *testing.T) {
+	// Make parent's best child not equal to child index, child leads to viable index and
+	// parents best child leads to viable index but parent's best child has more weight.
+	s := &Store{
+		justifiedEpoch: 1,
+		finalizedEpoch: 1,
+		nodes: []*Node{{bestChild: 1, justifiedEpoch: 1, finalizedEpoch: 1},
+			{bestDescendant: nonExistentNode, justifiedEpoch: 1, finalizedEpoch: 1, weight: 1},
+			{bestDescendant: nonExistentNode, justifiedEpoch: 1, finalizedEpoch: 1}}}
+
+	if err := s.updateBestChildAndDescendant(context.Background(), 0, 2); err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify parent's best child and best descendant are not changed.
+	if s.nodes[0].bestChild != 1 {
+		t.Error("Did not get correct best child index")
+	}
+	if s.nodes[0].bestDescendant != 0 {
+		t.Error("Did not get correct best descendant index")
+	}
+}
+
+func TestStore_UpdateBestChildAndDescendant_NoChangeAtLeaf(t *testing.T) {
+	// Make parent's best child to none and input child does not lead to viable index.
+	s := &Store{
+		justifiedEpoch: 1,
+		finalizedEpoch: 1,
+		nodes: []*Node{{bestChild: nonExistentNode, justifiedEpoch: 1, finalizedEpoch: 1},
+			{bestDescendant: nonExistentNode, justifiedEpoch: 1, finalizedEpoch: 1},
+			{bestDescendant: nonExistentNode}}}
+
+	if err := s.updateBestChildAndDescendant(context.Background(), 0, 2); err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify parent's best child and best descendant are not changed.
+	if s.nodes[0].bestChild != nonExistentNode {
+		t.Error("Did not get correct best child index")
+	}
+	if s.nodes[0].bestDescendant != 0 {
+		t.Error("Did not get correct best descendant index")
+	}
+}
+
+func TestStore_LeadsToViableHead(t *testing.T) {
 	tests := []struct {
 		n              *Node
 		justifiedEpoch uint64
@@ -35,7 +205,7 @@ func TestStore_leadsToViableHead(t *testing.T) {
 	}
 }
 
-func TestStore_viableForHead(t *testing.T) {
+func TestStore_ViableForHead(t *testing.T) {
 	tests := []struct {
 		n              *Node
 		justifiedEpoch uint64