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

     1  package backdata
     2  
     3  import (
     4  	"github.com/onflow/flow-go/model/flow"
     5  )
     6  
     7  // MapBackData implements a map-based generic memory BackData backed by a Go map.
     8  // Note that this implementation is NOT thread-safe, and the higher-level Backend is responsible for concurrency management.
     9  type MapBackData struct {
    10  	// NOTE: as a BackData implementation, MapBackData must be non-blocking.
    11  	// Concurrency management is done by overlay Backend.
    12  	entities map[flow.Identifier]flow.Entity
    13  }
    14  
    15  func NewMapBackData() *MapBackData {
    16  	bd := &MapBackData{
    17  		entities: make(map[flow.Identifier]flow.Entity),
    18  	}
    19  	return bd
    20  }
    21  
    22  // Has checks if backdata already contains the entity with the given identifier.
    23  func (b *MapBackData) Has(entityID flow.Identifier) bool {
    24  	_, exists := b.entities[entityID]
    25  	return exists
    26  }
    27  
    28  // Add adds the given entity to the backdata.
    29  func (b *MapBackData) Add(entityID flow.Identifier, entity flow.Entity) bool {
    30  	_, exists := b.entities[entityID]
    31  	if exists {
    32  		return false
    33  	}
    34  	b.entities[entityID] = entity
    35  	return true
    36  }
    37  
    38  // Remove removes the entity with the given identifier.
    39  func (b *MapBackData) Remove(entityID flow.Identifier) (flow.Entity, bool) {
    40  	entity, exists := b.entities[entityID]
    41  	if !exists {
    42  		return nil, false
    43  	}
    44  	delete(b.entities, entityID)
    45  	return entity, true
    46  }
    47  
    48  // Adjust adjusts the entity using the given function if the given identifier can be found.
    49  // Returns a bool which indicates whether the entity was updated as well as the updated entity.
    50  func (b *MapBackData) Adjust(entityID flow.Identifier, f func(flow.Entity) flow.Entity) (flow.Entity, bool) {
    51  	entity, ok := b.entities[entityID]
    52  	if !ok {
    53  		return nil, false
    54  	}
    55  	newentity := f(entity)
    56  	newentityID := newentity.ID()
    57  
    58  	delete(b.entities, entityID)
    59  	b.entities[newentityID] = newentity
    60  	return newentity, true
    61  }
    62  
    63  // AdjustWithInit adjusts the entity using the given function if the given identifier can be found. When the
    64  // entity is not found, it initializes the entity using the given init function and then applies the adjust function.
    65  // Args:
    66  // - entityID: the identifier of the entity to adjust.
    67  // - adjust: the function that adjusts the entity.
    68  // - init: the function that initializes the entity when it is not found.
    69  // Returns:
    70  //   - the adjusted entity.
    71  //
    72  // - a bool which indicates whether the entity was adjusted.
    73  func (b *MapBackData) AdjustWithInit(entityID flow.Identifier, adjust func(flow.Entity) flow.Entity, init func() flow.Entity) (flow.Entity, bool) {
    74  	if b.Has(entityID) {
    75  		return b.Adjust(entityID, adjust)
    76  	}
    77  	b.Add(entityID, init())
    78  	return b.Adjust(entityID, adjust)
    79  }
    80  
    81  // GetWithInit returns the given entity from the backdata. If the entity does not exist, it creates a new entity
    82  // using the factory function and stores it in the backdata.
    83  // Args:
    84  // - entityID: the identifier of the entity to get.
    85  // - init: the function that initializes the entity when it is not found.
    86  // Returns:
    87  //   - the entity.
    88  //   - a bool which indicates whether the entity was found (or created).
    89  func (b *MapBackData) GetWithInit(entityID flow.Identifier, init func() flow.Entity) (flow.Entity, bool) {
    90  	if b.Has(entityID) {
    91  		return b.ByID(entityID)
    92  	}
    93  	b.Add(entityID, init())
    94  	return b.ByID(entityID)
    95  }
    96  
    97  // ByID returns the given entity from the backdata.
    98  func (b *MapBackData) ByID(entityID flow.Identifier) (flow.Entity, bool) {
    99  	entity, exists := b.entities[entityID]
   100  	if !exists {
   101  		return nil, false
   102  	}
   103  	return entity, true
   104  }
   105  
   106  // Size returns the size of the backdata, i.e., total number of stored (entityId, entity)
   107  func (b *MapBackData) Size() uint {
   108  	return uint(len(b.entities))
   109  }
   110  
   111  // All returns all entities stored in the backdata.
   112  func (b *MapBackData) All() map[flow.Identifier]flow.Entity {
   113  	entities := make(map[flow.Identifier]flow.Entity)
   114  	for entityID, entity := range b.entities {
   115  		entities[entityID] = entity
   116  	}
   117  	return entities
   118  }
   119  
   120  // Identifiers returns the list of identifiers of entities stored in the backdata.
   121  func (b *MapBackData) Identifiers() flow.IdentifierList {
   122  	ids := make(flow.IdentifierList, len(b.entities))
   123  	i := 0
   124  	for entityID := range b.entities {
   125  		ids[i] = entityID
   126  		i++
   127  	}
   128  	return ids
   129  }
   130  
   131  // Entities returns the list of entities stored in the backdata.
   132  func (b *MapBackData) Entities() []flow.Entity {
   133  	entities := make([]flow.Entity, len(b.entities))
   134  	i := 0
   135  	for _, entity := range b.entities {
   136  		entities[i] = entity
   137  		i++
   138  	}
   139  	return entities
   140  }
   141  
   142  // Clear removes all entities from the backdata.
   143  func (b *MapBackData) Clear() {
   144  	b.entities = make(map[flow.Identifier]flow.Entity)
   145  }