github.com/blend/go-sdk@v1.20220411.3/cache/lru_heap.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package cache
     9  
    10  import "container/heap"
    11  
    12  var (
    13  	_ LRU = (*LRUHeap)(nil)
    14  )
    15  
    16  // NewLRUHeap creates a new, empty, LRU Heap.
    17  func NewLRUHeap() *LRUHeap {
    18  	return new(LRUHeap)
    19  }
    20  
    21  // LRUHeap is a fifo buffer that is backed by a pre-allocated array, instead of allocating
    22  // a whole new node object for each element (which saves GC churn).
    23  // Enqueue can be O(n), Dequeue can be O(1).
    24  type LRUHeap struct {
    25  	Values LRUHeapValues
    26  }
    27  
    28  // Len returns the length of the queue (as it is currently populated).
    29  // Actual memory footprint may be different.
    30  func (lrh *LRUHeap) Len() int {
    31  	return len(lrh.Values)
    32  }
    33  
    34  // Push adds an element to the heap.
    35  func (lrh *LRUHeap) Push(object *Value) {
    36  	heap.Push(&lrh.Values, object)
    37  }
    38  
    39  // Pop removes the first (oldest) element from the heap.
    40  func (lrh *LRUHeap) Pop() *Value {
    41  	if len(lrh.Values) == 0 {
    42  		return nil
    43  	}
    44  	return heap.Pop(&lrh.Values).(*Value)
    45  }
    46  
    47  // Fix updates a value by key.
    48  func (lrh *LRUHeap) Fix(newValue *Value) {
    49  	if len(lrh.Values) == 0 {
    50  		return
    51  	}
    52  	var i int
    53  	for index, value := range lrh.Values {
    54  		if value.Key == newValue.Key {
    55  			i = index
    56  			break
    57  		}
    58  	}
    59  	lrh.Values[i] = newValue
    60  	heap.Fix(&lrh.Values, i)
    61  }
    62  
    63  // Remove removes a value by key.
    64  func (lrh *LRUHeap) Remove(key interface{}) {
    65  	if len(lrh.Values) == 0 {
    66  		return
    67  	}
    68  	var i int
    69  	for index, value := range lrh.Values {
    70  		if value.Key == key {
    71  			i = index
    72  			break
    73  		}
    74  	}
    75  	heap.Remove(&lrh.Values, i)
    76  }
    77  
    78  // Peek returns the oldest value but does not dequeue it.
    79  func (lrh *LRUHeap) Peek() *Value {
    80  	if len(lrh.Values) == 0 {
    81  		return nil
    82  	}
    83  	return lrh.Values[0]
    84  }
    85  
    86  // Consume calls the consumer for each element in the buffer, while also dequeueing that entry.
    87  // The consumer should return `true` if it should remove the item and continue processing.
    88  // If `false` is returned, the current item will be left in place.
    89  func (lrh *LRUHeap) Consume(consumer func(value *Value) bool) {
    90  	if len(lrh.Values) == 0 {
    91  		return
    92  	}
    93  
    94  	len := len(lrh.Values)
    95  	for i := 0; i < len; i++ {
    96  		if !consumer(lrh.Peek()) {
    97  			return
    98  		}
    99  		lrh.Pop()
   100  	}
   101  }
   102  
   103  // Reset removes all elements from the heap, leaving an empty heap.
   104  func (lrh *LRUHeap) Reset() {
   105  	lrh.Values = nil
   106  }