Improve Field Trie Recomputation (#10884)

Co-authored-by: Raul Jordan <raul@prysmaticlabs.com>
Co-authored-by: Radosław Kapka <rkapka@wp.pl>

.github/CODEOWNERS

@@ -9,8 +9,8 @@ deps.bzl @prysmaticlabs/core-team
 
 # Radek and Nishant are responsible for changes that can affect the native state feature.
 # See https://www.notion.so/prysmaticlabs/Native-Beacon-State-Redesign-6cc9744b4ec1439bb34fa829b36a35c1
-/beacon-chain/state/fieldtrie/ @rkapka @nisdas
-/beacon-chain/state/v1/ @rkapka @nisdas
-/beacon-chain/state/v2/ @rkapka @nisdas
-/beacon-chain/state/v3/ @rkapka @nisdas
-/beacon-chain/state/state-native/ @rkapka @nisdas
+/beacon-chain/state/fieldtrie/ @rkapka @nisdas @rauljordan
+/beacon-chain/state/v1/ @rkapka @nisdas @rauljordan
+/beacon-chain/state/v2/ @rkapka @nisdas @rauljordan
+/beacon-chain/state/v3/ @rkapka @nisdas @rauljordan
+/beacon-chain/state/state-native/ @rkapka @nisdas @rauljordan

api/client/builder/types.go

@@ -63,7 +63,7 @@ func sszBytesToUint256(b []byte) Uint256 {
 
 // SSZBytes creates an ssz-style (little-endian byte slice) representation of the Uint256
 func (s Uint256) SSZBytes() []byte {
-	return bytesutil.ReverseByteOrder(s.Int.Bytes())
+	return bytesutil.PadTo(bytesutil.ReverseByteOrder(s.Int.Bytes()), 32)
 }
 
 var errUnmarshalUint256Failed = errors.New("unable to UnmarshalText into a Uint256 value")

api/client/builder/types_test.go

@@ -694,9 +694,10 @@ func TestMarshalBlindedBeaconBlockBodyBellatrix(t *testing.T) {
 }
 
 func TestRoundTripUint256(t *testing.T) {
-	vs := "452312848583266388373324160190187140051835877600158453279131187530910662656"
+	vs := "4523128485832663883733241601901871400518358776001584532791311875309106626"
 	u := stringToUint256(vs)
 	sb := u.SSZBytes()
+	require.Equal(t, 32, len(sb))
 	uu := sszBytesToUint256(sb)
 	require.Equal(t, true, bytes.Equal(u.SSZBytes(), uu.SSZBytes()))
 	require.Equal(t, vs, uu.String())

beacon-chain/db/kv/kv.go

@@ -105,7 +105,6 @@ func NewKVStore(ctx context.Context, dirPath string, config *Config) (*Store, er
 		}
 	}
 	datafile := KVStoreDatafilePath(dirPath)
-	start := time.Now()
 	log.Infof("Opening Bolt DB at %s", datafile)
 	boltDB, err := bolt.Open(
 		datafile,
@@ -116,40 +115,29 @@ func NewKVStore(ctx context.Context, dirPath string, config *Config) (*Store, er
 		},
 	)
 	if err != nil {
-		log.WithField("elapsed", time.Since(start)).Error("Failed to open Bolt DB")
 		if errors.Is(err, bolt.ErrTimeout) {
 			return nil, errors.New("cannot obtain database lock, database may be in use by another process")
 		}
 		return nil, err
 	}
-	log.WithField("elapsed", time.Since(start)).Info("Opened Bolt DB")
-
 	boltDB.AllocSize = boltAllocSize
-	start = time.Now()
-	log.Infof("Creating block cache...")
 	blockCache, err := ristretto.NewCache(&ristretto.Config{
 		NumCounters: 1000,           // number of keys to track frequency of (1000).
 		MaxCost:     BlockCacheSize, // maximum cost of cache (1000 Blocks).
 		BufferItems: 64,             // number of keys per Get buffer.
 	})
 	if err != nil {
-		log.WithField("elapsed", time.Since(start)).Error("Failed to create block cache")
 		return nil, err
 	}
-	log.WithField("elapsed", time.Since(start)).Info("Created block cache")
 
-	start = time.Now()
-	log.Infof("Creating validator cache...")
 	validatorCache, err := ristretto.NewCache(&ristretto.Config{
 		NumCounters: NumOfValidatorEntries, // number of entries in cache (2 Million).
 		MaxCost:     ValidatorEntryMaxCost, // maximum size of the cache (64Mb)
 		BufferItems: 64,                    // number of keys per Get buffer.
 	})
 	if err != nil {
-		log.WithField("elapsed", time.Since(start)).Error("Failed to to create validator cache")
 		return nil, err
 	}
-	log.WithField("elapsed", time.Since(start)).Info("Created validator cache")
 
 	kv := &Store{
 		db:                  boltDB,
@@ -159,8 +147,6 @@ func NewKVStore(ctx context.Context, dirPath string, config *Config) (*Store, er
 		stateSummaryCache:   newStateSummaryCache(),
 		ctx:                 ctx,
 	}
-	start = time.Now()
-	log.Infof("Updating DB and creating buckets...")
 	if err := kv.db.Update(func(tx *bolt.Tx) error {
 		return createBuckets(
 			tx,
@@ -195,13 +181,9 @@ func NewKVStore(ctx context.Context, dirPath string, config *Config) (*Store, er
 			registrationBucket,
 		)
 	}); err != nil {
-		log.WithField("elapsed", time.Since(start)).Error("Failed to update db and create buckets")
 		return nil, err
 	}
-	log.WithField("elapsed", time.Since(start)).Info("Updated db and created buckets")
-
 	err = prometheus.Register(createBoltCollector(kv.db))
-
 	return kv, err
 }
 

beacon-chain/state/fieldtrie/field_trie.go

@@ -19,12 +19,13 @@ var (
 // trie of the particular field.
 type FieldTrie struct {
 	*sync.RWMutex
-	reference   *stateutil.Reference
-	fieldLayers [][]*[32]byte
-	field       types.BeaconStateField
-	dataType    types.DataType
-	length      uint64
-	numOfElems  int
+	reference     *stateutil.Reference
+	fieldLayers   [][]*[32]byte
+	field         types.BeaconStateField
+	dataType      types.DataType
+	length        uint64
+	numOfElems    int
+	isTransferred bool
 }
 
 // NewFieldTrie is the constructor for the field trie data structure. It creates the corresponding
@@ -191,6 +192,43 @@ func (f *FieldTrie) CopyTrie() *FieldTrie {
 	}
 }
 
+// Length return the length of the whole field trie.
+func (f *FieldTrie) Length() uint64 {
+	return f.length
+}
+
+// TransferTrie starts the process of transferring all the
+// trie related data to a new trie. This is done if we
+// know that other states which hold references to this
+// trie will unlikely need it for recomputation. This helps
+// us save on a copy. Any caller of this method will need
+// to take care that this isn't called on an empty trie.
+func (f *FieldTrie) TransferTrie() *FieldTrie {
+	if f.fieldLayers == nil {
+		return &FieldTrie{
+			field:      f.field,
+			dataType:   f.dataType,
+			reference:  stateutil.NewRef(1),
+			RWMutex:    new(sync.RWMutex),
+			length:     f.length,
+			numOfElems: f.numOfElems,
+		}
+	}
+	f.isTransferred = true
+	nTrie := &FieldTrie{
+		fieldLayers: f.fieldLayers,
+		field:       f.field,
+		dataType:    f.dataType,
+		reference:   stateutil.NewRef(1),
+		RWMutex:     new(sync.RWMutex),
+		length:      f.length,
+		numOfElems:  f.numOfElems,
+	}
+	// Zero out field layers here.
+	f.fieldLayers = nil
+	return nTrie
+}
+
 // TrieRoot returns the corresponding root of the trie.
 func (f *FieldTrie) TrieRoot() ([32]byte, error) {
 	if f.Empty() {
@@ -222,7 +260,7 @@ func (f *FieldTrie) FieldReference() *stateutil.Reference {
 // Empty checks whether the underlying field trie is
 // empty or not.
 func (f *FieldTrie) Empty() bool {
-	return f == nil || len(f.fieldLayers) == 0
+	return f == nil || len(f.fieldLayers) == 0 || f.isTransferred
 }
 
 // InsertFieldLayer manually inserts a field layer. This method

beacon-chain/state/state-native/state_trie.go

@@ -740,17 +740,31 @@ func (b *BeaconState) rootSelector(ctx context.Context, field nativetypes.FieldI
 
 func (b *BeaconState) recomputeFieldTrie(index nativetypes.FieldIndex, elements interface{}) ([32]byte, error) {
 	fTrie := b.stateFieldLeaves[index]
+	fTrieMutex := fTrie.RWMutex
 	// We can't lock the trie directly because the trie's variable gets reassigned,
 	// and therefore we would call Unlock() on a different object.
-	fTrieMutex := fTrie.RWMutex
-	if fTrie.FieldReference().Refs() > 1 {
-		fTrieMutex.Lock()
+	fTrieMutex.Lock()
+
+	if fTrie.Empty() {
+		err := b.resetFieldTrie(index, elements, fTrie.Length())
+		if err != nil {
+			fTrieMutex.Unlock()
+			return [32]byte{}, err
+		}
+		// Reduce reference count as we are instantiating a new trie.
 		fTrie.FieldReference().MinusRef()
-		newTrie := fTrie.CopyTrie()
+		fTrieMutex.Unlock()
+		return b.stateFieldLeaves[index].TrieRoot()
+	}
+
+	if fTrie.FieldReference().Refs() > 1 {
+		fTrie.FieldReference().MinusRef()
+		newTrie := fTrie.TransferTrie()
 		b.stateFieldLeaves[index] = newTrie
 		fTrie = newTrie
-		fTrieMutex.Unlock()
 	}
+	fTrieMutex.Unlock()
+
 	// remove duplicate indexes
 	b.dirtyIndices[index] = slice.SetUint64(b.dirtyIndices[index])
 	// sort indexes again

beacon-chain/state/state-native/state_trie_test.go

@@ -383,3 +383,204 @@ func TestBeaconState_AppendValidator_DoesntMutateCopy(t *testing.T) {
 	_, ok := st1.ValidatorIndexByPubkey(bytesutil.ToBytes48(val.PublicKey))
 	assert.Equal(t, false, ok, "Expected no validator index to be present in st1 for the newly inserted pubkey")
 }
+
+func TestBeaconState_ValidatorMutation_Phase0(t *testing.T) {
+	testState, _ := util.DeterministicGenesisState(t, 400)
+	pbState, err := statenative.ProtobufBeaconStatePhase0(testState.InnerStateUnsafe())
+	require.NoError(t, err)
+	testState, err = statenative.InitializeFromProtoPhase0(pbState)
+	require.NoError(t, err)
+
+	_, err = testState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	// Reset tries
+	require.NoError(t, testState.UpdateValidatorAtIndex(200, new(ethpb.Validator)))
+	_, err = testState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	newState1 := testState.Copy()
+	_ = testState.Copy()
+
+	require.NoError(t, testState.UpdateValidatorAtIndex(15, &ethpb.Validator{
+		PublicKey:                  make([]byte, 48),
+		WithdrawalCredentials:      make([]byte, 32),
+		EffectiveBalance:           1111,
+		Slashed:                    false,
+		ActivationEligibilityEpoch: 1112,
+		ActivationEpoch:            1114,
+		ExitEpoch:                  1116,
+		WithdrawableEpoch:          1117,
+	}))
+
+	rt, err := testState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+	pbState, err = statenative.ProtobufBeaconStatePhase0(testState.InnerStateUnsafe())
+	require.NoError(t, err)
+
+	copiedTestState, err := statenative.InitializeFromProtoPhase0(pbState)
+	require.NoError(t, err)
+
+	rt2, err := copiedTestState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	assert.Equal(t, rt, rt2)
+
+	require.NoError(t, newState1.UpdateValidatorAtIndex(150, &ethpb.Validator{
+		PublicKey:                  make([]byte, 48),
+		WithdrawalCredentials:      make([]byte, 32),
+		EffectiveBalance:           2111,
+		Slashed:                    false,
+		ActivationEligibilityEpoch: 2112,
+		ActivationEpoch:            2114,
+		ExitEpoch:                  2116,
+		WithdrawableEpoch:          2117,
+	}))
+
+	rt, err = newState1.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+	pbState, err = statenative.ProtobufBeaconStatePhase0(newState1.InnerStateUnsafe())
+	require.NoError(t, err)
+
+	copiedTestState, err = statenative.InitializeFromProtoPhase0(pbState)
+	require.NoError(t, err)
+
+	rt2, err = copiedTestState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	assert.Equal(t, rt, rt2)
+}
+
+func TestBeaconState_ValidatorMutation_Altair(t *testing.T) {
+	testState, _ := util.DeterministicGenesisStateAltair(t, 400)
+	pbState, err := statenative.ProtobufBeaconStateAltair(testState.InnerStateUnsafe())
+	require.NoError(t, err)
+	testState, err = statenative.InitializeFromProtoAltair(pbState)
+	require.NoError(t, err)
+
+	_, err = testState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	// Reset tries
+	require.NoError(t, testState.UpdateValidatorAtIndex(200, new(ethpb.Validator)))
+	_, err = testState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	newState1 := testState.Copy()
+	_ = testState.Copy()
+
+	require.NoError(t, testState.UpdateValidatorAtIndex(15, &ethpb.Validator{
+		PublicKey:                  make([]byte, 48),
+		WithdrawalCredentials:      make([]byte, 32),
+		EffectiveBalance:           1111,
+		Slashed:                    false,
+		ActivationEligibilityEpoch: 1112,
+		ActivationEpoch:            1114,
+		ExitEpoch:                  1116,
+		WithdrawableEpoch:          1117,
+	}))
+
+	rt, err := testState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+	pbState, err = statenative.ProtobufBeaconStateAltair(testState.InnerStateUnsafe())
+	require.NoError(t, err)
+
+	copiedTestState, err := statenative.InitializeFromProtoAltair(pbState)
+	require.NoError(t, err)
+
+	rt2, err := copiedTestState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	assert.Equal(t, rt, rt2)
+
+	require.NoError(t, newState1.UpdateValidatorAtIndex(150, &ethpb.Validator{
+		PublicKey:                  make([]byte, 48),
+		WithdrawalCredentials:      make([]byte, 32),
+		EffectiveBalance:           2111,
+		Slashed:                    false,
+		ActivationEligibilityEpoch: 2112,
+		ActivationEpoch:            2114,
+		ExitEpoch:                  2116,
+		WithdrawableEpoch:          2117,
+	}))
+
+	rt, err = newState1.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+	pbState, err = statenative.ProtobufBeaconStateAltair(newState1.InnerStateUnsafe())
+	require.NoError(t, err)
+
+	copiedTestState, err = statenative.InitializeFromProtoAltair(pbState)
+	require.NoError(t, err)
+
+	rt2, err = copiedTestState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	assert.Equal(t, rt, rt2)
+}
+
+func TestBeaconState_ValidatorMutation_Bellatrix(t *testing.T) {
+	testState, _ := util.DeterministicGenesisStateBellatrix(t, 400)
+	pbState, err := statenative.ProtobufBeaconStateBellatrix(testState.InnerStateUnsafe())
+	require.NoError(t, err)
+	testState, err = statenative.InitializeFromProtoBellatrix(pbState)
+	require.NoError(t, err)
+
+	_, err = testState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	// Reset tries
+	require.NoError(t, testState.UpdateValidatorAtIndex(200, new(ethpb.Validator)))
+	_, err = testState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	newState1 := testState.Copy()
+	_ = testState.Copy()
+
+	require.NoError(t, testState.UpdateValidatorAtIndex(15, &ethpb.Validator{
+		PublicKey:                  make([]byte, 48),
+		WithdrawalCredentials:      make([]byte, 32),
+		EffectiveBalance:           1111,
+		Slashed:                    false,
+		ActivationEligibilityEpoch: 1112,
+		ActivationEpoch:            1114,
+		ExitEpoch:                  1116,
+		WithdrawableEpoch:          1117,
+	}))
+
+	rt, err := testState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+	pbState, err = statenative.ProtobufBeaconStateBellatrix(testState.InnerStateUnsafe())
+	require.NoError(t, err)
+
+	copiedTestState, err := statenative.InitializeFromProtoBellatrix(pbState)
+	require.NoError(t, err)
+
+	rt2, err := copiedTestState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	assert.Equal(t, rt, rt2)
+
+	require.NoError(t, newState1.UpdateValidatorAtIndex(150, &ethpb.Validator{
+		PublicKey:                  make([]byte, 48),
+		WithdrawalCredentials:      make([]byte, 32),
+		EffectiveBalance:           2111,
+		Slashed:                    false,
+		ActivationEligibilityEpoch: 2112,
+		ActivationEpoch:            2114,
+		ExitEpoch:                  2116,
+		WithdrawableEpoch:          2117,
+	}))
+
+	rt, err = newState1.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+	pbState, err = statenative.ProtobufBeaconStateBellatrix(newState1.InnerStateUnsafe())
+	require.NoError(t, err)
+
+	copiedTestState, err = statenative.InitializeFromProtoBellatrix(pbState)
+	require.NoError(t, err)
+
+	rt2, err = copiedTestState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	assert.Equal(t, rt, rt2)
+}

beacon-chain/state/v1/state_trie.go

@@ -390,17 +390,31 @@ func (b *BeaconState) rootSelector(ctx context.Context, field types.FieldIndex)
 
 func (b *BeaconState) recomputeFieldTrie(index types.FieldIndex, elements interface{}) ([32]byte, error) {
 	fTrie := b.stateFieldLeaves[index]
+	fTrieMutex := fTrie.RWMutex
 	// We can't lock the trie directly because the trie's variable gets reassigned,
 	// and therefore we would call Unlock() on a different object.
-	fTrieMutex := fTrie.RWMutex
-	if fTrie.FieldReference().Refs() > 1 {
-		fTrieMutex.Lock()
+	fTrieMutex.Lock()
+
+	if fTrie.Empty() {
+		err := b.resetFieldTrie(index, elements, fTrie.Length())
+		if err != nil {
+			fTrieMutex.Unlock()
+			return [32]byte{}, err
+		}
+		// Reduce reference count as we are instantiating a new trie.
 		fTrie.FieldReference().MinusRef()
-		newTrie := fTrie.CopyTrie()
+		fTrieMutex.Unlock()
+		return b.stateFieldLeaves[index].TrieRoot()
+	}
+
+	if fTrie.FieldReference().Refs() > 1 {
+		fTrie.FieldReference().MinusRef()
+		newTrie := fTrie.TransferTrie()
 		b.stateFieldLeaves[index] = newTrie
 		fTrie = newTrie
-		fTrieMutex.Unlock()
 	}
+	fTrieMutex.Unlock()
+
 	// remove duplicate indexes
 	b.dirtyIndices[index] = slice.SetUint64(b.dirtyIndices[index])
 	// sort indexes again

beacon-chain/state/v1/state_trie_test.go

@@ -269,3 +269,70 @@ func TestBeaconState_AppendValidator_DoesntMutateCopy(t *testing.T) {
 	_, ok := st1.ValidatorIndexByPubkey(bytesutil.ToBytes48(val.PublicKey))
 	assert.Equal(t, false, ok, "Expected no validator index to be present in st1 for the newly inserted pubkey")
 }
+
+func TestBeaconState_ValidatorMutation_Phase0(t *testing.T) {
+	testState, _ := util.DeterministicGenesisState(t, 400)
+	pbState, err := v1.ProtobufBeaconState(testState.InnerStateUnsafe())
+	require.NoError(t, err)
+	testState, err = v1.InitializeFromProto(pbState)
+	require.NoError(t, err)
+
+	_, err = testState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	// Reset tries
+	require.NoError(t, testState.UpdateValidatorAtIndex(200, new(ethpb.Validator)))
+	_, err = testState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	newState1 := testState.Copy()
+	_ = testState.Copy()
+
+	require.NoError(t, testState.UpdateValidatorAtIndex(15, &ethpb.Validator{
+		PublicKey:                  make([]byte, 48),
+		WithdrawalCredentials:      make([]byte, 32),
+		EffectiveBalance:           1111,
+		Slashed:                    false,
+		ActivationEligibilityEpoch: 1112,
+		ActivationEpoch:            1114,
+		ExitEpoch:                  1116,
+		WithdrawableEpoch:          1117,
+	}))
+
+	rt, err := testState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+	pbState, err = v1.ProtobufBeaconState(testState.InnerStateUnsafe())
+	require.NoError(t, err)
+
+	copiedTestState, err := v1.InitializeFromProto(pbState)
+	require.NoError(t, err)
+
+	rt2, err := copiedTestState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	assert.Equal(t, rt, rt2)
+
+	require.NoError(t, newState1.UpdateValidatorAtIndex(150, &ethpb.Validator{
+		PublicKey:                  make([]byte, 48),
+		WithdrawalCredentials:      make([]byte, 32),
+		EffectiveBalance:           2111,
+		Slashed:                    false,
+		ActivationEligibilityEpoch: 2112,
+		ActivationEpoch:            2114,
+		ExitEpoch:                  2116,
+		WithdrawableEpoch:          2117,
+	}))
+
+	rt, err = newState1.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+	pbState, err = v1.ProtobufBeaconState(newState1.InnerStateUnsafe())
+	require.NoError(t, err)
+
+	copiedTestState, err = v1.InitializeFromProto(pbState)
+	require.NoError(t, err)
+
+	rt2, err = copiedTestState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	assert.Equal(t, rt, rt2)
+}

beacon-chain/state/v2/BUILD.bazel

@@ -71,6 +71,7 @@ go_test(
         "references_test.go",
         "setters_test.go",
         "state_fuzz_test.go",
+        "state_test.go",
         "state_trie_test.go",
     ],
     data = glob(["testdata/**"]),

beacon-chain/state/v2/state_test.go

@@ -0,0 +1,79 @@
+package v2_test
+
+import (
+	"context"
+	"testing"
+
+	v2 "github.com/prysmaticlabs/prysm/beacon-chain/state/v2"
+	ethpb "github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1"
+	"github.com/prysmaticlabs/prysm/testing/assert"
+	"github.com/prysmaticlabs/prysm/testing/require"
+	"github.com/prysmaticlabs/prysm/testing/util"
+)
+
+func TestBeaconState_ValidatorMutation_Altair(t *testing.T) {
+	testState, _ := util.DeterministicGenesisStateAltair(t, 400)
+	pbState, err := v2.ProtobufBeaconState(testState.InnerStateUnsafe())
+	require.NoError(t, err)
+	testState, err = v2.InitializeFromProto(pbState)
+	require.NoError(t, err)
+
+	_, err = testState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	// Reset tries
+	require.NoError(t, testState.UpdateValidatorAtIndex(200, new(ethpb.Validator)))
+	_, err = testState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	newState1 := testState.Copy()
+	_ = testState.Copy()
+
+	require.NoError(t, testState.UpdateValidatorAtIndex(15, &ethpb.Validator{
+		PublicKey:                  make([]byte, 48),
+		WithdrawalCredentials:      make([]byte, 32),
+		EffectiveBalance:           1111,
+		Slashed:                    false,
+		ActivationEligibilityEpoch: 1112,
+		ActivationEpoch:            1114,
+		ExitEpoch:                  1116,
+		WithdrawableEpoch:          1117,
+	}))
+
+	rt, err := testState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+	pbState, err = v2.ProtobufBeaconState(testState.InnerStateUnsafe())
+	require.NoError(t, err)
+
+	copiedTestState, err := v2.InitializeFromProto(pbState)
+	require.NoError(t, err)
+
+	rt2, err := copiedTestState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	assert.Equal(t, rt, rt2)
+
+	require.NoError(t, newState1.UpdateValidatorAtIndex(150, &ethpb.Validator{
+		PublicKey:                  make([]byte, 48),
+		WithdrawalCredentials:      make([]byte, 32),
+		EffectiveBalance:           2111,
+		Slashed:                    false,
+		ActivationEligibilityEpoch: 2112,
+		ActivationEpoch:            2114,
+		ExitEpoch:                  2116,
+		WithdrawableEpoch:          2117,
+	}))
+
+	rt, err = newState1.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+	pbState, err = v2.ProtobufBeaconState(newState1.InnerStateUnsafe())
+	require.NoError(t, err)
+
+	copiedTestState, err = v2.InitializeFromProto(pbState)
+	require.NoError(t, err)
+
+	rt2, err = copiedTestState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	assert.Equal(t, rt, rt2)
+}

beacon-chain/state/v2/state_trie.go

@@ -378,17 +378,31 @@ func (b *BeaconState) rootSelector(ctx context.Context, field types.FieldIndex)
 
 func (b *BeaconState) recomputeFieldTrie(index types.FieldIndex, elements interface{}) ([32]byte, error) {
 	fTrie := b.stateFieldLeaves[index]
+	fTrieMutex := fTrie.RWMutex
 	// We can't lock the trie directly because the trie's variable gets reassigned,
 	// and therefore we would call Unlock() on a different object.
-	fTrieMutex := fTrie.RWMutex
-	if fTrie.FieldReference().Refs() > 1 {
-		fTrieMutex.Lock()
+	fTrieMutex.Lock()
+
+	if fTrie.Empty() {
+		err := b.resetFieldTrie(index, elements, fTrie.Length())
+		if err != nil {
+			fTrieMutex.Unlock()
+			return [32]byte{}, err
+		}
+		// Reduce reference count as we are instantiating a new trie.
 		fTrie.FieldReference().MinusRef()
-		newTrie := fTrie.CopyTrie()
+		fTrieMutex.Unlock()
+		return b.stateFieldLeaves[index].TrieRoot()
+	}
+
+	if fTrie.FieldReference().Refs() > 1 {
+		fTrie.FieldReference().MinusRef()
+		newTrie := fTrie.TransferTrie()
 		b.stateFieldLeaves[index] = newTrie
 		fTrie = newTrie
-		fTrieMutex.Unlock()
 	}
+	fTrieMutex.Unlock()
+
 	// remove duplicate indexes
 	b.dirtyIndices[index] = slice.SetUint64(b.dirtyIndices[index])
 	// sort indexes again

beacon-chain/state/v3/BUILD.bazel

@@ -73,6 +73,7 @@ go_test(
         "references_test.go",
         "setters_test.go",
         "state_fuzz_test.go",
+        "state_test.go",
         "state_trie_test.go",
     ],
     embed = [":go_default_library"],

beacon-chain/state/v3/state_test.go

@@ -0,0 +1,79 @@
+package v3_test
+
+import (
+	"context"
+	"testing"
+
+	v3 "github.com/prysmaticlabs/prysm/beacon-chain/state/v3"
+	ethpb "github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1"
+	"github.com/prysmaticlabs/prysm/testing/assert"
+	"github.com/prysmaticlabs/prysm/testing/require"
+	"github.com/prysmaticlabs/prysm/testing/util"
+)
+
+func TestBeaconState_ValidatorMutation_Bellatrix(t *testing.T) {
+	testState, _ := util.DeterministicGenesisStateBellatrix(t, 400)
+	pbState, err := v3.ProtobufBeaconState(testState.InnerStateUnsafe())
+	require.NoError(t, err)
+	testState, err = v3.InitializeFromProto(pbState)
+	require.NoError(t, err)
+
+	_, err = testState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	// Reset tries
+	require.NoError(t, testState.UpdateValidatorAtIndex(200, new(ethpb.Validator)))
+	_, err = testState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	newState1 := testState.Copy()
+	_ = testState.Copy()
+
+	require.NoError(t, testState.UpdateValidatorAtIndex(15, &ethpb.Validator{
+		PublicKey:                  make([]byte, 48),
+		WithdrawalCredentials:      make([]byte, 32),
+		EffectiveBalance:           1111,
+		Slashed:                    false,
+		ActivationEligibilityEpoch: 1112,
+		ActivationEpoch:            1114,
+		ExitEpoch:                  1116,
+		WithdrawableEpoch:          1117,
+	}))
+
+	rt, err := testState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+	pbState, err = v3.ProtobufBeaconState(testState.InnerStateUnsafe())
+	require.NoError(t, err)
+
+	copiedTestState, err := v3.InitializeFromProtoUnsafe(pbState)
+	require.NoError(t, err)
+
+	rt2, err := copiedTestState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	assert.Equal(t, rt, rt2)
+
+	require.NoError(t, newState1.UpdateValidatorAtIndex(150, &ethpb.Validator{
+		PublicKey:                  make([]byte, 48),
+		WithdrawalCredentials:      make([]byte, 32),
+		EffectiveBalance:           2111,
+		Slashed:                    false,
+		ActivationEligibilityEpoch: 2112,
+		ActivationEpoch:            2114,
+		ExitEpoch:                  2116,
+		WithdrawableEpoch:          2117,
+	}))
+
+	rt, err = newState1.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+	pbState, err = v3.ProtobufBeaconState(newState1.InnerStateUnsafe())
+	require.NoError(t, err)
+
+	copiedTestState, err = v3.InitializeFromProto(pbState)
+	require.NoError(t, err)
+
+	rt2, err = copiedTestState.HashTreeRoot(context.Background())
+	require.NoError(t, err)
+
+	assert.Equal(t, rt, rt2)
+}

beacon-chain/state/v3/state_trie.go

@@ -376,14 +376,31 @@ func (b *BeaconState) rootSelector(field types.FieldIndex) ([32]byte, error) {
 
 func (b *BeaconState) recomputeFieldTrie(index types.FieldIndex, elements interface{}) ([32]byte, error) {
 	fTrie := b.stateFieldLeaves[index]
-	if fTrie.FieldReference().Refs() > 1 {
-		fTrie.Lock()
-		defer fTrie.Unlock()
+	fTrieMutex := fTrie.RWMutex
+	// We can't lock the trie directly because the trie's variable gets reassigned,
+	// and therefore we would call Unlock() on a different object.
+	fTrieMutex.Lock()
+
+	if fTrie.Empty() {
+		err := b.resetFieldTrie(index, elements, fTrie.Length())
+		if err != nil {
+			fTrieMutex.Unlock()
+			return [32]byte{}, err
+		}
+		// Reduce reference count as we are instantiating a new trie.
 		fTrie.FieldReference().MinusRef()
-		newTrie := fTrie.CopyTrie()
+		fTrieMutex.Unlock()
+		return b.stateFieldLeaves[index].TrieRoot()
+	}
+
+	if fTrie.FieldReference().Refs() > 1 {
+		fTrie.FieldReference().MinusRef()
+		newTrie := fTrie.TransferTrie()
 		b.stateFieldLeaves[index] = newTrie
 		fTrie = newTrie
 	}
+	fTrieMutex.Unlock()
+
 	// remove duplicate indexes
 	b.dirtyIndices[index] = slice.SetUint64(b.dirtyIndices[index])
 	// sort indexes again