golang.org/x/tools/gopls@v0.15.3/internal/util/lru/lru.go (about)

     1  // Copyright 2023 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // The lru package implements a fixed-size in-memory LRU cache.
     6  package lru
     7  
     8  import (
     9  	"container/heap"
    10  	"fmt"
    11  	"sync"
    12  )
    13  
    14  // A Cache is a fixed-size in-memory LRU cache.
    15  type Cache struct {
    16  	capacity int
    17  
    18  	mu    sync.Mutex
    19  	used  int            // used capacity, in user-specified units
    20  	m     map[any]*entry // k/v lookup
    21  	lru   queue          // min-atime priority queue of *entry
    22  	clock int64          // clock time, incremented whenever the cache is updated
    23  }
    24  
    25  type entry struct {
    26  	key   any
    27  	value any
    28  	size  int   // caller-specified size
    29  	atime int64 // last access / set time
    30  	index int   // index of entry in the heap slice
    31  }
    32  
    33  // New creates a new Cache with the given capacity, which must be positive.
    34  //
    35  // The cache capacity uses arbitrary units, which are specified during the Set
    36  // operation.
    37  func New(capacity int) *Cache {
    38  	if capacity == 0 {
    39  		panic("zero capacity")
    40  	}
    41  
    42  	return &Cache{
    43  		capacity: capacity,
    44  		m:        make(map[any]*entry),
    45  	}
    46  }
    47  
    48  // Get retrieves the value for the specified key, or nil if the key is not
    49  // found.
    50  //
    51  // If the key is found, its access time is updated.
    52  func (c *Cache) Get(key any) any {
    53  	c.mu.Lock()
    54  	defer c.mu.Unlock()
    55  
    56  	c.clock++ // every access updates the clock
    57  
    58  	if e, ok := c.m[key]; ok { // cache hit
    59  		e.atime = c.clock
    60  		heap.Fix(&c.lru, e.index)
    61  		return e.value
    62  	}
    63  
    64  	return nil
    65  }
    66  
    67  // Set stores a value for the specified key, using its given size to update the
    68  // current cache size, evicting old entries as necessary to fit in the cache
    69  // capacity.
    70  //
    71  // Size must be a non-negative value. If size is larger than the cache
    72  // capacity, the value is not stored and the cache is not modified.
    73  func (c *Cache) Set(key, value any, size int) {
    74  	if size < 0 {
    75  		panic(fmt.Sprintf("size must be non-negative, got %d", size))
    76  	}
    77  	if size > c.capacity {
    78  		return // uncacheable
    79  	}
    80  
    81  	c.mu.Lock()
    82  	defer c.mu.Unlock()
    83  
    84  	c.clock++
    85  
    86  	// Remove the existing cache entry for key, if it exists.
    87  	e, ok := c.m[key]
    88  	if ok {
    89  		c.used -= e.size
    90  		heap.Remove(&c.lru, e.index)
    91  		delete(c.m, key)
    92  	}
    93  
    94  	// Evict entries until the new value will fit.
    95  	newUsed := c.used + size
    96  	if newUsed < 0 {
    97  		return // integer overflow; return silently
    98  	}
    99  	c.used = newUsed
   100  	for c.used > c.capacity {
   101  		// evict oldest entry
   102  		e = heap.Pop(&c.lru).(*entry)
   103  		c.used -= e.size
   104  		delete(c.m, e.key)
   105  	}
   106  
   107  	// Store the new value.
   108  	// Opt: e is evicted, so it can be reused to reduce allocation.
   109  	if e == nil {
   110  		e = new(entry)
   111  	}
   112  	e.key = key
   113  	e.value = value
   114  	e.size = size
   115  	e.atime = c.clock
   116  	c.m[e.key] = e
   117  	heap.Push(&c.lru, e)
   118  
   119  	if len(c.m) != len(c.lru) {
   120  		panic("map and LRU are inconsistent")
   121  	}
   122  }
   123  
   124  // -- priority queue boilerplate --
   125  
   126  // queue is a min-atime priority queue of cache entries.
   127  type queue []*entry
   128  
   129  func (q queue) Len() int { return len(q) }
   130  
   131  func (q queue) Less(i, j int) bool { return q[i].atime < q[j].atime }
   132  
   133  func (q queue) Swap(i, j int) {
   134  	q[i], q[j] = q[j], q[i]
   135  	q[i].index = i
   136  	q[j].index = j
   137  }
   138  
   139  func (q *queue) Push(x any) {
   140  	e := x.(*entry)
   141  	e.index = len(*q)
   142  	*q = append(*q, e)
   143  }
   144  
   145  func (q *queue) Pop() any {
   146  	last := len(*q) - 1
   147  	e := (*q)[last]
   148  	(*q)[last] = nil // aid GC
   149  	*q = (*q)[:last]
   150  	return e
   151  }