Part 6 of proto array fork choice - update weight (#4636)

beacon-chain/forkchoice/protoarray/nodes.go

@@ -3,10 +3,75 @@ package protoarray
 import (
 	"bytes"
 	"context"
+	"math"
 
+	"github.com/prysmaticlabs/prysm/shared/params"
 	"go.opencensus.io/trace"
 )
 
+// applyWeightChanges iterates backwards through the nodes in store. It checks all nodes parent
+// and its best child. For each node, it updates the weight with input delta and
+// back propagate the nodes delta to its parents delta. After scoring changes,
+// the best child is then updated along with best descendant.
+func (s *Store) applyWeightChanges(ctx context.Context, justifiedEpoch uint64, finalizedEpoch uint64, delta []int) error {
+	ctx, span := trace.StartSpan(ctx, "protoArrayForkChoice.applyWeightChanges")
+	defer span.End()
+
+	// The length of the nodes can not be different than length of the delta.
+	if len(s.nodes) != len(delta) {
+		return errInvalidDeltaLength
+	}
+
+	// Update the justified / finalized epochs in store if necessary.
+	if s.justifiedEpoch != justifiedEpoch || s.finalizedEpoch != finalizedEpoch {
+		s.justifiedEpoch = justifiedEpoch
+		s.finalizedEpoch = finalizedEpoch
+	}
+
+	// Iterate backwards through all index to node in store.
+	for i := len(s.nodes) - 1; i >= 0; i-- {
+		n := s.nodes[i]
+
+		// There is no need to adjust the balances or manage parent of the zero hash, it
+		// is an alias to the genesis block.
+		if n.root == params.BeaconConfig().ZeroHash {
+			continue
+		}
+
+		nodeDelta := delta[i]
+
+		if nodeDelta < 0 {
+			// A node's weight can not be negative but the delta can be negative.
+			if int(n.weight)+nodeDelta < 0 {
+				n.weight = 0
+			} else {
+				// Subtract node's weight.
+				n.weight -= uint64(math.Abs(float64(nodeDelta)))
+			}
+		} else {
+			// Add node's weight.
+			n.weight += uint64(nodeDelta)
+		}
+
+		s.nodes[i] = n
+
+		// Update parent's best child and descendent if the node has a known parent.
+		if n.parent != nonExistentNode {
+			// Protection against node parent index out of bound. This should not happen.
+			if int(n.parent) >= len(delta) {
+				return errInvalidParentDelta
+			}
+			// Back propagate the nodes delta to its parent.
+			delta[n.parent] += nodeDelta
+			if err := s.updateBestChildAndDescendant(ctx, n.parent, uint64(i)); err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}
+
 // 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.

beacon-chain/forkchoice/protoarray/nodes_test.go

@@ -5,6 +5,102 @@ import (
 	"testing"
 )
 
+func TestStore_ApplyScoreChanges_InvalidDeltaLength(t *testing.T) {
+	s := &Store{}
+
+	// This will fail because node indices has length of 0, and delta list has a length of 1.
+	if err := s.applyWeightChanges(context.Background(), 0, 0, []int{1}); err.Error() != errInvalidDeltaLength.Error() {
+		t.Error("Did not get wanted error")
+	}
+}
+
+func TestStore_ApplyScoreChanges_UpdateEpochs(t *testing.T) {
+	s := &Store{}
+
+	// The justified and finalized epochs in Store should be updated to 1 and 1 given the following input.
+	if err := s.applyWeightChanges(context.Background(), 1, 1, []int{}); err != nil {
+		t.Error("Did not get wanted error")
+	}
+
+	if s.justifiedEpoch != 1 {
+		t.Error("Did not update justified epoch")
+	}
+	if s.finalizedEpoch != 1 {
+		t.Error("Did not update justified epoch")
+	}
+}
+
+func TestStore_ApplyScoreChanges_UpdateWeightsPositiveDelta(t *testing.T) {
+	// Construct 3 nodes with weight 100 on each node. The 3 nodes linked to each other.
+	s := &Store{nodes: []*Node{
+		{root: [32]byte{'A'}, weight: 100},
+		{root: [32]byte{'A'}, weight: 100},
+		{parent: 1, root: [32]byte{'A'}, weight: 100}}}
+
+	// Each node gets one unique vote. The weight should look like 103 <- 102 <- 101 because
+	// they get propagated back.
+	if err := s.applyWeightChanges(context.Background(), 0, 0, []int{1, 1, 1}); err != nil {
+		t.Fatal(err)
+	}
+
+	if s.nodes[0].weight != 103 {
+		t.Error("Did not get correct weight")
+	}
+	if s.nodes[1].weight != 102 {
+		t.Error("Did not get correct weight")
+	}
+	if s.nodes[2].weight != 101 {
+		t.Error("Did not get correct weight")
+	}
+}
+
+func TestStore_ApplyScoreChanges_UpdateWeightsNegativeDelta(t *testing.T) {
+	// Construct 3 nodes with weight 100 on each node. The 3 nodes linked to each other.
+	s := &Store{nodes: []*Node{
+		{root: [32]byte{'A'}, weight: 100},
+		{root: [32]byte{'A'}, weight: 100},
+		{parent: 1, root: [32]byte{'A'}, weight: 100}}}
+
+	// Each node gets one unique vote which contributes to negative delta.
+	// The weight should look like 97 <- 98 <- 99 because they get propagated back.
+	if err := s.applyWeightChanges(context.Background(), 0, 0, []int{-1, -1, -1}); err != nil {
+		t.Fatal(err)
+	}
+
+	if s.nodes[0].weight != 97 {
+		t.Error("Did not get correct weight")
+	}
+	if s.nodes[1].weight != 98 {
+		t.Error("Did not get correct weight")
+	}
+	if s.nodes[2].weight != 99 {
+		t.Error("Did not get correct weight")
+	}
+}
+
+func TestStore_ApplyScoreChanges_UpdateWeightsMixedDelta(t *testing.T) {
+	// Construct 3 nodes with weight 100 on each node. The 3 nodes linked to each other.
+	s := &Store{nodes: []*Node{
+		{root: [32]byte{'A'}, weight: 100},
+		{root: [32]byte{'A'}, weight: 100},
+		{parent: 1, root: [32]byte{'A'}, weight: 100}}}
+
+	// Each node gets one mixed vote. The weight should look like 100 <- 200 <- 250.
+	if err := s.applyWeightChanges(context.Background(), 0, 0, []int{-100, -50, 150}); err != nil {
+		t.Fatal(err)
+	}
+
+	if s.nodes[0].weight != 100 {
+		t.Error("Did not get correct weight")
+	}
+	if s.nodes[1].weight != 200 {
+		t.Error("Did not get correct weight")
+	}
+	if s.nodes[2].weight != 250 {
+		t.Error("Did not get correct weight")
+	}
+}
+
 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}