Take #1

beacon-chain/blockchain/forkchoice/proto-array/BUILD.bazel

@@ -2,7 +2,13 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library")
 
 go_library(
     name = "go_default_library",
-    srcs = ["type.go"],
+    srcs = [
+        "forkchoice.go",
+        "helpers.go",
+        "nodes.go",
+        "type.go",
+    ],
     importpath = "github.com/prysmaticlabs/prysm/beacon-chain/blockchain/forkchoice/proto-array",
     visibility = ["//visibility:public"],
+    deps = ["//shared/params:go_default_library"],
 )

beacon-chain/blockchain/forkchoice/proto-array/forkchoice.go

@@ -0,0 +1,29 @@
+package proto_array
+
+import "github.com/prysmaticlabs/prysm/shared/params"
+
+func (f *ForkChoice) New(justifiedEpoch uint64, finalizedEpoch uint64, finalizedRoot [32]byte) {
+	f.store = &Store{
+		justifiedEpoch: justifiedEpoch,
+		finalizedEpoch: finalizedEpoch,
+		finalizedRoot:  finalizedRoot,
+		nodes:          make([]Node, 0),
+		nodeIndices:    make(map[[32]byte]uint64),
+	}
+
+	f.store.Insert(finalizedRoot, params.BeaconConfig().ZeroHash, justifiedEpoch, finalizedEpoch)
+
+	f.balances = make([]uint64, 0)
+	f.votes = make([]Vote, 0)
+}
+
+func (f *ForkChoice) ProcessAttestation(validatorIndex uint64, blockRoot [32]byte, blockEpoch uint64) {
+	if blockEpoch > f.votes[validatorIndex].nextEpoch {
+		f.votes[validatorIndex].nextEpoch = blockEpoch
+		f.votes[validatorIndex].nextRoot = blockRoot
+	}
+}
+
+func (f *ForkChoice) ProcessBlock(blockRoot [32]byte, parentRoot [32]byte, finalizedEpoch uint64, justifiedEpoch uint64) {
+	f.store.Insert(blockRoot, parentRoot, justifiedEpoch, finalizedEpoch)
+}

beacon-chain/blockchain/forkchoice/proto-array/helpers.go

@@ -0,0 +1,49 @@
+package proto_array
+
+import (
+	"errors"
+
+	"github.com/prysmaticlabs/prysm/shared/params"
+)
+
+func computeDeltas(indices map[[32]byte]uint64, votes []Vote, oldBalances []uint64, newBalances []uint64) ([]int, error) {
+	deltas := make([]int, len(indices))
+
+	for validatorIndex, vote := range votes {
+		oldBalance := uint64(0)
+		newBalance := uint64(0)
+
+		// Skip if validator has never voted or voted for zero hash (ie genesis block)
+		if vote.currentRoot == params.BeaconConfig().ZeroHash || vote.nextRoot == params.BeaconConfig().ZeroHash {
+			continue
+		}
+
+		// If the validator did not exist in `old` or `new balance` list before, the balance is just 0.
+		if validatorIndex < len(oldBalances) {
+			oldBalance = oldBalances[validatorIndex]
+		}
+		if validatorIndex < len(newBalances) {
+			newBalance = newBalances[validatorIndex]
+		}
+
+		// Perform delta only if vote has changed and balance has changed.
+		if vote.currentRoot != vote.nextRoot || oldBalance != newBalance {
+			// Ignore the vote if it's not known in `indices`
+			nextDeltaIndex, ok := indices[vote.nextRoot]
+			if !ok {
+				return nil, errors.New("vote is not a key in indices")
+			}
+			deltas[nextDeltaIndex] += int(newBalance)
+
+			currentDeltaIndex, ok := indices[vote.currentRoot]
+			if !ok {
+				return nil, errors.New("vote is not a key in indices")
+			}
+			deltas[currentDeltaIndex] -= int(oldBalance)
+		}
+
+		vote.currentRoot = vote.nextRoot
+	}
+
+	return deltas
+}

beacon-chain/blockchain/forkchoice/proto-array/nodes.go

