State: Refactor Reference and ValidatorMapHandler to stateutil pkg (#8589)

* Starting

* Fix tests

* Gazelle

beacon-chain/state/BUILD.bazel

@@ -27,7 +27,6 @@ go_library(
         "//tools/pcli:__pkg__",
     ],
     deps = [
-        "//beacon-chain/core/state/stateutils:go_default_library",
         "//beacon-chain/state/interface:go_default_library",
         "//beacon-chain/state/stateutil:go_default_library",
         "//proto/beacon/p2p/v1:go_default_library",

beacon-chain/state/field_trie.go

@@ -17,7 +17,7 @@ import (
 // trie of the particular field.
 type FieldTrie struct {
 	*sync.RWMutex
-	*reference
+	reference   *stateutil.Reference
 	fieldLayers [][]*[32]byte
 	field       fieldIndex
 }
@@ -29,7 +29,7 @@ func NewFieldTrie(field fieldIndex, elements interface{}, length uint64) (*Field
 	if elements == nil {
 		return &FieldTrie{
 			field:     field,
-			reference: &reference{refs: 1},
+			reference: stateutil.NewRef(1),
 			RWMutex:   new(sync.RWMutex),
 		}, nil
 	}
@@ -46,14 +46,14 @@ func NewFieldTrie(field fieldIndex, elements interface{}, length uint64) (*Field
 		return &FieldTrie{
 			fieldLayers: stateutil.ReturnTrieLayer(fieldRoots, length),
 			field:       field,
-			reference:   &reference{refs: 1},
+			reference:   stateutil.NewRef(1),
 			RWMutex:     new(sync.RWMutex),
 		}, nil
 	case compositeArray:
 		return &FieldTrie{
 			fieldLayers: stateutil.ReturnTrieLayerVariable(fieldRoots, length),
 			field:       field,
-			reference:   &reference{refs: 1},
+			reference:   stateutil.NewRef(1),
 			RWMutex:     new(sync.RWMutex),
 		}, nil
 	default:
@@ -105,7 +105,7 @@ func (f *FieldTrie) CopyTrie() *FieldTrie {
 	if f.fieldLayers == nil {
 		return &FieldTrie{
 			field:     f.field,
-			reference: &reference{refs: 1},
+			reference: stateutil.NewRef(1),
 			RWMutex:   new(sync.RWMutex),
 		}
 	}
@@ -117,7 +117,7 @@ func (f *FieldTrie) CopyTrie() *FieldTrie {
 	return &FieldTrie{
 		fieldLayers: dstFieldTrie,
 		field:       f.field,
-		reference:   &reference{refs: 1},
+		reference:   stateutil.NewRef(1),
 		RWMutex:     new(sync.RWMutex),
 	}
 }

beacon-chain/state/getters.go

@@ -524,23 +524,23 @@ func (b *BeaconState) ValidatorAtIndexReadOnly(idx types.ValidatorIndex) (iface.
 
 // ValidatorIndexByPubkey returns a given validator by its 48-byte public key.
 func (b *BeaconState) ValidatorIndexByPubkey(key [48]byte) (types.ValidatorIndex, bool) {
-	if b == nil || b.valMapHandler == nil || b.valMapHandler.valIdxMap == nil {
+	if b == nil || b.valMapHandler == nil || b.valMapHandler.IsNil() {
 		return 0, false
 	}
 	b.lock.RLock()
 	defer b.lock.RUnlock()
-	idx, ok := b.valMapHandler.valIdxMap[key]
+	idx, ok := b.valMapHandler.Get(key)
 	return idx, ok
 }
 
 func (b *BeaconState) validatorIndexMap() map[[48]byte]types.ValidatorIndex {
-	if b == nil || b.valMapHandler == nil || b.valMapHandler.valIdxMap == nil {
+	if b == nil || b.valMapHandler == nil || b.valMapHandler.IsNil() {
 		return map[[48]byte]types.ValidatorIndex{}
 	}
 	b.lock.RLock()
 	defer b.lock.RUnlock()
 
-	return b.valMapHandler.copy().valIdxMap
+	return b.valMapHandler.Copy().ValidatorIndexMap()
 }
 
 // PubkeyAtIndex returns the pubkey at the given

beacon-chain/state/references_test.go

@@ -20,24 +20,24 @@ func TestStateReferenceSharing_Finalizer(t *testing.T) {
 
 	a, err := InitializeFromProtoUnsafe(&p2ppb.BeaconState{RandaoMixes: [][]byte{[]byte("foo")}})
 	require.NoError(t, err)
-	assert.Equal(t, uint(1), a.sharedFieldReferences[randaoMixes].refs, "Expected a single reference for RANDAO mixes")
+	assert.Equal(t, uint(1), a.sharedFieldReferences[randaoMixes].Refs(), "Expected a single reference for RANDAO mixes")
 
 	func() {
 		// Create object in a different scope for GC
 		b := a.Copy()
-		assert.Equal(t, uint(2), a.sharedFieldReferences[randaoMixes].refs, "Expected 2 references to RANDAO mixes")
+		assert.Equal(t, uint(2), a.sharedFieldReferences[randaoMixes].Refs(), "Expected 2 references to RANDAO mixes")
 		_ = b
 	}()
 
 	runtime.GC() // Should run finalizer on object b
-	assert.Equal(t, uint(1), a.sharedFieldReferences[randaoMixes].refs, "Expected 1 shared reference to RANDAO mixes!")
+	assert.Equal(t, uint(1), a.sharedFieldReferences[randaoMixes].Refs(), "Expected 1 shared reference to RANDAO mixes!")
 
 	copied := a.Copy()
 	b, ok := copied.(*BeaconState)
 	require.Equal(t, true, ok)
-	assert.Equal(t, uint(2), b.sharedFieldReferences[randaoMixes].refs, "Expected 2 shared references to RANDAO mixes")
+	assert.Equal(t, uint(2), b.sharedFieldReferences[randaoMixes].Refs(), "Expected 2 shared references to RANDAO mixes")
 	require.NoError(t, b.UpdateRandaoMixesAtIndex(0, []byte("bar")))
-	if b.sharedFieldReferences[randaoMixes].refs != 1 || a.sharedFieldReferences[randaoMixes].refs != 1 {
+	if b.sharedFieldReferences[randaoMixes].Refs() != 1 || a.sharedFieldReferences[randaoMixes].Refs() != 1 {
 		t.Error("Expected 1 shared reference to RANDAO mix for both a and b")
 	}
 }
@@ -318,7 +318,7 @@ func TestValidatorReferences_RemainsConsistent(t *testing.T) {
 // assertRefCount checks whether reference count for a given state
 // at a given index is equal to expected amount.
 func assertRefCount(t *testing.T, b *BeaconState, idx fieldIndex, want uint) {
-	if cnt := b.sharedFieldReferences[idx].refs; cnt != want {
+	if cnt := b.sharedFieldReferences[idx].Refs(); cnt != want {
 		t.Errorf("Unexpected count of references for index %d, want: %v, got: %v", idx, want, cnt)
 	}
 }

beacon-chain/state/setters.go

@@ -8,7 +8,7 @@ import (
 	types "github.com/prysmaticlabs/eth2-types"
 	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
 	"github.com/prysmaticlabs/go-bitfield"
-	coreutils "github.com/prysmaticlabs/prysm/beacon-chain/core/state/stateutils"
+	"github.com/prysmaticlabs/prysm/beacon-chain/state/stateutil"
 	pbp2p "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
 	"github.com/prysmaticlabs/prysm/shared/bytesutil"
 	"github.com/prysmaticlabs/prysm/shared/hashutil"
@@ -107,7 +107,7 @@ func (b *BeaconState) SetBlockRoots(val [][]byte) error {
 	defer b.lock.Unlock()
 
 	b.sharedFieldReferences[blockRoots].MinusRef()
-	b.sharedFieldReferences[blockRoots] = &reference{refs: 1}
+	b.sharedFieldReferences[blockRoots] = stateutil.NewRef(1)
 
 	b.state.BlockRoots = val
 	b.markFieldAsDirty(blockRoots)
@@ -133,7 +133,7 @@ func (b *BeaconState) UpdateBlockRootAtIndex(idx uint64, blockRoot [32]byte) err
 		r = make([][]byte, len(b.state.BlockRoots))
 		copy(r, b.state.BlockRoots)
 		ref.MinusRef()
-		b.sharedFieldReferences[blockRoots] = &reference{refs: 1}
+		b.sharedFieldReferences[blockRoots] = stateutil.NewRef(1)
 	}
 
 	r[idx] = blockRoot[:]
@@ -154,7 +154,7 @@ func (b *BeaconState) SetStateRoots(val [][]byte) error {
 	defer b.lock.Unlock()
 
 	b.sharedFieldReferences[stateRoots].MinusRef()
-	b.sharedFieldReferences[stateRoots] = &reference{refs: 1}
+	b.sharedFieldReferences[stateRoots] = stateutil.NewRef(1)
 
 	b.state.StateRoots = val
 	b.markFieldAsDirty(stateRoots)
@@ -186,7 +186,7 @@ func (b *BeaconState) UpdateStateRootAtIndex(idx uint64, stateRoot [32]byte) err
 		r = make([][]byte, len(b.state.StateRoots))
 		copy(r, b.state.StateRoots)
 		ref.MinusRef()
-		b.sharedFieldReferences[stateRoots] = &reference{refs: 1}
+		b.sharedFieldReferences[stateRoots] = stateutil.NewRef(1)
 	}
 
 	r[idx] = stateRoot[:]
@@ -207,7 +207,7 @@ func (b *BeaconState) SetHistoricalRoots(val [][]byte) error {
 	defer b.lock.Unlock()
 
 	b.sharedFieldReferences[historicalRoots].MinusRef()
-	b.sharedFieldReferences[historicalRoots] = &reference{refs: 1}
+	b.sharedFieldReferences[historicalRoots] = stateutil.NewRef(1)
 
 	b.state.HistoricalRoots = val
 	b.markFieldAsDirty(historicalRoots)
@@ -237,7 +237,7 @@ func (b *BeaconState) SetEth1DataVotes(val []*ethpb.Eth1Data) error {
 	defer b.lock.Unlock()
 
 	b.sharedFieldReferences[eth1DataVotes].MinusRef()
-	b.sharedFieldReferences[eth1DataVotes] = &reference{refs: 1}
+	b.sharedFieldReferences[eth1DataVotes] = stateutil.NewRef(1)
 
 	b.state.Eth1DataVotes = val
 	b.markFieldAsDirty(eth1DataVotes)
@@ -260,7 +260,7 @@ func (b *BeaconState) AppendEth1DataVotes(val *ethpb.Eth1Data) error {
 		votes = make([]*ethpb.Eth1Data, len(b.state.Eth1DataVotes))
 		copy(votes, b.state.Eth1DataVotes)
 		b.sharedFieldReferences[eth1DataVotes].MinusRef()
-		b.sharedFieldReferences[eth1DataVotes] = &reference{refs: 1}
+		b.sharedFieldReferences[eth1DataVotes] = stateutil.NewRef(1)
 	}
 
 	b.state.Eth1DataVotes = append(votes, val)
@@ -293,13 +293,10 @@ func (b *BeaconState) SetValidators(val []*ethpb.Validator) error {
 
 	b.state.Validators = val
 	b.sharedFieldReferences[validators].MinusRef()
-	b.sharedFieldReferences[validators] = &reference{refs: 1}
+	b.sharedFieldReferences[validators] = stateutil.NewRef(1)
 	b.markFieldAsDirty(validators)
 	b.rebuildTrie[validators] = true
-	b.valMapHandler = &validatorMapHandler{
-		valIdxMap: coreutils.ValidatorIndexMap(b.state.Validators),
-		mapRef:    &reference{refs: 1},
-	}
+	b.valMapHandler = stateutil.NewValMapHandler(b.state.Validators)
 	return nil
 }
 
@@ -314,7 +311,7 @@ func (b *BeaconState) ApplyToEveryValidator(f func(idx int, val *ethpb.Validator
 	if ref := b.sharedFieldReferences[validators]; ref.Refs() > 1 {
 		v = b.validatorsReferences()
 		ref.MinusRef()
-		b.sharedFieldReferences[validators] = &reference{refs: 1}
+		b.sharedFieldReferences[validators] = stateutil.NewRef(1)
 	}
 	b.lock.Unlock()
 	var changedVals []uint64
@@ -355,7 +352,7 @@ func (b *BeaconState) UpdateValidatorAtIndex(idx types.ValidatorIndex, val *ethp
 	if ref := b.sharedFieldReferences[validators]; ref.Refs() > 1 {
 		v = b.validatorsReferences()
 		ref.MinusRef()
-		b.sharedFieldReferences[validators] = &reference{refs: 1}
+		b.sharedFieldReferences[validators] = stateutil.NewRef(1)
 	}
 
 	v[idx] = val
@@ -376,7 +373,7 @@ func (b *BeaconState) SetBalances(val []uint64) error {
 	defer b.lock.Unlock()
 
 	b.sharedFieldReferences[balances].MinusRef()
-	b.sharedFieldReferences[balances] = &reference{refs: 1}
+	b.sharedFieldReferences[balances] = stateutil.NewRef(1)
 
 	b.state.Balances = val
 	b.markFieldAsDirty(balances)
@@ -399,7 +396,7 @@ func (b *BeaconState) UpdateBalancesAtIndex(idx types.ValidatorIndex, val uint64
 	if b.sharedFieldReferences[balances].Refs() > 1 {
 		bals = b.balances()
 		b.sharedFieldReferences[balances].MinusRef()
-		b.sharedFieldReferences[balances] = &reference{refs: 1}
+		b.sharedFieldReferences[balances] = stateutil.NewRef(1)
 	}
 
 	bals[idx] = val
@@ -418,7 +415,7 @@ func (b *BeaconState) SetRandaoMixes(val [][]byte) error {
 	defer b.lock.Unlock()
 
 	b.sharedFieldReferences[randaoMixes].MinusRef()
-	b.sharedFieldReferences[randaoMixes] = &reference{refs: 1}
+	b.sharedFieldReferences[randaoMixes] = stateutil.NewRef(1)
 
 	b.state.RandaoMixes = val
 	b.markFieldAsDirty(randaoMixes)
@@ -444,7 +441,7 @@ func (b *BeaconState) UpdateRandaoMixesAtIndex(idx uint64, val []byte) error {
 		mixes = make([][]byte, len(b.state.RandaoMixes))
 		copy(mixes, b.state.RandaoMixes)
 		b.sharedFieldReferences[randaoMixes].MinusRef()
-		b.sharedFieldReferences[randaoMixes] = &reference{refs: 1}
+		b.sharedFieldReferences[randaoMixes] = stateutil.NewRef(1)
 	}
 
 	mixes[idx] = val
@@ -465,7 +462,7 @@ func (b *BeaconState) SetSlashings(val []uint64) error {
 	defer b.lock.Unlock()
 
 	b.sharedFieldReferences[slashings].MinusRef()
-	b.sharedFieldReferences[slashings] = &reference{refs: 1}
+	b.sharedFieldReferences[slashings] = stateutil.NewRef(1)
 
 	b.state.Slashings = val
 	b.markFieldAsDirty(slashings)
@@ -488,7 +485,7 @@ func (b *BeaconState) UpdateSlashingsAtIndex(idx, val uint64) error {
 	if b.sharedFieldReferences[slashings].Refs() > 1 {
 		s = b.slashings()
 		b.sharedFieldReferences[slashings].MinusRef()
-		b.sharedFieldReferences[slashings] = &reference{refs: 1}
+		b.sharedFieldReferences[slashings] = stateutil.NewRef(1)
 	}
 
 	s[idx] = val
@@ -509,7 +506,7 @@ func (b *BeaconState) SetPreviousEpochAttestations(val []*pbp2p.PendingAttestati
 	defer b.lock.Unlock()
 
 	b.sharedFieldReferences[previousEpochAttestations].MinusRef()
-	b.sharedFieldReferences[previousEpochAttestations] = &reference{refs: 1}
+	b.sharedFieldReferences[previousEpochAttestations] = stateutil.NewRef(1)
 
 	b.state.PreviousEpochAttestations = val
 	b.markFieldAsDirty(previousEpochAttestations)
@@ -527,7 +524,7 @@ func (b *BeaconState) SetCurrentEpochAttestations(val []*pbp2p.PendingAttestatio
 	defer b.lock.Unlock()
 
 	b.sharedFieldReferences[currentEpochAttestations].MinusRef()
-	b.sharedFieldReferences[currentEpochAttestations] = &reference{refs: 1}
+	b.sharedFieldReferences[currentEpochAttestations] = stateutil.NewRef(1)
 
 	b.state.CurrentEpochAttestations = val
 	b.markFieldAsDirty(currentEpochAttestations)
@@ -549,7 +546,7 @@ func (b *BeaconState) AppendHistoricalRoots(root [32]byte) error {
 		roots = make([][]byte, len(b.state.HistoricalRoots))
 		copy(roots, b.state.HistoricalRoots)
 		b.sharedFieldReferences[historicalRoots].MinusRef()
-		b.sharedFieldReferences[historicalRoots] = &reference{refs: 1}
+		b.sharedFieldReferences[historicalRoots] = stateutil.NewRef(1)
 	}
 
 	b.state.HistoricalRoots = append(roots, root[:])
@@ -572,7 +569,7 @@ func (b *BeaconState) AppendCurrentEpochAttestations(val *pbp2p.PendingAttestati
 		atts = make([]*pbp2p.PendingAttestation, len(b.state.CurrentEpochAttestations))
 		copy(atts, b.state.CurrentEpochAttestations)
 		b.sharedFieldReferences[currentEpochAttestations].MinusRef()
-		b.sharedFieldReferences[currentEpochAttestations] = &reference{refs: 1}
+		b.sharedFieldReferences[currentEpochAttestations] = stateutil.NewRef(1)
 	}
 
 	b.state.CurrentEpochAttestations = append(atts, val)
@@ -595,7 +592,7 @@ func (b *BeaconState) AppendPreviousEpochAttestations(val *pbp2p.PendingAttestat
 		atts = make([]*pbp2p.PendingAttestation, len(b.state.PreviousEpochAttestations))
 		copy(atts, b.state.PreviousEpochAttestations)
 		b.sharedFieldReferences[previousEpochAttestations].MinusRef()
-		b.sharedFieldReferences[previousEpochAttestations] = &reference{refs: 1}
+		b.sharedFieldReferences[previousEpochAttestations] = stateutil.NewRef(1)
 	}
 
 	b.state.PreviousEpochAttestations = append(atts, val)
@@ -618,7 +615,7 @@ func (b *BeaconState) AppendValidator(val *ethpb.Validator) error {
 	if b.sharedFieldReferences[validators].Refs() > 1 {
 		vals = b.validatorsReferences()
 		b.sharedFieldReferences[validators].MinusRef()
-		b.sharedFieldReferences[validators] = &reference{refs: 1}
+		b.sharedFieldReferences[validators] = stateutil.NewRef(1)
 	}
 
 	// append validator to slice
@@ -626,12 +623,12 @@ func (b *BeaconState) AppendValidator(val *ethpb.Validator) error {
 	valIdx := types.ValidatorIndex(len(b.state.Validators) - 1)
 
 	// Copy if this is a shared validator map
-	if ref := b.valMapHandler.mapRef; ref.Refs() > 1 {
-		valMap := b.valMapHandler.copy()
+	if ref := b.valMapHandler.MapRef(); ref.Refs() > 1 {
+		valMap := b.valMapHandler.Copy()
 		ref.MinusRef()
 		b.valMapHandler = valMap
 	}
-	b.valMapHandler.valIdxMap[bytesutil.ToBytes48(val.PublicKey)] = valIdx
+	b.valMapHandler.Set(bytesutil.ToBytes48(val.PublicKey), valIdx)
 
 	b.markFieldAsDirty(validators)
 	b.addDirtyIndices(validators, []uint64{uint64(valIdx)})
@@ -651,7 +648,7 @@ func (b *BeaconState) AppendBalance(bal uint64) error {
 	if b.sharedFieldReferences[balances].Refs() > 1 {
 		bals = b.balances()
 		b.sharedFieldReferences[balances].MinusRef()
-		b.sharedFieldReferences[balances] = &reference{refs: 1}
+		b.sharedFieldReferences[balances] = stateutil.NewRef(1)
 	}
 
 	b.state.Balances = append(bals, bal)

beacon-chain/state/state_test.go

@@ -6,6 +6,7 @@ import (
 	"testing"
 
 	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
+	"github.com/prysmaticlabs/prysm/beacon-chain/state/stateutil"
 	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
 	"github.com/prysmaticlabs/prysm/shared/bytesutil"
 	"github.com/prysmaticlabs/prysm/shared/params"
@@ -31,11 +32,13 @@ func TestValidatorMap_DistinctCopy(t *testing.T) {
 			WithdrawableEpoch:          1,
 		})
 	}
-	handler := newValHandler(vals)
-	newHandler := handler.copy()
+	handler := stateutil.NewValMapHandler(vals)
+	newHandler := handler.Copy()
 	wantedPubkey := strconv.Itoa(22)
-	handler.valIdxMap[bytesutil.ToBytes48([]byte(wantedPubkey))] = 27
-	assert.NotEqual(t, handler.valIdxMap[bytesutil.ToBytes48([]byte(wantedPubkey))], newHandler.valIdxMap[bytesutil.ToBytes48([]byte(wantedPubkey))], "Values are supposed to be unequal due to copy")
+	handler.Set(bytesutil.ToBytes48([]byte(wantedPubkey)), 27)
+	val1, _ := handler.Get(bytesutil.ToBytes48([]byte(wantedPubkey)))
+	val2, _ := newHandler.Get(bytesutil.ToBytes48([]byte(wantedPubkey)))
+	assert.NotEqual(t, val1, val2, "Values are supposed to be unequal due to copy")
 }
 
 func TestBeaconState_NoDeadlock(t *testing.T) {
@@ -75,7 +78,7 @@ func TestBeaconState_NoDeadlock(t *testing.T) {
 					f.fieldLayers = make([][]*[32]byte, 10)
 				}
 				f.Unlock()
-				f.AddRef()
+				f.reference.AddRef()
 			}
 		}
 		wg.Done()

beacon-chain/state/state_trie.go

@@ -37,9 +37,9 @@ func InitializeFromProtoUnsafe(st *pbp2p.BeaconState) (*BeaconState, error) {
 		dirtyFields:           make(map[fieldIndex]interface{}, fieldCount),
 		dirtyIndices:          make(map[fieldIndex][]uint64, fieldCount),
 		stateFieldLeaves:      make(map[fieldIndex]*FieldTrie, fieldCount),
-		sharedFieldReferences: make(map[fieldIndex]*reference, 10),
+		sharedFieldReferences: make(map[fieldIndex]*stateutil.Reference, 10),
 		rebuildTrie:           make(map[fieldIndex]bool, fieldCount),
-		valMapHandler:         newValHandler(st.Validators),
+		valMapHandler:         stateutil.NewValMapHandler(st.Validators),
 	}
 
 	for i := 0; i < fieldCount; i++ {
@@ -48,22 +48,22 @@ func InitializeFromProtoUnsafe(st *pbp2p.BeaconState) (*BeaconState, error) {
 		b.dirtyIndices[fieldIndex(i)] = []uint64{}
 		b.stateFieldLeaves[fieldIndex(i)] = &FieldTrie{
 			field:     fieldIndex(i),
-			reference: &reference{refs: 1},
+			reference: stateutil.NewRef(1),
 			RWMutex:   new(sync.RWMutex),
 		}
 	}
 
 	// Initialize field reference tracking for shared data.
-	b.sharedFieldReferences[randaoMixes] = &reference{refs: 1}
-	b.sharedFieldReferences[stateRoots] = &reference{refs: 1}
-	b.sharedFieldReferences[blockRoots] = &reference{refs: 1}
-	b.sharedFieldReferences[previousEpochAttestations] = &reference{refs: 1}
-	b.sharedFieldReferences[currentEpochAttestations] = &reference{refs: 1}
-	b.sharedFieldReferences[slashings] = &reference{refs: 1}
-	b.sharedFieldReferences[eth1DataVotes] = &reference{refs: 1}
-	b.sharedFieldReferences[validators] = &reference{refs: 1}
-	b.sharedFieldReferences[balances] = &reference{refs: 1}
-	b.sharedFieldReferences[historicalRoots] = &reference{refs: 1}
+	b.sharedFieldReferences[randaoMixes] = stateutil.NewRef(1)
+	b.sharedFieldReferences[stateRoots] = stateutil.NewRef(1)
+	b.sharedFieldReferences[blockRoots] = stateutil.NewRef(1)
+	b.sharedFieldReferences[previousEpochAttestations] = stateutil.NewRef(1)
+	b.sharedFieldReferences[currentEpochAttestations] = stateutil.NewRef(1)
+	b.sharedFieldReferences[slashings] = stateutil.NewRef(1)
+	b.sharedFieldReferences[eth1DataVotes] = stateutil.NewRef(1)
+	b.sharedFieldReferences[validators] = stateutil.NewRef(1)
+	b.sharedFieldReferences[balances] = stateutil.NewRef(1)
+	b.sharedFieldReferences[historicalRoots] = stateutil.NewRef(1)
 
 	return b, nil
 }
@@ -111,7 +111,7 @@ func (b *BeaconState) Copy() iface.BeaconState {
 		dirtyFields:           make(map[fieldIndex]interface{}, fieldCount),
 		dirtyIndices:          make(map[fieldIndex][]uint64, fieldCount),
 		rebuildTrie:           make(map[fieldIndex]bool, fieldCount),
-		sharedFieldReferences: make(map[fieldIndex]*reference, 10),
+		sharedFieldReferences: make(map[fieldIndex]*stateutil.Reference, 10),
 		stateFieldLeaves:      make(map[fieldIndex]*FieldTrie, fieldCount),
 
 		// Copy on write validator index map.
@@ -124,7 +124,7 @@ func (b *BeaconState) Copy() iface.BeaconState {
 	}
 
 	// Increment ref for validator map
-	b.valMapHandler.mapRef.AddRef()
+	b.valMapHandler.AddRef()
 
 	for i := range b.dirtyFields {
 		dst.dirtyFields[i] = true
@@ -144,7 +144,7 @@ func (b *BeaconState) Copy() iface.BeaconState {
 		dst.stateFieldLeaves[fldIdx] = fieldTrie
 		if fieldTrie.reference != nil {
 			fieldTrie.Lock()
-			fieldTrie.AddRef()
+			fieldTrie.reference.AddRef()
 			fieldTrie.Unlock()
 		}
 	}
@@ -165,7 +165,7 @@ func (b *BeaconState) Copy() iface.BeaconState {
 		for field, v := range b.sharedFieldReferences {
 			v.MinusRef()
 			if b.stateFieldLeaves[field].reference != nil {
-				b.stateFieldLeaves[field].MinusRef()
+				b.stateFieldLeaves[field].reference.MinusRef()
 			}
 		}
 	})
@@ -214,7 +214,7 @@ func (b *BeaconState) FieldReferencesCount() map[string]uint64 {
 		refMap[i.String()] = uint64(f.Refs())
 	}
 	for i, f := range b.stateFieldLeaves {
-		numOfRefs := uint64(f.Refs())
+		numOfRefs := uint64(f.reference.Refs())
 		f.RLock()
 		if len(f.fieldLayers) != 0 {
 			refMap[i.String()+"_trie"] = numOfRefs
@@ -369,10 +369,10 @@ func (b *BeaconState) rootSelector(field fieldIndex) ([32]byte, error) {
 
 func (b *BeaconState) recomputeFieldTrie(index fieldIndex, elements interface{}) ([32]byte, error) {
 	fTrie := b.stateFieldLeaves[index]
-	if fTrie.Refs() > 1 {
+	if fTrie.reference.Refs() > 1 {
 		fTrie.Lock()
 		defer fTrie.Unlock()
-		fTrie.MinusRef()
+		fTrie.reference.MinusRef()
 		newTrie := fTrie.CopyTrie()
 		b.stateFieldLeaves[index] = newTrie
 		fTrie = newTrie

beacon-chain/state/stateutil/BUILD.bazel

@@ -7,8 +7,10 @@ go_library(
         "arrays.go",
         "attestations.go",
         "blocks.go",
+        "reference.go",
         "state_root.go",
         "trie_helpers.go",
+        "validator_map_handler.go",
         "validators.go",
     ],
     importpath = "github.com/prysmaticlabs/prysm/beacon-chain/state/stateutil",
@@ -25,6 +27,7 @@ go_library(
         "//validator/client:__pkg__",
     ],
     deps = [
+        "//beacon-chain/core/state/stateutils:go_default_library",
         "//proto/beacon/p2p/v1:go_default_library",
         "//shared/bytesutil:go_default_library",
         "//shared/featureconfig:go_default_library",
@@ -34,6 +37,7 @@ go_library(
         "//shared/trieutil:go_default_library",
         "@com_github_dgraph_io_ristretto//:go_default_library",
         "@com_github_pkg_errors//:go_default_library",
+        "@com_github_prysmaticlabs_eth2_types//:go_default_library",
         "@com_github_prysmaticlabs_ethereumapis//eth/v1alpha1:go_default_library",
     ],
 )
@@ -43,6 +47,7 @@ go_test(
     srcs = [
         "attestations_test.go",
         "benchmark_test.go",
+        "reference_bench_test.go",
         "state_root_test.go",
         "stateutil_test.go",
         "trie_helpers_test.go",

beacon-chain/state/stateutil/reference.go

@@ -0,0 +1,45 @@
+package stateutil
+
+import "sync"
+
+// Reference structs are shared across BeaconState copies to understand when the state must use
+// copy-on-write for shared fields or may modify a field in place when it holds the only reference
+// to the field value. References are tracked in a map of fieldIndex -> *reference. Whenever a state
+// releases their reference to the field value, they must decrement the refs. Likewise whenever a
+// copy is performed then the state must increment the refs counter.
+type Reference struct {
+	refs uint
+	lock sync.RWMutex
+}
+
+// NewRef initializes the Reference struct.
+func NewRef(refs uint) *Reference {
+	return &Reference{
+		refs: refs,
+	}
+}
+
+// Refs returns the reference number.
+func (r *Reference) Refs() uint {
+	r.lock.RLock()
+	defer r.lock.RUnlock()
+	return r.refs
+}
+
+// AddRef adds 1 to the reference number.
+func (r *Reference) AddRef() {
+	r.lock.Lock()
+	r.refs++
+	r.lock.Unlock()
+}
+
+// MinusRef subtracts 1 to the reference number.
+func (r *Reference) MinusRef() {
+	r.lock.Lock()
+	// Do not reduce further if object
+	// already has 0 reference to prevent underflow.
+	if r.refs > 0 {
+		r.refs--
+	}
+	r.lock.Unlock()
+}

beacon-chain/state/stateutil/reference_bench_test.go

@@ -1,4 +1,4 @@
-package state
+package stateutil
 
 import (
 	"math"
@@ -6,7 +6,7 @@ import (
 )
 
 func BenchmarkReference_MinusRef(b *testing.B) {
-	ref := &reference{
+	ref := &Reference{
 		refs: math.MaxUint64,
 	}
 	for i := 0; i < b.N; i++ {

beacon-chain/state/stateutil/validator_map_handler.go

@@ -0,0 +1,71 @@
+package stateutil
+
+import (
+	types "github.com/prysmaticlabs/eth2-types"
+	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
+	coreutils "github.com/prysmaticlabs/prysm/beacon-chain/core/state/stateutils"
+)
+
+// ValidatorMapHandler is a container to hold the map and a reference tracker for how many
+// states shared this.
+type ValidatorMapHandler struct {
+	valIdxMap map[[48]byte]types.ValidatorIndex
+	mapRef    *Reference
+}
+
+// NewValMapHandler returns a new validator map handler.
+func NewValMapHandler(vals []*ethpb.Validator) *ValidatorMapHandler {
+	return &ValidatorMapHandler{
+		valIdxMap: coreutils.ValidatorIndexMap(vals),
+		mapRef:    &Reference{refs: 1},
+	}
+}
+
+// Copy the whole map and returns a map handler with the copied map.
+func (v *ValidatorMapHandler) AddRef() {
+	v.mapRef.AddRef()
+}
+
+// ValidatorIndexMap returns the validator index map.
+func (v *ValidatorMapHandler) ValidatorIndexMap() map[[48]byte]types.ValidatorIndex {
+	return v.valIdxMap
+}
+
+// MapRef returns the map reference.
+func (v *ValidatorMapHandler) MapRef() *Reference {
+	return v.mapRef
+}
+
+// IsNil returns true if the underlying validator index map is nil.
+func (v *ValidatorMapHandler) IsNil() bool {
+	return v.mapRef == nil
+}
+
+// Copy the whole map and returns a map handler with the copied map.
+func (v *ValidatorMapHandler) Copy() *ValidatorMapHandler {
+	if v == nil || v.valIdxMap == nil {
+		return &ValidatorMapHandler{valIdxMap: map[[48]byte]types.ValidatorIndex{}, mapRef: new(Reference)}
+	}
+	m := make(map[[48]byte]types.ValidatorIndex, len(v.valIdxMap))
+	for k, v := range v.valIdxMap {
+		m[k] = v
+	}
+	return &ValidatorMapHandler{
+		valIdxMap: m,
+		mapRef:    &Reference{refs: 1},
+	}
+}
+
+// Get the validator index using the corresponding public key.
+func (v *ValidatorMapHandler) Get(key [48]byte) (types.ValidatorIndex, bool) {
+	idx, ok := v.valIdxMap[key]
+	if !ok {
+		return 0, false
+	}
+	return idx, true
+}
+
+// Set the validator index using the corresponding public key.
+func (v *ValidatorMapHandler) Set(key [48]byte, index types.ValidatorIndex) {
+	v.valIdxMap[key] = index
+}

beacon-chain/state/types.go

@@ -4,10 +4,8 @@ import (
 	"sync"
 
 	"github.com/pkg/errors"
-	types "github.com/prysmaticlabs/eth2-types"
-	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
-	coreutils "github.com/prysmaticlabs/prysm/beacon-chain/core/state/stateutils"
 	iface "github.com/prysmaticlabs/prysm/beacon-chain/state/interface"
+	"github.com/prysmaticlabs/prysm/beacon-chain/state/stateutil"
 	pbp2p "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
 	"github.com/prysmaticlabs/prysm/shared/params"
 )
@@ -74,16 +72,6 @@ const (
 // to its corresponding data type.
 var fieldMap map[fieldIndex]dataType
 
-// Reference structs are shared across BeaconState copies to understand when the state must use
-// copy-on-write for shared fields or may modify a field in place when it holds the only reference
-// to the field value. References are tracked in a map of fieldIndex -> *reference. Whenever a state
-// releases their reference to the field value, they must decrement the refs. Likewise whenever a
-// copy is performed then the state must increment the refs counter.
-type reference struct {
-	refs uint
-	lock sync.RWMutex
-}
-
 // ErrNilInnerState returns when the inner state is nil and no copy set or get
 // operations can be performed on state.
 var ErrNilInnerState = errors.New("nil inner state")
@@ -97,9 +85,9 @@ type BeaconState struct {
 	dirtyIndices          map[fieldIndex][]uint64
 	stateFieldLeaves      map[fieldIndex]*FieldTrie
 	rebuildTrie           map[fieldIndex]bool
-	valMapHandler         *validatorMapHandler
+	valMapHandler         *stateutil.ValidatorMapHandler
 	merkleLayers          [][][]byte
-	sharedFieldReferences map[fieldIndex]*reference
+	sharedFieldReferences map[fieldIndex]*stateutil.Reference
 }
 
 // String returns the name of the field index.
@@ -151,61 +139,3 @@ func (f fieldIndex) String() string {
 		return ""
 	}
 }
-
-// ReadOnlyValidator returns a wrapper that only allows fields from a validator
-// to be read, and prevents any modification of internal validator fields.
-type ReadOnlyValidator struct {
-	validator *ethpb.Validator
-}
-
-func (r *reference) Refs() uint {
-	r.lock.RLock()
-	defer r.lock.RUnlock()
-	return r.refs
-}
-
-func (r *reference) AddRef() {
-	r.lock.Lock()
-	r.refs++
-	r.lock.Unlock()
-}
-
-func (r *reference) MinusRef() {
-	r.lock.Lock()
-	// Do not reduce further if object
-	// already has 0 reference to prevent underflow.
-	if r.refs > 0 {
-		r.refs--
-	}
-	r.lock.Unlock()
-}
-
-// a container to hold the map and a reference tracker for how many
-// states shared this.
-type validatorMapHandler struct {
-	valIdxMap map[[48]byte]types.ValidatorIndex
-	mapRef    *reference
-}
-
-// A constructor for the map handler.
-func newValHandler(vals []*ethpb.Validator) *validatorMapHandler {
-	return &validatorMapHandler{
-		valIdxMap: coreutils.ValidatorIndexMap(vals),
-		mapRef:    &reference{refs: 1},
-	}
-}
-
-// copies the whole map and returns a map handler with the copied map.
-func (v *validatorMapHandler) copy() *validatorMapHandler {
-	if v == nil || v.valIdxMap == nil {
-		return &validatorMapHandler{valIdxMap: map[[48]byte]types.ValidatorIndex{}, mapRef: new(reference)}
-	}
-	m := make(map[[48]byte]types.ValidatorIndex, len(v.valIdxMap))
-	for k, v := range v.valIdxMap {
-		m[k] = v
-	}
-	return &validatorMapHandler{
-		valIdxMap: m,
-		mapRef:    &reference{refs: 1},
-	}
-}

beacon-chain/state/validator_getters.go

@@ -2,8 +2,15 @@ package state
 
 import (
 	types "github.com/prysmaticlabs/eth2-types"
+	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
 )
 
+// ReadOnlyValidator returns a wrapper that only allows fields from a validator
+// to be read, and prevents any modification of internal validator fields.
+type ReadOnlyValidator struct {
+	validator *ethpb.Validator
+}
+
 // EffectiveBalance returns the effective balance of the
 // read only validator.
 func (v ReadOnlyValidator) EffectiveBalance() uint64 {