github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/module/mempool/stdmap/assignments.go (about)

     1  package stdmap
     2  
     3  import (
     4  	chunkmodels "github.com/onflow/flow-go/model/chunks"
     5  	"github.com/onflow/flow-go/model/flow"
     6  )
     7  
     8  // Assignments implements the chunk assignment memory pool.
     9  type Assignments struct {
    10  	*Backend
    11  }
    12  
    13  // NewAssignments creates a new memory pool for Assignments.
    14  func NewAssignments(limit uint) (*Assignments, error) {
    15  	a := &Assignments{
    16  		Backend: NewBackend(WithLimit(limit)),
    17  	}
    18  	return a, nil
    19  }
    20  
    21  // Has checks whether the Assignment with the given hash is currently in
    22  // the memory pool.
    23  func (a *Assignments) Has(assignmentID flow.Identifier) bool {
    24  	return a.Backend.Has(assignmentID)
    25  
    26  }
    27  
    28  // ByID retrieves the chunk assignment from mempool based on provided ID
    29  func (a *Assignments) ByID(assignmentID flow.Identifier) (*chunkmodels.Assignment, bool) {
    30  	entity, exists := a.Backend.ByID(assignmentID)
    31  	if !exists {
    32  		return nil, false
    33  	}
    34  	adp := entity.(*chunkmodels.AssignmentDataPack)
    35  	return adp.Assignment(), true
    36  }
    37  
    38  // Add adds an Assignment to the mempool.
    39  func (a *Assignments) Add(fingerprint flow.Identifier, assignment *chunkmodels.Assignment) bool {
    40  	return a.Backend.Add(chunkmodels.NewAssignmentDataPack(fingerprint, assignment))
    41  }
    42  
    43  // Remove will remove the given Assignment from the memory pool; it will
    44  // return true if the Assignment was known and removed.
    45  func (a *Assignments) Remove(assignmentID flow.Identifier) bool {
    46  	return a.Backend.Remove(assignmentID)
    47  }
    48  
    49  // Size will return the current size of the memory pool.
    50  func (a *Assignments) Size() uint {
    51  	return a.Backend.Size()
    52  }
    53  
    54  // All returns all chunk data packs from the pool.
    55  func (a *Assignments) All() []*chunkmodels.Assignment {
    56  	entities := a.Backend.All()
    57  	assignments := make([]*chunkmodels.Assignment, 0, len(entities))
    58  	for _, entity := range entities {
    59  		assignments = append(assignments, entity.(*chunkmodels.AssignmentDataPack).Assignment())
    60  	}
    61  	return assignments
    62  }