@@ -0,0 +1,117 @@
+package proto_array
+
+import "bytes"
+
+// Insert registers a new block with the fork choice.
+func (s Store) Insert(root [32]byte, parent [32]byte, justifiedEpoch uint64, finalizedEpoch uint64) {
+	s.nodeIndicesLock.Lock()
+	defer s.nodeIndicesLock.Unlock()
+
+	index := len(s.nodes)
+	parentIndex, ok := s.nodeIndices[parent]
+	// Mark genesis block's parent as non existent.
+	if !ok {
+		parentIndex = nonExistentNode
+	}
+
+	n := Node{
+		root: root,
+		parent: parentIndex,
+		justifiedEpoch: justifiedEpoch,
+		finalizedEpoch: finalizedEpoch,
+		bestChild: nonExistentNode,
+		bestDescendant: nonExistentNode,
+		weight: 0,
+	}
+
+	s.nodeIndices[root] = uint64(index)
+	s.nodes = append(s.nodes, n)
+
+	if n.parent != nonExistentNode {
+		s.UpdateBestChildAndDescendant(parentIndex, uint64(index))
+	}
+}
+
+// UpdateBestChildAndDescendant updates parent node's best child and descendent.
+// It looks at parent node and child node and potentially modifies parent's best
+// child and best descendent values.
+// There are four outcomes:
+// - The child is already the best child but it's now invalid due to a FFG change and should be removed.
+// - The child is already the best child and the parent is updated with the new best descendant.
+// - The child is not the best child but becomes the best child.
+// - The child is not the best child and does not become best child.
+func (s Store) UpdateBestChildAndDescendant(parentIndex uint64, childIndex uint64) {
+	parent := s.nodes[parentIndex]
+	child := s.nodes[childIndex]
+
+	childLeadsToViableHead := s.LeadsToViableHead(child)
+
+	// Define 3 variables for the 3 options mentioned above. This is to
+	// set `parent.bestChild` and `parent.bestDescendent` to. These
+	// aliases are to assist readability.
+	changeToNone := []uint64{nonExistentNode, nonExistentNode}
+	changeToChild := []uint64{childIndex, child.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.
+			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.
+			newParentChild = changeToChild
+		} else {
+			bestChild := s.nodes[parent.bestChild]
+			bestChildLeadsToViableHead := s.LeadsToViableHead(bestChild)
+
+			if childLeadsToViableHead && !bestChildLeadsToViableHead {
+				// The child leads to a viable head, but the current best child doesnt.
+				newParentChild = changeToChild
+			} else if !childLeadsToViableHead && bestChildLeadsToViableHead {
+				// The best child leads to viable head, but the child doesnt.
+				newParentChild = noChange
+			} else if child.weight == bestChild.weight {
+				// 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 {
+			// There's no current best child and the child is viable.
+			newParentChild = changeToChild
+		} else {
+			// There's no current best child and the child is not viable.
+			newParentChild = noChange
+		}
+	}
+
+	parent.bestChild = newParentChild[0]
+	parent.bestDescendant = newParentChild[1]
+	s.nodes[parentIndex] = parent
+}
+
+func (s Store) LeadsToViableHead(node Node) bool {
+	bestDescendentIndex := node.bestDescendant
+	bestDescendentNode := s.nodes[bestDescendentIndex]
+	return s.ViableForHead(bestDescendentNode) || s.ViableForHead(node)
+}
+
+func (s Store) ViableForHead(node Node) bool {
+	justified := s.justifiedEpoch == node.justifiedEpoch || s.justifiedEpoch == 0
+	finalized := s.finalizedEpoch == node.finalizedEpoch || s.finalizedEpoch == 0
+	return justified && finalized
+}

beacon-chain/blockchain/forkchoice/proto-array/type.go

@@ -1,9 +1,11 @@
 package proto_array
 
+import "sync"
+
 // ForkChoice defines the overall fork choice store which includes block nodes, validator's latest votes and balances.
 type ForkChoice struct {
 	store    *Store
-	votes    []*Vote  // tracks individual validator's latest vote.
+	votes    []Vote  // tracks individual validator's latest vote.
 	balances []uint64 // tracks individual validator's effective balances.
 }
 
@@ -15,7 +17,8 @@ type Store struct {
 	finalizedEpoch     uint64              // latest finalized epoch in store.
 	finalizedRoot      [32]byte            // latest finalized root in store.
 	nodes              []Node              // list of block nodes, each node is a representation of one block.
-	indices            map[[32]byte]uint64 // root of block node and its index in the nodes list.
+	nodeIndices        map[[32]byte]uint64 // the root of block node and the nodes index in the list.
+	nodeIndicesLock    sync.RWMutex
 }
 
 // Node defines the individual block which includes its parent, ancestor and how much weight accounted for it.
@@ -35,3 +38,5 @@ type Vote struct {
 	nextRoot    [32]byte
 	nextEpoch   uint64
 }
+
+const nonExistentNode = ^uint64(0)