Rust でLRUキャッシュ用のコレクションを自作してみる

これは Rust Advent Calendar 2018 ２日目の記事です。

## はじめに

Rust 書き始めた頃は[所有権](https://doc.rust-lang.org/book/2018-edition/ch04-00-understanding-ownership.html)に苦しめられてきた(個人の感想)。でも、 Rust を書きなれてくると大抵のアルゴリズムやデータ構造は実装できる気になってくる。

そこで(?)、Rust を素直に使うだけだと書き難いことで有名(*要出典)な[双方向連結リスト](https://ja.wikipedia.org/wiki/%E9%80%A3%E7%B5%90%E3%83%AA%E3%82%B9%E3%83%88)と連想配列を使って実装されることが多いLRUキャッシュを作ってみる。

今から作るLRUキャッシュは実用的ではないし(APIが少ない)、テストも不十分だと思う。
ここでの目的はそういうことではなく「`unsafe`ってちょっとなー」とか思っている人向けの、 `unsafe` や Rust での生ポインタの扱いの練習の意味合いが強い。

作ったコードはここにある: https://github.com/hinohi/lru_collections

作るコレクションの名前は `LruMap` とする。

## LruMapのAPIを決める

+ 最大サイズを指定できる
+ エントリーを追加したときに最大サイズを超えていたら、最近最も使われていないエントリーを破棄する
+ APIは [Rust HashMap](https://doc.rust-lang.org/std/collections/struct.HashMap.html#methods) から最低限だけ選択
    + エントリーの追加
    + エントリーの取得
+ **エントリーの追加・取得が O(1) で行えること**
+ `struct` としては `struct LruMap<K, V>`

(今回はエントリーの内容による順序には無頓着なので) `HashMap` のメソッド一覧を眺めて、最低限のメソッドとして以下を実装することにした。

 
### 新しいエントリーを追加するメソッド

[HashMap.insert](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.insert)と同じで、すでに存在する key に対して再度 `insert` した場合には以前に入っていた value を `Some(V)` として返す。

```rust
fn insert(&mut self, k: K, v: V) -> Option<V>;
```

### 値を取得するメソッド

key に対応するエントリーが存在する場合は対応する value を `Some(&V)` として返す。エントリーが存在しないときは `None`。
LRU を実現するために、値にアクセスするだけでもコレクションへの変更が必要なので、 `&mut self` が要求される。

```rust
fn get<Q>(&mut self, k: &Q) -> Option<&V>;
where
    K: Borrow<Q>,
    Q: Hash + Eq + ?Sized,
```

### その他のメソッド

最低限だけ用意する。

```rust
fn len(&self) -> usize;
fn is_empty(&self) -> bool;
```

## データ構造を決める

LRUキャッシュは以下の2要素からなる。

+ エントリーが利用された順番を保持するための双方向連結リスト
+ 値の取得を高速に行うためのリストのノードへのポインタを保持する連想配列

連想配列は Rust 標準ライブラリの `HashMap` を使うことにする。
双方向連結リストは Rust の標準ライブラリに [LinkedList](https://doc.rust-lang.org/std/collections/struct.LinkedList.html) というものがある。しかし `LinkedList` の実装の詳細はプライベートなので、「リストのノードへのポインタから辿ってリストを操作する」ことができない。今回は完全に自作することにする。

次にデータの保持場所を考える。
双方向連結リストに対して実際に行う操作を列挙すると以下の操作が必要だとわかった。

+ リストの先頭にノードを追加
+ リストの末尾のノードを削除して、削除されたノードに対応する key を返却
+ ノードへのポインタを指定して、リスト中の任意の場所のノードをリストの先頭に移動

エントリー `K, V` に対して、 `K` は連想配列と連結リストの双方で保持する必要があり、 `V` はどちらが保持しても良いことがわかる。
今回は `K: Hash + Eq + Clone` とすることで `K` を両方が保持、 `V` を連結リストが保持するように実装した。

実際の定義は以下

```rust
use std::ptr::NonNull;

struct LinkedList<K, V> {
    head: Option<NonNull<Node<K, V>>>,
    tail: Option<NonNull<Node<K, V>>>,
}

struct Node<K, V> {
    next: Option<NonNull<Node<K, V>>>,
    prev: Option<NonNull<Node<K, V>>>,
    key: K,
    value: V,
}

pub struct LruMap<K, V>
where
    K: Hash + Eq + Clone,
{
    max_size: usize,
    map: HashMap<K, NonNull<Node<K, V>>>,
    list: MyLinkedList<K, V>,
}
```

(`K` への参照を連想配列のキーにしたいが間に合わなかった(できるのかな？))

`LinkedList` は標準ライブラリのものと名前が被っているけどまあいいよね。
`Option<NonNull<Node<_>>>` というのは標準ライブラリの `LinkedList` のコードを参考にした。

## 実装していく

### カスタマイズド双方向連結リスト

標準ライブラリのコードが大変に参考になる。
標準の `LinkedList` のコードを眺めつつ今回の `LinkedList` を作成していく(名前が紛らわしい)。

例えば「リストの先頭にノードを追加」は以下のように実装した。

```rust
fn push_front(&mut self, mut node: Box<Node<K, V>>) -> NonNull<Node<K, V>> {
    node.next = self.head;
    node.prev = None;
    unsafe {
        let node = Some(Box::into_raw_non_null(node));
        match self.head {
            None => self.tail = node,
            Some(mut head) => head.as_mut().prev = node,
        }
        self.head = node;
    }
    self.head.unwrap()  // Some(NonNull<_>) を NonNull<_> にしている
}
```

1. `Node<_>` の所有権を受け取り
2. リストの先頭に追加し(この時に `Box<_>` を `NonNull` に変換している)
3. 追加したノードへのポインタを返す

ポインターを作成するだけならセーフだが参照解決を行うのは `unsafe` なので大部分が `unsafe` で囲われている。
当初は `push_front` にノードへのポインタを渡していたが、こうするとなんか色々うまくいかなかったので、`Box<_>` を渡してノードのポインタを返すようになった。返されたノードへのポインタは連想配列の value として登録される。

ちなみに、`push_front` を使う側は、何も `unsafe` なことをしなくていいので、 関数ボディの半分以上が `unsafe` だが `push_front` 自身は `unsafe` ではない。


次に「ノードへのポインタを指定して、リスト中の任意の場所のノードをリストの先頭に移動」のメソッドを示す。

```rust
unsafe fn unlink_and_push_front(&mut self, node: *mut Node<K, V>) {
    let node = node.as_mut().unwrap();
    match node.prev {
        Some(mut prev) => prev.as_mut().next = node.next,
        // this node is the head node
        // nothing to do
        None => return,
    }

    match node.next {
        Some(mut next) => next.as_mut().prev = node.prev,
        // this node is the tail node
        // node.prev is Some<_> in this branch
        None => self.tail = node.prev,
    };

    node.next = self.head;
    node.prev = None;

    let node = Some(node.into());
    self.head.unwrap().as_mut().prev = node;
    self.head = node;
}
```

生ポインタを引数にとり、生ポインタを参照解決してただの `&mut` にしている。
残りの部分は多分 C でも　Rust でも同じような感じになり、特に面白さはない(のに紙面はとる)。

というわけで、 `Rust` でも双方向連結リストを書くことはできた。

### LruMap

`insert` のコードを以下に示す。

```rust
use std::mem;

pub fn insert(&mut self, k: K, v: V) -> Option<V> {
    // TODO use entry API
    if self.map.contains_key(&k) {
        unsafe {
            let mut node = self.map[&k];
            self.list.unlink_and_push_front(node.as_ptr());
            let node = node.as_mut();
            return Some(mem::replace(&mut node.value, v));
        }
    }

    // insert new node
    let node = Node::new(k.clone(), v);
    let ptr = self.list.push_front(Box::new(node));
    self.map.insert(k, ptr);

    // check size
    if self.max_size == 0 || self.map.len() <= self.max_size {
        return None;
    }

    // drop oldest node
    let tail = self.list.pop_back().unwrap();
    self.map.remove(&tail.key);
    None
}
```

まずはエントリーの存在確認を行い、結果で分岐する。

エントリーがすでに存在した場合は既存のエントリーを上書きする。
今回のように 1. 存在確認 2. 対応する値の取得と更新 はまとめて `HashMap.entry.and_modified` などで行うのが普通だが、今回は `v` を `move` するために `move || `　とすると `self` まで `move` されてしまい、後続がうまく書けなかったので諦めた。無念。
おもむろに `unsafe` を開始して `unlink_and_push_front` で連結リストの順序を入れ替え、 `mem::replace` でノードが保持している value を書き換え、古い結果を返している。

エントリーが未登録だったときは

1. `Node<K, V>` を作成し
2. リストの先頭にノードを追加して、ノードへのポインタをもらい
3. ポインタを `HashMap` に登録する

また、エントリーを追加したことによりコレクションのサイズが指定された最大サイズを超えていれば

1. リストの末尾をpopして
2. popされたノードから key を取得し
2. `HashMap`　からもエントリーを削除する

最後に、`get` のコードは以下の通り。

```rust
fn get<Q>(&mut self, k: &Q) -> Option<&V>
where
    K: Borrow<Q>,
    Q: Hash + Eq + ?Sized,
{
    let ptr = match self.map.get(k) {
        None => return None,
        Some(ptr) => ptr,
    };
    unsafe {
        self.list.unlink_and_push_front(ptr.as_ptr());
        Some(&ptr.as_ref().value)
    }
}
```

途中の `ptr` の型は `&NonNull<Node<K, V>>` になっている、ちょっと面白い(?)。
指定されてエントリーが存在していたら、`unsafe` を開始して、連結リストの順序を入れ替え、 value への参照を返している。

(実は関数から参照を返せる条件がよくわかってなかったりするので、 `&ptr.as_ref().value` でいいのかーってなってたりする。誰か教えて)


これで `LruMap` もできた。

## 軽くベンチ

せっかくなので。

```rust
#[bench]
fn bench_get_from_e6(b: &mut Bencher) {
    let mut m = LruMap::new(1_000_000);
    for i in 0..1_000_000 {
        m.insert(i, i);
    }
    b.iter(move || {
        for i in (0..100).rev() {
            m.get(&i);
        }
        for i in 0..100 {
            m.get(&i);
        }
    });
}
```

こんなのを `e2` `e4` `e6` と作った。
(内部のアクセス順序微妙や)

結果

```text
running 3 tests
test bench_get_from_e2 ... bench:       5,289 ns/iter (+/- 1,408)
test bench_get_from_e4 ... bench:       4,494 ns/iter (+/- 229)
test bench_get_from_e6 ... bench:       4,456 ns/iter (+/- 230)
```

エントリーの数に関係なく O(1) で取得はできてる。

## まとめと感想

+ LRUキャッシュ用のコレクション `LruMap` を作った
+ 標準ライブラリの `LinkedList` を参考にすれば自作の双方向連結リストもそんなに難しくない
+ `LruMap` は内部で key を複製して二重に保持しているの無駄なので直したい
+ 作成中はコンパイル通ったのに実行したら `SIGBUS` や `SIGSEGV` が頻発して、`unsafe` 嫌いってなってた
    + `Box` と `NonNull` 行き来していると C よりもデバッグしにくい印象？