LRU Cache in Go

cache.go

package cache

type Cache interface {
	Get(key string) interface{}
	Set(key string, value interface{})
}

type CacheNode struct {
	Key   string
	Value interface{}
}

lru.go

package cache

import (
	"errors"
	"sync"
)

// LRUCache structure
type LRUCache struct {
	Cache

	Capacity  int
	Map       map[string]*lruCacheNode
	FirstNode *lruCacheNode
	LastNode  *lruCacheNode

	mu sync.Mutex
}

type lruCacheNode struct {
	CacheNode

	PrevNode *lruCacheNode
	NextNode *lruCacheNode
}

// NewLRUCache returns a new LRUCache
func NewLRUCache(capacity int) (*LRUCache, error) {
	if capacity <= 0 {
		return nil, errors.New("invalid capacity")
	}

	return &LRUCache{
		Capacity:  capacity,
		Map:       make(map[string]*lruCacheNode),
		FirstNode: nil,
		LastNode:  nil,
	}, nil
}

// Get returns the value of the key
func (c *LRUCache) Get(key string) interface{} {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.Map[key] == nil {
		return nil
	}

	node := c.Map[key]
	c.remove(node)
	c.addNodeToFirst(node)

	return node.Value
}

// Set sets a key and value pair to the cache
func (c *LRUCache) Set(key string, value interface{}) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.Map[key] == nil {
		node := &lruCacheNode{
			CacheNode: CacheNode{
				Key:   key,
				Value: value,
			},
			PrevNode: nil,
			NextNode: nil,
		}
		c.addNodeToFirst(node)
	} else {
		node := c.Map[key]
		c.remove(node)
		c.addNodeToFirst(node)
	}

	if len(c.Map) > c.Capacity {
		c.remove(c.LastNode)
	}
}

func (c *LRUCache) addNodeToFirst(node *lruCacheNode) {
	c.Map[node.Key] = node

	node.PrevNode = nil
	node.NextNode = c.FirstNode

	if c.FirstNode != nil {
		c.FirstNode.PrevNode = node
	}
	c.FirstNode = node

	if c.LastNode == nil {
		c.LastNode = node
	}
}

func (c *LRUCache) remove(node *lruCacheNode) {
	if c.Map[node.Key] != nil {
		delete(c.Map, node.Key)
	}

	if node.PrevNode != nil {
		node.PrevNode.NextNode = node.NextNode
	}
	if node.NextNode != nil {
		node.NextNode.PrevNode = node.PrevNode
	}

	if c.FirstNode != nil && c.FirstNode.Key == node.Key {
		c.FirstNode = c.FirstNode.NextNode
	}
	if c.LastNode != nil && c.LastNode.Key == node.Key {
		c.LastNode = c.LastNode.PrevNode
	}
}

lru_test.go

package cache

import (
	"errors"
	"testing"
)

func TestLRUCache(t *testing.T) {
	cache, err := NewLRUCache(2)
	if err != nil {
		t.Error(err)
	}

	cache.Set("1", "a")
	cache.Set("2", "b")
	cache.Set("1", "a")
	cache.Set("3", "c")
	r1 := cache.Get("1")
	if r1 != "a" {
		t.Error(errors.New("wrong"))
	}
	r2 := cache.Get("2")
	if r2 != nil {
		t.Error(errors.New("wrong"))
	}

	if cache.FirstNode != cache.Map["1"] {
		t.Error(errors.New("wrong"))
	}
	if cache.FirstNode.PrevNode != nil {
		t.Error(errors.New("wrong"))
	}
	if cache.FirstNode.NextNode != cache.Map["3"] {
		t.Error(errors.New("wrong"))
	}

	if cache.LastNode != cache.Map["3"] {
		t.Error(errors.New("wrong"))
	}
	if cache.LastNode.PrevNode != cache.Map["1"] {
		t.Error(errors.New("wrong"))
	}
	if cache.LastNode.NextNode != nil {
		t.Error(errors.New("wrong"))
	}
}