github.com/prysmaticlabs/prysm@v1.4.4/beacon-chain/state/stateutil/validator_map_handler.go (about)

     1  package stateutil
     2  
     3  import (
     4  	"sync"
     5  
     6  	types "github.com/prysmaticlabs/eth2-types"
     7  	coreutils "github.com/prysmaticlabs/prysm/beacon-chain/core/state/stateutils"
     8  	ethpb "github.com/prysmaticlabs/prysm/proto/eth/v1alpha1"
     9  )
    10  
    11  // ValidatorMapHandler is a container to hold the map and a reference tracker for how many
    12  // states shared this.
    13  type ValidatorMapHandler struct {
    14  	valIdxMap map[[48]byte]types.ValidatorIndex
    15  	mapRef    *Reference
    16  	*sync.RWMutex
    17  }
    18  
    19  // NewValMapHandler returns a new validator map handler.
    20  func NewValMapHandler(vals []*ethpb.Validator) *ValidatorMapHandler {
    21  	return &ValidatorMapHandler{
    22  		valIdxMap: coreutils.ValidatorIndexMap(vals),
    23  		mapRef:    &Reference{refs: 1},
    24  		RWMutex:   new(sync.RWMutex),
    25  	}
    26  }
    27  
    28  // AddRef copies the whole map and returns a map handler with the copied map.
    29  func (v *ValidatorMapHandler) AddRef() {
    30  	v.mapRef.AddRef()
    31  }
    32  
    33  // IsNil returns true if the underlying validator index map is nil.
    34  func (v *ValidatorMapHandler) IsNil() bool {
    35  	return v.mapRef == nil || v.valIdxMap == nil
    36  }
    37  
    38  // Copy the whole map and returns a map handler with the copied map.
    39  func (v *ValidatorMapHandler) Copy() *ValidatorMapHandler {
    40  	if v == nil || v.valIdxMap == nil {
    41  		return &ValidatorMapHandler{valIdxMap: map[[48]byte]types.ValidatorIndex{}, mapRef: new(Reference), RWMutex: new(sync.RWMutex)}
    42  	}
    43  	v.RLock()
    44  	defer v.RUnlock()
    45  	m := make(map[[48]byte]types.ValidatorIndex, len(v.valIdxMap))
    46  	for k, v := range v.valIdxMap {
    47  		m[k] = v
    48  	}
    49  	return &ValidatorMapHandler{
    50  		valIdxMap: m,
    51  		mapRef:    &Reference{refs: 1},
    52  		RWMutex:   new(sync.RWMutex),
    53  	}
    54  }
    55  
    56  // Get the validator index using the corresponding public key.
    57  func (v *ValidatorMapHandler) Get(key [48]byte) (types.ValidatorIndex, bool) {
    58  	v.RLock()
    59  	defer v.RUnlock()
    60  	idx, ok := v.valIdxMap[key]
    61  	if !ok {
    62  		return 0, false
    63  	}
    64  	return idx, true
    65  }
    66  
    67  // Set the validator index using the corresponding public key.
    68  func (v *ValidatorMapHandler) Set(key [48]byte, index types.ValidatorIndex) {
    69  	v.Lock()
    70  	defer v.Unlock()
    71  	v.valIdxMap[key] = index
    72  }