Add Balance Field Trie

beacon-chain/state/fieldtrie/BUILD.bazel

@@ -12,6 +12,8 @@ go_library(
         "//beacon-chain/state/stateutil:go_default_library",
         "//beacon-chain/state/types:go_default_library",
         "//crypto/hash:go_default_library",
+        "//encoding/bytesutil: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",
@@ -26,6 +28,7 @@ go_test(
     ],
     embed = [":go_default_library"],
     deps = [
+        "//beacon-chain/state/stateutil:go_default_library",
         "//beacon-chain/state/types:go_default_library",
         "//beacon-chain/state/v1:go_default_library",
         "//config/params: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,18 +27,19 @@ 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)
 	if err != nil {
 		return nil, err
 	}
-	if err := validateElements(field, elements, length); err != nil {
+	if err := validateElements(field, dataType, elements, length); err != nil {
 		return nil, err
 	}
 	switch dataType {
@@ -53,15 +55,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())
@@ -92,13 +96,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.ElemsInChunk()
+		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())
 	}
@@ -109,11 +140,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))
@@ -128,6 +160,7 @@ func (f *FieldTrie) CopyTrie() *FieldTrie {
 		reference:   stateutil.NewRef(1),
 		RWMutex:     new(sync.RWMutex),
 		length:      f.length,
+		numOfElems:  f.numOfElems,
 	}
 }
 
@@ -139,6 +172,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,27 +1,45 @@
 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/bytesutil"
+	"github.com/prysmaticlabs/prysm/encoding/ssz"
 	ethpb "github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1"
 	"github.com/prysmaticlabs/prysm/runtime/version"
 )
 
 func (f *FieldTrie) validateIndices(idxs []uint64) error {
+	length := f.length
+	if f.dataType == types.CompressedArray {
+		comLength, err := f.field.ElemsInChunk()
+		if err != nil {
+			return err
+		}
+		length *= comLength
+	}
 	for _, idx := range idxs {
-		if idx >= f.length {
-			return errors.Errorf("invalid index for field %s: %d >= length %d", f.field.String(version.Phase0), idx, f.length)
+		if idx >= length {
+			return errors.Errorf("invalid index for field %s: %d >= length %d", f.field.String(version.Phase0), idx, length)
 		}
 	}
 	return nil
 }
 
-func validateElements(field types.FieldIndex, elements interface{}, length uint64) error {
+func validateElements(field types.FieldIndex, dataType types.DataType, elements interface{}, length uint64) error {
+	if dataType == types.CompressedArray {
+		comLength, err := field.ElemsInChunk()
+		if err != nil {
+			return err
+		}
+		length *= comLength
+	}
 	val := reflect.ValueOf(elements)
 	if val.Len() > int(length) {
 		return errors.Errorf("elements length is larger than expected for field %s: %d > %d", field.String(version.Phase0), val.Len(), length)
@@ -38,35 +56,109 @@ func fieldConverters(field types.FieldIndex, indices []uint64, elements interfac
 			return nil, errors.Errorf("Wanted type of %v but got %v",
 				reflect.TypeOf([][]byte{}).Name(), reflect.TypeOf(elements).Name())
 		}
-		return stateutil.HandleByteArrays(val, indices, convertAll)
+		return handleByteArrays(val, indices, convertAll)
 	case types.Eth1DataVotes:
 		val, ok := elements.([]*ethpb.Eth1Data)
 		if !ok {
 			return nil, errors.Errorf("Wanted type of %v but got %v",
 				reflect.TypeOf([]*ethpb.Eth1Data{}).Name(), reflect.TypeOf(elements).Name())
 		}
-		return HandleEth1DataSlice(val, indices, convertAll)
+		return handleEth1DataSlice(val, indices, convertAll)
 	case types.Validators:
 		val, ok := elements.([]*ethpb.Validator)
 		if !ok {
 			return nil, errors.Errorf("Wanted type of %v but got %v",
 				reflect.TypeOf([]*ethpb.Validator{}).Name(), reflect.TypeOf(elements).Name())
 		}
-		return stateutil.HandleValidatorSlice(val, indices, convertAll)
+		return handleValidatorSlice(val, indices, convertAll)
 	case types.PreviousEpochAttestations, types.CurrentEpochAttestations:
 		val, ok := elements.([]*ethpb.PendingAttestation)
 		if !ok {
 			return nil, errors.Errorf("Wanted type of %v but got %v",
 				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())
 	}
 }
 
-// HandleEth1DataSlice processes a list of eth1data and indices into the appropriate roots.
-func HandleEth1DataSlice(val []*ethpb.Eth1Data, indices []uint64, convertAll bool) ([][32]byte, error) {
+// handleByteArrays computes and returns byte arrays in a slice of root format.
+func handleByteArrays(val [][]byte, indices []uint64, convertAll bool) ([][32]byte, error) {
+	length := len(indices)
+	if convertAll {
+		length = len(val)
+	}
+	roots := make([][32]byte, 0, length)
+	rootCreator := func(input []byte) {
+		newRoot := bytesutil.ToBytes32(input)
+		roots = append(roots, newRoot)
+	}
+	if convertAll {
+		for i := range val {
+			rootCreator(val[i])
+		}
+		return roots, nil
+	}
+	if len(val) > 0 {
+		for _, idx := range indices {
+			if idx > uint64(len(val))-1 {
+				return nil, fmt.Errorf("index %d greater than number of byte arrays %d", idx, len(val))
+			}
+			rootCreator(val[idx])
+		}
+	}
+	return roots, nil
+}
+
+// handleValidatorSlice returns the validator indices in a slice of root format.
+func handleValidatorSlice(val []*ethpb.Validator, indices []uint64, convertAll bool) ([][32]byte, error) {
+	length := len(indices)
+	if convertAll {
+		length = len(val)
+	}
+	roots := make([][32]byte, 0, length)
+	hasher := hash.CustomSHA256Hasher()
+	rootCreator := func(input *ethpb.Validator) error {
+		newRoot, err := stateutil.ValidatorRootWithHasher(hasher, input)
+		if err != nil {
+			return err
+		}
+		roots = append(roots, newRoot)
+		return nil
+	}
+	if convertAll {
+		for i := range val {
+			err := rootCreator(val[i])
+			if err != nil {
+				return nil, err
+			}
+		}
+		return roots, nil
+	}
+	if len(val) > 0 {
+		for _, idx := range indices {
+			if idx > uint64(len(val))-1 {
+				return nil, fmt.Errorf("index %d greater than number of validators %d", idx, len(val))
+			}
+			err := rootCreator(val[idx])
+			if err != nil {
+				return nil, err
+			}
+		}
+	}
+	return roots, nil
+}
+
+// handleEth1DataSlice processes a list of eth1data and indices into the appropriate roots.
+func handleEth1DataSlice(val []*ethpb.Eth1Data, indices []uint64, convertAll bool) ([][32]byte, error) {
 	length := len(indices)
 	if convertAll {
 		length = len(val)
@@ -141,3 +233,48 @@ 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 _, b := range val {
+			balanceBuf := make([]byte, 8)
+			binary.LittleEndian.PutUint64(balanceBuf, b)
+			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.ElemsInChunk()
+		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 are adding chunks in sets of 4, if the set is at the edge of the array
+				// then you will need to zero out the rest of the chunk. Ex : 41 indexes,
+				// so 41 % 4 = 1 . There are 3 indexes, which do not exist yet but we
+				// have to add in as a root. These 3 indexes are then given 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
+}