cache - cache.go

  // Package cache implements a simple LRU or MRU cache store.
  package cache
  
  import "sync"
  
  type item[T comparable, U any] struct {
  	key        T
  	data       U
  	prev, next *item[T, U]
  }
  
  type mru[T comparable, U any] struct{}
  
  func (mru[T, U]) Remove(first, _ **item[T, U]) *item[T, U] {
  	ret := *first
  
  	(*first).next.prev = nil
  	*first = (*first).next
  	(*first).prev = nil
  
  	return ret
  }
  
  type lru[T comparable, U any] struct {
  	first, last *item[T, U]
  }
  
  func (lru[T, U]) Remove(_, last **item[T, U]) *item[T, U] {
  	ret := *last
  
  	(*last).prev.next = nil
  	*last = (*last).prev
  	(*last).next = nil
  
  	return ret
  }
  
  type ru[T comparable, U any] interface {
  	Remove(first, last **item[T, U]) *item[T, U]
  }
  
  type cache[T comparable, U any, V ru[T, U]] struct {
  	mu          sync.Mutex
  	maxItems    uint64
  	data        map[T]*item[T, U]
  	first, last *item[T, U]
  	remover     V
  }
  
  // Cache is an interface containing the methods used by the LRU and MRU types.
  type Cache[T comparable, U any] interface {
  	Get(T) (U, bool)
  	Set(T, U) bool
  	Remove(T) bool
  	MaxItems() uint64
  	Count() uint64
  	Clear()
  }
  
  // LRU implements a Least-Recently-Used cache, in which a maximum number of
  // items are store, with new items evicting the item that was accessed longest
  // ago.
  type LRU[T comparable, U any] struct {
  	cache[T, U, lru[T, U]]
  }
  
  // NewLRU create a new LRU cache with the specified manximum number of items.
  func NewLRU[T comparable, U any](maxItems uint64) *LRU[T, U] {
  	return &LRU[T, U]{
  		cache: cache[T, U, lru[T, U]]{
  			maxItems: maxItems,
  			data:     make(map[T]*item[T, U], maxItems+1),
  		},
  	}
  }
  
  // MRU implements a Most-Recently-Used cache, in which a maximum number of items
  // are store, with new items evicting the item that was accessed longest ago.
  type MRU[T comparable, U any] struct {
  	cache[T, U, mru[T, U]]
  }
  
  // NewMRU create a new MRU cache with the specified manximum number of items.
  func NewMRU[T comparable, U any](maxItems uint64) *MRU[T, U] {
  	return &MRU[T, U]{
  		cache: cache[T, U, mru[T, U]]{
  			maxItems: maxItems,
  			data:     make(map[T]*item[T, U]),
  		},
  	}
  }
  
  // Get retrieves an item from the cache by its ID.
  //
  // If the item is currently in the store, it returns the stored data and true.
  //
  // If the item does not exist, it returns the zero value for that type and
  // false.
  func (c *cache[T, U, V]) Get(id T) (U, bool) {
  	c.mu.Lock()
  	defer c.mu.Unlock()
  
  	item, ok := c.data[id]
  	if !ok {
  		var data U
  
  		return data, false
  	}
  
  	c.set(item, false)
  
  	return item.data, true
  }
  
  // Set sets the given ID to the given data.
  //
  // Returns false if the ID already exists, true otherwise.
  func (c *cache[T, U, V]) Set(id T, data U) bool {
  	c.mu.Lock()
  	defer c.mu.Unlock()
  
  	it, ok := c.data[id]
  	if ok {
  		it.data = data
  
  		c.set(it, false)
  
  		return false
  	}
  
  	it = &item[T, U]{key: id, data: data}
  	c.data[id] = it
  
  	c.set(it, uint64(len(c.data)) > c.maxItems)
  
  	return true
  }
  
  func (c *cache[T, U, V]) set(it *item[T, U], remove bool) {
  	if c.first == it {
  		return
  	}
  
  	if remove {
  		if toRemove := c.remover.Remove(&c.first, &c.last); toRemove != nil {
  			delete(c.data, toRemove.key)
  		}
  	}
  
  	if c.last == it {
  		c.last = it.prev
  	}
  
  	if it.prev != nil {
  		it.prev.next = it.next
  	}
  
  	if it.next != nil {
  		it.next.prev = it.prev
  	}
  
  	if c.first != nil {
  		c.first.prev = it
  		it.next = c.first
  	} else {
  		c.last = it
  	}
  
  	it.prev = nil
  	c.first = it
  }
  
  // Remove removes the given ID, and its associated data, from the store.
  //
  // Returns false if the ID did not exist, true otherwise.
  func (c *cache[T, U, V]) Remove(id T) bool {
  	c.mu.Lock()
  	defer c.mu.Unlock()
  
  	it, ok := c.data[id]
  	if !ok {
  		return false
  	}
  
  	if it.prev != nil {
  		it.prev.next = it.next
  	}
  
  	if it.next != nil {
  		it.next.prev = it.prev
  	}
  
  	if c.first == it {
  		c.first = it.next
  	}
  
  	if c.last == it {
  		c.last = it.prev
  	}
  
  	delete(c.data, id)
  
  	return true
  }
  
  // MaxItems returns the maximum number of items this cache can store, as
  // specified during creation.
  func (c *cache[T, U, V]) MaxItems() uint64 {
  	return c.maxItems
  }
  
  // Count returns the number of items currently stored in the cache.
  func (c *cache[T, U, V]) Count() uint64 {
  	c.mu.Lock()
  	defer c.mu.Unlock()
  
  	return uint64(len(c.data))
  }
  
  // Clear empties the cache.
  func (c *cache[T, U, V]) Clear() {
  	c.mu.Lock()
  	defer c.mu.Unlock()
  
  	clear(c.data)
  
  	c.first = nil
  	c.last = nil
  }