GoでLRU Cacheを実装する

LRU Cacheは何かをキャッシュする際によく使うデータ構造の一つだと思う。よく使う一方でその実装はやったことがなかったので、今回Goで実装してみたよ、という話。

LRUCacheとは？

Least Recently Used Cache のこと。一定のキャパシティを持つもので、キャパシティを超える場合使用された時間がもっとも古い要素から削除される。詳しくはWikipediaを見てもらうと良いと思う。

Goでの実装とその解説

まずは単純に実装してみた。テストも含めると上記の3つのファイルから構成されている。lru_cache.goには以下のようなGetとPutメソッドがinterfaceとして定義されており、その実装がdefault.goにある。

type LRUCache interface {
    // Get returns value for `key`. Returns -1 if not found
    Get(key int) int
    // Put sets value with key
    Put(key, value int)
}

使用例

example_test.go で使用例が書かれており、これを

   cache := lru_cache.NewDefault(2)
    cache.Put(1, 1)
    cache.Put(2, 2)

    fmt.Println(cache.Get(1)) // References key `1`

    cache.Put(3, 3) // This operation evicts key `2`
    // Output:
    // 1
    // -1
    fmt.Println(cache.Get(2))

上のコードにあるように、cache.Put(3, 3) によって一番参照が古いkey 2 がcacheから削除される。結果として、cache.Get(2)は-1を返す。これがLRU Cacheの基本的な動作である。

キャッシュのアイテム

キャッシュするアイテムはmapを使って管理する。この時、「いつ参照されたのか」を判断するために、mapのvalueには以下のようなitem構造体を保存する。

type item struct {
    value int
    age   int
}

なお、item.ageはそのアイテムがGetで参照された時またはPutで追加された時のageとなる。そのため、LRU Cacheの実装をしている defaultLRUCache ではGet/Putが呼ばれるたびにcurrentAgeというフィールドをインクリメントする。

type defaultLRUCache struct {
    capacity   int
    values     map[int]*item
    currentAge int
    mutex      *sync.Mutex
}

Get

Getの実装はシンプルで、以下を行うだけである

keyが存在しない場合: -1を返す
keyが存在する場合: 参照されたitemのageを更新し、currentAgeをインクリメントしてreturnする

func (c *defaultLRUCache) Get(key int) int {
    i, ok := c.values[key]
    if !ok {
        return -1
    }
    c.mutex.Lock()
    i.age = c.currentAge
    // `Get` also increment current age
    c.currentAge++
    c.mutex.Unlock()
    return i.value
}

Put

Getに比べるとPutはやや複雑である。

すでにitemがある場合: mapに存在するitemのageを更新する
itemがない場合
- LRU Cacheのcapacityを超えている場合: mapの中で一番ageが小さいitemを見つけて削除する
- 上記に加えて、mapにitemをセットする

という操作を行う。またそれぞれの操作で currentAgeをインクリメントする。

func (c *defaultLRUCache) Put(key int, value int) {
    c.mutex.Lock()
    defer c.mutex.Unlock()

    if i, ok := c.values[key]; ok {
        // If the key exists, update its value and increment its age for this key
        i.value = value
        i.age = c.currentAge
        c.currentAge++
    } else {
        if len(c.values) >= c.capacity {
            // Search key with least age when over capacity before setting key and value
            leastAge := math.MaxInt32
            leastAgeKey := 0
            for key, item := range c.values {
                if item.age < leastAge {
                    leastAge = item.age
                    leastAgeKey = key
                }
            }
            if leastAgeKey != 0 {
                // Evict least age key from cache
                delete(c.values, leastAgeKey)
            }
        }
        // Set key and value to cache
        c.values[key] = &item{
            value: value,
            age:   c.currentAge,
        }
        c.currentAge++
    }
}