github.com/MetalBlockchain/metalgo@v1.11.9/utils/heap/map.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package heap
     5  
     6  import (
     7  	"container/heap"
     8  
     9  	"github.com/MetalBlockchain/metalgo/utils"
    10  )
    11  
    12  var _ heap.Interface = (*indexedQueue[int, int])(nil)
    13  
    14  func MapValues[K comparable, V any](m Map[K, V]) []V {
    15  	result := make([]V, 0, m.Len())
    16  	for _, e := range m.queue.entries {
    17  		result = append(result, e.v)
    18  	}
    19  	return result
    20  }
    21  
    22  // NewMap returns a heap without duplicates ordered by its values
    23  func NewMap[K comparable, V any](less func(a, b V) bool) Map[K, V] {
    24  	return Map[K, V]{
    25  		queue: &indexedQueue[K, V]{
    26  			queue: queue[entry[K, V]]{
    27  				less: func(a, b entry[K, V]) bool {
    28  					return less(a.v, b.v)
    29  				},
    30  			},
    31  			index: make(map[K]int),
    32  		},
    33  	}
    34  }
    35  
    36  type Map[K comparable, V any] struct {
    37  	queue *indexedQueue[K, V]
    38  }
    39  
    40  // Push returns the evicted previous value if present
    41  func (m *Map[K, V]) Push(k K, v V) (V, bool) {
    42  	if i, ok := m.queue.index[k]; ok {
    43  		prev := m.queue.entries[i]
    44  		m.queue.entries[i].v = v
    45  		heap.Fix(m.queue, i)
    46  		return prev.v, true
    47  	}
    48  
    49  	heap.Push(m.queue, entry[K, V]{k: k, v: v})
    50  	return utils.Zero[V](), false
    51  }
    52  
    53  func (m *Map[K, V]) Pop() (K, V, bool) {
    54  	if m.Len() == 0 {
    55  		return utils.Zero[K](), utils.Zero[V](), false
    56  	}
    57  
    58  	popped := heap.Pop(m.queue).(entry[K, V])
    59  	return popped.k, popped.v, true
    60  }
    61  
    62  func (m *Map[K, V]) Peek() (K, V, bool) {
    63  	if m.Len() == 0 {
    64  		return utils.Zero[K](), utils.Zero[V](), false
    65  	}
    66  
    67  	entry := m.queue.entries[0]
    68  	return entry.k, entry.v, true
    69  }
    70  
    71  func (m *Map[K, V]) Len() int {
    72  	return m.queue.Len()
    73  }
    74  
    75  func (m *Map[K, V]) Remove(k K) (V, bool) {
    76  	if i, ok := m.queue.index[k]; ok {
    77  		removed := heap.Remove(m.queue, i).(entry[K, V])
    78  		return removed.v, true
    79  	}
    80  	return utils.Zero[V](), false
    81  }
    82  
    83  func (m *Map[K, V]) Contains(k K) bool {
    84  	_, ok := m.queue.index[k]
    85  	return ok
    86  }
    87  
    88  func (m *Map[K, V]) Get(k K) (V, bool) {
    89  	if i, ok := m.queue.index[k]; ok {
    90  		got := m.queue.entries[i]
    91  		return got.v, true
    92  	}
    93  	return utils.Zero[V](), false
    94  }
    95  
    96  func (m *Map[K, V]) Fix(k K) {
    97  	if i, ok := m.queue.index[k]; ok {
    98  		heap.Fix(m.queue, i)
    99  	}
   100  }
   101  
   102  type indexedQueue[K comparable, V any] struct {
   103  	queue[entry[K, V]]
   104  	index map[K]int
   105  }
   106  
   107  func (h *indexedQueue[K, V]) Swap(i, j int) {
   108  	h.entries[i], h.entries[j] = h.entries[j], h.entries[i]
   109  	h.index[h.entries[i].k], h.index[h.entries[j].k] = i, j
   110  }
   111  
   112  func (h *indexedQueue[K, V]) Push(x any) {
   113  	entry := x.(entry[K, V])
   114  	h.entries = append(h.entries, entry)
   115  	h.index[entry.k] = len(h.index)
   116  }
   117  
   118  func (h *indexedQueue[K, V]) Pop() any {
   119  	end := len(h.entries) - 1
   120  
   121  	popped := h.entries[end]
   122  	h.entries[end] = entry[K, V]{}
   123  	h.entries = h.entries[:end]
   124  
   125  	delete(h.index, popped.k)
   126  	return popped
   127  }
   128  
   129  type entry[K any, V any] struct {
   130  	k K
   131  	v V
   132  }