github.com/dim4egster/coreth@v0.10.2/plugin/evm/tx_heap.go (about)

     1  // (c) 2020-2021, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package evm
     5  
     6  import (
     7  	"container/heap"
     8  
     9  	"github.com/dim4egster/qmallgo/ids"
    10  )
    11  
    12  // txEntry is used to track the [gasPrice] transactions pay to be included in
    13  // the mempool.
    14  type txEntry struct {
    15  	id       ids.ID
    16  	gasPrice uint64
    17  	tx       *Tx
    18  	index    int
    19  }
    20  
    21  // internalTxHeap is used to track pending atomic transactions by [gasPrice]
    22  type internalTxHeap struct {
    23  	isMinHeap bool
    24  	items     []*txEntry
    25  	lookup    map[ids.ID]*txEntry
    26  }
    27  
    28  func newInternalTxHeap(items int, isMinHeap bool) *internalTxHeap {
    29  	return &internalTxHeap{
    30  		isMinHeap: isMinHeap,
    31  		items:     make([]*txEntry, 0, items),
    32  		lookup:    map[ids.ID]*txEntry{},
    33  	}
    34  }
    35  
    36  func (th internalTxHeap) Len() int { return len(th.items) }
    37  
    38  func (th internalTxHeap) Less(i, j int) bool {
    39  	if th.isMinHeap {
    40  		return th.items[i].gasPrice < th.items[j].gasPrice
    41  	}
    42  	return th.items[i].gasPrice > th.items[j].gasPrice
    43  }
    44  
    45  func (th internalTxHeap) Swap(i, j int) {
    46  	th.items[i], th.items[j] = th.items[j], th.items[i]
    47  	th.items[i].index = i
    48  	th.items[j].index = j
    49  }
    50  
    51  func (th *internalTxHeap) Push(x interface{}) {
    52  	entry := x.(*txEntry)
    53  	if th.Has(entry.id) {
    54  		return
    55  	}
    56  	th.items = append(th.items, entry)
    57  	th.lookup[entry.id] = entry
    58  }
    59  
    60  func (th *internalTxHeap) Pop() interface{} {
    61  	n := len(th.items)
    62  	item := th.items[n-1]
    63  	th.items[n-1] = nil // avoid memory leak
    64  	th.items = th.items[0 : n-1]
    65  	delete(th.lookup, item.id)
    66  	return item
    67  }
    68  
    69  func (th *internalTxHeap) Get(id ids.ID) (*txEntry, bool) {
    70  	entry, ok := th.lookup[id]
    71  	if !ok {
    72  		return nil, false
    73  	}
    74  	return entry, true
    75  }
    76  
    77  func (th *internalTxHeap) Has(id ids.ID) bool {
    78  	_, has := th.Get(id)
    79  	return has
    80  }
    81  
    82  type txHeap struct {
    83  	maxHeap *internalTxHeap
    84  	minHeap *internalTxHeap
    85  }
    86  
    87  func newTxHeap(maxSize int) *txHeap {
    88  	return &txHeap{
    89  		maxHeap: newInternalTxHeap(maxSize, false),
    90  		minHeap: newInternalTxHeap(maxSize, true),
    91  	}
    92  }
    93  
    94  func (th *txHeap) Push(tx *Tx, gasPrice uint64) {
    95  	txID := tx.ID()
    96  	oldLen := th.Len()
    97  	heap.Push(th.maxHeap, &txEntry{
    98  		id:       txID,
    99  		gasPrice: gasPrice,
   100  		tx:       tx,
   101  		index:    oldLen,
   102  	})
   103  	heap.Push(th.minHeap, &txEntry{
   104  		id:       txID,
   105  		gasPrice: gasPrice,
   106  		tx:       tx,
   107  		index:    oldLen,
   108  	})
   109  }
   110  
   111  // Assumes there is non-zero items in [txHeap]
   112  func (th *txHeap) PeekMax() (*Tx, uint64) {
   113  	txEntry := th.maxHeap.items[0]
   114  	return txEntry.tx, txEntry.gasPrice
   115  }
   116  
   117  // Assumes there is non-zero items in [txHeap]
   118  func (th *txHeap) PeekMin() (*Tx, uint64) {
   119  	txEntry := th.minHeap.items[0]
   120  	return txEntry.tx, txEntry.gasPrice
   121  }
   122  
   123  // Assumes there is non-zero items in [txHeap]
   124  func (th *txHeap) PopMax() *Tx {
   125  	return th.Remove(th.maxHeap.items[0].id)
   126  }
   127  
   128  // Assumes there is non-zero items in [txHeap]
   129  func (th *txHeap) PopMin() *Tx {
   130  	return th.Remove(th.minHeap.items[0].id)
   131  }
   132  
   133  func (th *txHeap) Remove(id ids.ID) *Tx {
   134  	maxEntry, ok := th.maxHeap.Get(id)
   135  	if !ok {
   136  		return nil
   137  	}
   138  	heap.Remove(th.maxHeap, maxEntry.index)
   139  
   140  	minEntry, ok := th.minHeap.Get(id)
   141  	if !ok {
   142  		// This should never happen, as that would mean the heaps are out of
   143  		// sync.
   144  		return nil
   145  	}
   146  	return heap.Remove(th.minHeap, minEntry.index).(*txEntry).tx
   147  }
   148  
   149  func (th *txHeap) Len() int {
   150  	return th.maxHeap.Len()
   151  }
   152  
   153  func (th *txHeap) Get(id ids.ID) (*Tx, bool) {
   154  	txEntry, ok := th.maxHeap.Get(id)
   155  	if !ok {
   156  		return nil, false
   157  	}
   158  	return txEntry.tx, true
   159  }
   160  
   161  func (th *txHeap) Has(id ids.ID) bool {
   162  	return th.maxHeap.Has(id)
   163  }