Add Balance Field Trie

beacon-chain/state/fieldtrie/BUILD.bazel

@@ -12,6 +12,7 @@ go_library(
         "//beacon-chain/state/stateutil:go_default_library",
         "//beacon-chain/state/types:go_default_library",
         "//crypto/hash:go_default_library",
+        "//encoding/ssz:go_default_library",
         "//proto/prysm/v1alpha1:go_default_library",
         "//runtime/version:go_default_library",
         "@com_github_pkg_errors//:go_default_library",

beacon-chain/state/fieldtrie/field_trie.go

@@ -18,6 +18,7 @@ type FieldTrie struct {
 	field       types.FieldIndex
 	dataType    types.DataType
 	length      uint64
+	numOfElems  int
 }
 
 // NewFieldTrie is the constructor for the field trie data structure. It creates the corresponding
@@ -26,11 +27,12 @@ type FieldTrie struct {
 func NewFieldTrie(field types.FieldIndex, dataType types.DataType, elements interface{}, length uint64) (*FieldTrie, error) {
 	if elements == nil {
 		return &FieldTrie{
-			field:     field,
-			dataType:  dataType,
-			reference: stateutil.NewRef(1),
-			RWMutex:   new(sync.RWMutex),
-			length:    length,
+			field:      field,
+			dataType:   dataType,
+			reference:  stateutil.NewRef(1),
+			RWMutex:    new(sync.RWMutex),
+			length:     length,
+			numOfElems: 0,
 		}, nil
 	}
 	fieldRoots, err := fieldConverters(field, []uint64{}, elements, true)
@@ -49,15 +51,17 @@ func NewFieldTrie(field types.FieldIndex, dataType types.DataType, elements inte
 			reference:   stateutil.NewRef(1),
 			RWMutex:     new(sync.RWMutex),
 			length:      length,
+			numOfElems:  reflect.ValueOf(elements).Len(),
 		}, nil
-	case types.CompositeArray:
+	case types.CompositeArray, types.CompressedArray:
 		return &FieldTrie{
 			fieldLayers: stateutil.ReturnTrieLayerVariable(fieldRoots, length),
 			field:       field,
 			dataType:    dataType,
 			reference:   stateutil.NewRef(1),
 			RWMutex:     new(sync.RWMutex),
 			length:      length,
+			numOfElems:  reflect.ValueOf(elements).Len(),
 		}, nil
 	default:
 		return nil, errors.Errorf("unrecognized data type in field map: %v", reflect.TypeOf(dataType).Name())
@@ -88,13 +92,40 @@ func (f *FieldTrie) RecomputeTrie(indices []uint64, elements interface{}) ([32]b
 		if err != nil {
 			return [32]byte{}, err
 		}
+		f.numOfElems = reflect.ValueOf(elements).Len()
 		return fieldRoot, nil
 	case types.CompositeArray:
 		fieldRoot, f.fieldLayers, err = stateutil.RecomputeFromLayerVariable(fieldRoots, indices, f.fieldLayers)
 		if err != nil {
 			return [32]byte{}, err
 		}
+		f.numOfElems = reflect.ValueOf(elements).Len()
 		return stateutil.AddInMixin(fieldRoot, uint64(len(f.fieldLayers[0])))
+	case types.CompressedArray:
+		numOfElems, err := f.field.CompressedLength()
+		if err != nil {
+			return [32]byte{}, err
+		}
+		// We remove the duplicates here in order to prevent
+		// duplicated insertions into the trie.
+		newIndices := []uint64{}
+		indexExists := make(map[uint64]bool)
+		newRoots := make([][32]byte, 0, len(fieldRoots)/int(numOfElems))
+		for i, idx := range indices {
+			startIdx := idx / numOfElems
+			if indexExists[startIdx] {
+				continue
+			}
+			newIndices = append(newIndices, startIdx)
+			indexExists[startIdx] = true
+			newRoots = append(newRoots, fieldRoots[i])
+		}
+		fieldRoot, f.fieldLayers, err = stateutil.RecomputeFromLayerVariable(newRoots, newIndices, f.fieldLayers)
+		if err != nil {
+			return [32]byte{}, err
+		}
+		f.numOfElems = reflect.ValueOf(elements).Len()
+		return stateutil.AddInMixin(fieldRoot, uint64(f.numOfElems))
 	default:
 		return [32]byte{}, errors.Errorf("unrecognized data type in field map: %v", reflect.TypeOf(f.dataType).Name())
 	}
@@ -105,11 +136,12 @@ func (f *FieldTrie) RecomputeTrie(indices []uint64, elements interface{}) ([32]b
 func (f *FieldTrie) CopyTrie() *FieldTrie {
 	if f.fieldLayers == nil {
 		return &FieldTrie{
-			field:     f.field,
-			dataType:  f.dataType,
-			reference: stateutil.NewRef(1),
-			RWMutex:   new(sync.RWMutex),
-			length:    f.length,
+			field:      f.field,
+			dataType:   f.dataType,
+			reference:  stateutil.NewRef(1),
+			RWMutex:    new(sync.RWMutex),
+			length:     f.length,
+			numOfElems: f.numOfElems,
 		}
 	}
 	dstFieldTrie := make([][]*[32]byte, len(f.fieldLayers))
@@ -124,6 +156,7 @@ func (f *FieldTrie) CopyTrie() *FieldTrie {
 		reference:   stateutil.NewRef(1),
 		RWMutex:     new(sync.RWMutex),
 		length:      f.length,
+		numOfElems:  f.numOfElems,
 	}
 }
 
@@ -135,6 +168,9 @@ func (f *FieldTrie) TrieRoot() ([32]byte, error) {
 	case types.CompositeArray:
 		trieRoot := *f.fieldLayers[len(f.fieldLayers)-1][0]
 		return stateutil.AddInMixin(trieRoot, uint64(len(f.fieldLayers[0])))
+	case types.CompressedArray:
+		trieRoot := *f.fieldLayers[len(f.fieldLayers)-1][0]
+		return stateutil.AddInMixin(trieRoot, uint64(f.numOfElems))
 	default:
 		return [32]byte{}, errors.Errorf("unrecognized data type in field map: %v", reflect.TypeOf(f.dataType).Name())
 	}

beacon-chain/state/fieldtrie/field_trie_helpers.go

@@ -1,13 +1,15 @@
 package fieldtrie
 
 import (
+	"encoding/binary"
 	"fmt"
 	"reflect"
 
 	"github.com/pkg/errors"
 	"github.com/prysmaticlabs/prysm/beacon-chain/state/stateutil"
 	"github.com/prysmaticlabs/prysm/beacon-chain/state/types"
 	"github.com/prysmaticlabs/prysm/crypto/hash"
+	"github.com/prysmaticlabs/prysm/encoding/ssz"
 	ethpb "github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1"
 	"github.com/prysmaticlabs/prysm/runtime/version"
 )
@@ -60,6 +62,13 @@ func fieldConverters(field types.FieldIndex, indices []uint64, elements interfac
 				reflect.TypeOf([]*ethpb.PendingAttestation{}).Name(), reflect.TypeOf(elements).Name())
 		}
 		return handlePendingAttestation(val, indices, convertAll)
+	case types.Balances:
+		val, ok := elements.([]uint64)
+		if !ok {
+			return nil, errors.Errorf("Wanted type of %v but got %v",
+				reflect.TypeOf([]uint64{}).Name(), reflect.TypeOf(elements).Name())
+		}
+		return handleBalanceSlice(val, indices, convertAll)
 	default:
 		return [][32]byte{}, errors.Errorf("got unsupported type of %v", reflect.TypeOf(elements).Name())
 	}
@@ -141,3 +150,47 @@ func handlePendingAttestation(val []*ethpb.PendingAttestation, indices []uint64,
 	}
 	return roots, nil
 }
+
+func handleBalanceSlice(val []uint64, indices []uint64, convertAll bool) ([][32]byte, error) {
+	if convertAll {
+		balancesMarshaling := make([][]byte, 0)
+		for i := 0; i < len(val); i++ {
+			balanceBuf := make([]byte, 8)
+			binary.LittleEndian.PutUint64(balanceBuf, val[i])
+			balancesMarshaling = append(balancesMarshaling, balanceBuf)
+		}
+		balancesChunks, err := ssz.PackByChunk(balancesMarshaling)
+		if err != nil {
+			return [][32]byte{}, errors.Wrap(err, "could not pack balances into chunks")
+		}
+		return balancesChunks, nil
+	}
+	if len(val) > 0 {
+		numOfElems, err := types.Balances.CompressedLength()
+		if err != nil {
+			return nil, err
+		}
+		roots := [][32]byte{}
+		for _, idx := range indices {
+			// We split the indexes into their relevant groups. Balances
+			// are compressed according to 4 values -> 1 chunk.
+			startIdx := idx / numOfElems
+			startGroup := startIdx * numOfElems
+			chunk := [32]byte{}
+			sizeOfElem := len(chunk) / int(numOfElems)
+			for i, j := 0, startGroup; j < startGroup+numOfElems; i, j = i+sizeOfElem, j+1 {
+				wantedVal := uint64(0)
+				// We only select from a field value, if the
+				// index exists in our element list. If it doesn't
+				// we assume a zero value.
+				if int(j) < len(val) {
+					wantedVal = val[j]
+				}
+				binary.LittleEndian.PutUint64(chunk[i:i+sizeOfElem], wantedVal)
+			}
+			roots = append(roots, chunk)
+		}
+		return roots, nil
+	}
+	return [][32]byte{}, nil
+}

beacon-chain/state/types/BUILD.bazel

@@ -5,7 +5,10 @@ go_library(
     srcs = ["types.go"],
     importpath = "github.com/prysmaticlabs/prysm/beacon-chain/state/types",
     visibility = ["//beacon-chain:__subpackages__"],
-    deps = ["//runtime/version:go_default_library"],
+    deps = [
+        "//runtime/version:go_default_library",
+        "@com_github_pkg_errors//:go_default_library",
+    ],
 )
 
 go_test(

beacon-chain/state/types/types.go

@@ -1,6 +1,7 @@
 package types
 
 import (
+	"github.com/pkg/errors"
 	"github.com/prysmaticlabs/prysm/runtime/version"
 )
 
@@ -18,6 +19,10 @@ const (
 	// CompositeArray represents a variable length array with
 	// a non primitive type.
 	CompositeArray
+	// CompressedArray represents a variable length array which
+	// can pack multiple elements into a leaf of the underlying
+	// trie.
+	CompressedArray
 )
 
 // String returns the name of the field index.
@@ -82,6 +87,17 @@ func (f FieldIndex) String(stateVersion int) string {
 	}
 }
 
+// CompressedLength returns the compressed length of the type(number of
+// elements that are able to be packed).
+func (f FieldIndex) CompressedLength() (uint64, error) {
+	switch f {
+	case Balances:
+		return 4, nil
+	default:
+		return 0, errors.Errorf("field %d doesn't support element compression", f)
+	}
+}
+
 // Below we define a set of useful enum values for the field
 // indices of the beacon state. For example, genesisTime is the
 // 0th field of the beacon state. This is helpful when we are

beacon-chain/state/v1/BUILD.bazel

@@ -87,6 +87,7 @@ go_test(
         "//beacon-chain/state:go_default_library",
         "//beacon-chain/state/stateutil:go_default_library",
         "//beacon-chain/state/types:go_default_library",
+        "//config/features:go_default_library",
         "//config/params:go_default_library",
         "//encoding/bytesutil:go_default_library",
         "//proto/prysm/v1alpha1:go_default_library",

beacon-chain/state/v1/setters_misc.go

@@ -5,6 +5,7 @@ import (
 	types "github.com/prysmaticlabs/eth2-types"
 	"github.com/prysmaticlabs/prysm/beacon-chain/state/stateutil"
 	stateTypes "github.com/prysmaticlabs/prysm/beacon-chain/state/types"
+	"github.com/prysmaticlabs/prysm/config/features"
 	"github.com/prysmaticlabs/prysm/crypto/hash"
 	ethpb "github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1"
 	"google.golang.org/protobuf/proto"
@@ -172,6 +173,10 @@ func (b *BeaconState) addDirtyIndices(index stateTypes.FieldIndex, indices []uin
 	if b.rebuildTrie[index] {
 		return
 	}
+	// Exit early if balance trie computation isn't enabled.
+	if !features.Get().EnableBalanceTrieComputation && index == balances {
+		return
+	}
 	totalIndicesLen := len(b.dirtyIndices[index]) + len(indices)
 	if totalIndicesLen > indicesLimit {
 		b.rebuildTrie[index] = true

beacon-chain/state/v1/setters_validator.go

@@ -103,6 +103,7 @@ func (b *BeaconState) SetBalances(val []uint64) error {
 
 	b.state.Balances = val
 	b.markFieldAsDirty(balances)
+	b.rebuildTrie[balances] = true
 	return nil
 }
 
@@ -128,6 +129,7 @@ func (b *BeaconState) UpdateBalancesAtIndex(idx types.ValidatorIndex, val uint64
 	bals[idx] = val
 	b.state.Balances = bals
 	b.markFieldAsDirty(balances)
+	b.addDirtyIndices(balances, []uint64{uint64(idx)})
 	return nil
 }
 
@@ -219,6 +221,8 @@ func (b *BeaconState) AppendBalance(bal uint64) error {
 	}
 
 	b.state.Balances = append(bals, bal)
+	balIdx := len(b.state.Balances) - 1
 	b.markFieldAsDirty(balances)
+	b.addDirtyIndices(balances, []uint64{uint64(balIdx)})
 	return nil
 }