s_cherry_e/sr/internal/extend/lru/simple_lru.go

package lru

import (
	"container/list"
	"fmt"
)

// EvictCallback is used to get a callback when a cache entry is evicted
type EvictCallback func(key any, value any)

// SimpleLRU implements a non-thread safe fixed size LRU cache
type SimpleLRU struct {
	size      int
	evictList *list.List
	items     map[any]*list.Element
	onEvict   EvictCallback
}

// entry is used to hold a value in the evictList
type entry struct {
	key   any
	value any
}

func New(size int) LRU {
	return NewSimple(size, nil)
}

// NewSimple constructs an LRU of the given size
func NewSimple(size int, onEvict EvictCallback) LRU {
	if size < 1 {
		panic(fmt.Errorf("must provide a positive size. size = %d", size))
	}

	c := &SimpleLRU{
		size:      size,
		evictList: list.New(),
		items:     make(map[any]*list.Element),
		onEvict:   onEvict,
	}
	return c
}

// Purge is used to completely clear the cache.
func (c *SimpleLRU) Purge() {
	for k, v := range c.items {
		if c.onEvict != nil {
			c.onEvict(k, v.Value.(*entry).value)
		}
		delete(c.items, k)
	}
	c.evictList.Init()
}

// Add adds a value to the cache.  Returns true if an eviction occurred.
func (c *SimpleLRU) Add(key, value any) (evicted bool) {
	// Check for existing item
	if ent, ok := c.items[key]; ok {
		c.evictList.MoveToFront(ent)
		ent.Value.(*entry).value = value
		return false
	}

	// NewActor new item
	ent := &entry{key, value}
	frontEntry := c.evictList.PushFront(ent)
	c.items[key] = frontEntry

	evict := c.evictList.Len() > c.size
	// Verify size not exceeded
	if evict {
		c.removeOldest()
	}
	return evict
}

// Get looks up a key's value from the cache.
func (c *SimpleLRU) Get(key any) (value any, ok bool) {
	if ent, ok := c.items[key]; ok {
		c.evictList.MoveToFront(ent)
		if ent.Value.(*entry) == nil {
			return nil, false
		}
		return ent.Value.(*entry).value, true
	}
	return
}

// Contains checks if a key is in the cache, without updating the recent-ness
// or deleting it for being stale.
func (c *SimpleLRU) Contains(key any) (ok bool) {
	_, ok = c.items[key]
	return ok
}

// Peek returns the key value (or undefined if not found) without updating
// the "recently used"-ness of the key.
func (c *SimpleLRU) Peek(key any) (value any, ok bool) {
	var ent *list.Element
	if ent, ok = c.items[key]; ok {
		return ent.Value.(*entry).value, true
	}
	return nil, ok
}

// Remove removes the provided key from the cache, returning if the
// key was contained.
func (c *SimpleLRU) Remove(key any) (present bool) {
	if ent, ok := c.items[key]; ok {
		c.removeElement(ent)
		return true
	}
	return false
}

// RemoveOldest removes the oldest item from the cache.
func (c *SimpleLRU) RemoveOldest() (key any, value any, ok bool) {
	ent := c.evictList.Back()
	if ent != nil {
		c.removeElement(ent)
		kv := ent.Value.(*entry)
		return kv.key, kv.value, true
	}
	return nil, nil, false
}

// GetOldest returns the oldest entry
func (c *SimpleLRU) GetOldest() (key any, value any, ok bool) {
	ent := c.evictList.Back()
	if ent != nil {
		kv := ent.Value.(*entry)
		return kv.key, kv.value, true
	}
	return nil, nil, false
}

// Keys returns a slice of the keys in the cache, from oldest to newest.
func (c *SimpleLRU) Keys() []any {
	keys := make([]any, len(c.items))
	i := 0
	for ent := c.evictList.Back(); ent != nil; ent = ent.Prev() {
		keys[i] = ent.Value.(*entry).key
		i++
	}
	return keys
}

func (c *SimpleLRU) Values() []any {
	values := make([]any, len(c.items))
	i := 0
	for ent := c.evictList.Back(); ent != nil; ent = ent.Prev() {
		values[i] = ent.Value.(*entry).value
		i++
	}
	return values
}

// Len returns the number of items in the cache.
func (c *SimpleLRU) Len() int {
	return c.evictList.Len()
}

// Resize changes the cache size.
func (c *SimpleLRU) Resize(size int) (evicted int) {
	diff := c.Len() - size
	if diff < 0 {
		diff = 0
	}
	for i := 0; i < diff; i++ {
		c.removeOldest()
	}
	c.size = size
	return diff
}

// removeOldest removes the oldest item from the cache.
func (c *SimpleLRU) removeOldest() {
	ent := c.evictList.Back()
	if ent != nil {
		c.removeElement(ent)
	}
}

// removeElement is used to remove a given list element from the cache
func (c *SimpleLRU) removeElement(e *list.Element) {
	c.evictList.Remove(e)
	kv := e.Value.(*entry)
	delete(c.items, kv.key)
	if c.onEvict != nil {
		c.onEvict(kv.key, kv.value)
	}
}