typed-arenaで双方向リンクリストを作ってみた

この記事は Rust Advent Calender 2024シリーズ2 8日目の記事です。

はじめに

Rustが流行ってますね。

Rustを特徴づける機能の一つにshared xor mutableというものがあるのをご存知でしょうか？
これは複数のリソースから参照されている値を可変にすることはできないというもので、このようにすることで"いつどこで値が変更されるかわからない"という状況を防ぐことができます。

shared xor mutableのおかげでコードは見通しがよくなり、『どこで変数が書き換えられているかわからない』というありがちな困り事を防ぐことができます。とてもありがたいものなのです。

が、一つ困ったことがあります。
それは『相互に参照するデータ構造を扱うのが面倒』というものです。

たとえば、以下のような双方リストを作りたいとして

[A] <--> [B] <--> [C]

これを実装すると、内部可変性パターンか、unsafeを使う必要があります。

内部可変性パターンを使った場合、構造体は

use std::rc::Rc;
use std::cell::{Ref, RefMut, RefCell};

pub struct List<T> {
    head: Link<T>,
    tail: Link<T>,
}

type Link<T> = Option<Rc<RefCell<Node<T>>>>;

struct Node<T> {
    elem: T,
    next: Link<T>,
    prev: Link<T>,
}

となり、メソッドは

impl<T> Node<T> {
    fn new(elem: T) -> Rc<RefCell<Self>> {
        Rc::new(RefCell::new(Node {
            elem: elem,
            prev: None,
            next: None,
        }))
    }
}

impl<T> List<T> {
    pub fn new() -> Self {
        List { head: None, tail: None }
    }

    pub fn push_front(&mut self, elem: T) {
        let new_head = Node::new(elem);
        match self.head.take() {
            Some(old_head) => {
                old_head.borrow_mut().prev = Some(new_head.clone());
                new_head.borrow_mut().next = Some(old_head);
                self.head = Some(new_head);
            }
            None => {
                self.tail = Some(new_head.clone());
                self.head = Some(new_head);
            }
        }
    }

    pub fn push_back(&mut self, elem: T) {
        let new_tail = Node::new(elem);
        match self.tail.take() {
            Some(old_tail) => {
                old_tail.borrow_mut().next = Some(new_tail.clone());
                new_tail.borrow_mut().prev = Some(old_tail);
                self.tail = Some(new_tail);
            }
            None => {
                self.head = Some(new_tail.clone());
                self.tail = Some(new_tail);
            }
        }
    }
...(以下略)

と、やや煩雑になってしまいます。

unsafeだとすっきり書けますが、borrow checkerが働かないのでメモリリークのリスクがあります。(コード例は割愛します。詳しくはこちらが参考になると思います。)

では他に良い方法はないのでしょうか？

そこでtyped-arenaというcrateがあります。

これを使うとうまくいきそうです。ですが、調べてみたところ事例がありません。ということで調査してみました。（本当は使ってみたかっただけですが）

typed-arenaってなに？

typed-arenaは割り当てられたオブジェクトを、同じタイミングでdeallocateすることができるcrateです。
https://docs.rs/typed-arena/latest/typed_arena/

これを使うと、このようにallocateできます。（これはサンプルコードそのままですが）
これはa, bどちらのデータもライフタイムがarenaと同じ期間になります。そのため、相互参照できるのです。

use std::cell::Cell;
use typed_arena::Arena;

struct CycleParticipant<'a> {
    other: Cell<Option<&'a CycleParticipant<'a>>>,
}

let arena = Arena::new();

let a = arena.alloc(CycleParticipant { other: Cell::new(None) });
let b = arena.alloc(CycleParticipant { other: Cell::new(None) });

a.other.set(Some(b));
b.other.set(Some(a));

当然ですが、arenaとライフタイムが同じになるため途中でarenaをmoveすることはできません。

fn main() {
    let arena = Arena::new();
    let a = arena.alloc(CycleParticipant {
        other: Cell::new(None),
    });

    let b = arena.alloc(CycleParticipant {
        other: Cell::new(None),
    });
    a.other.set(Some(b));
    drop(arena);
    b.other.set(Some(a));
}

これもダメです。(a, bはそのスコープの間生存するので)

fn main() {
    let arena = Arena::new();
    let a = arena.alloc(CycleParticipant {
        other: Cell::new(None),
    });

    let b = arena.alloc(CycleParticipant {
        other: Cell::new(None),
    });
    a.other.set(Some(b));
    b.other.set(Some(a));
    drop(arena);
}

双方向リンクリストを作ってみる

それでは双方向リンクリストを作ってみましょう。
内部可変性パターン版,unsafe版はLearning Rust With Entirely Too Many Linked Lists
を参考にしてください。双方向リンクリストの説明もそちらに譲ります。

Listの実装

まずListの構造体を定義しましょう。
内部のデータを可変にしたいので、RefCellを使用します。(RefCellは内部データをborrow_mutで可変参照できるスマートポインタです。)

use typed_arena::Arena;
use std::cell::RefCell;

pub struct List<'a, T> {
    arena: &'a Arena<RefCell<Node<'a, T>>>,
    head: Option<&'a RefCell<Node<'a, T>>>,
    tail: Option<&'a RefCell<Node<'a, T>>>,
    length: usize,
}

Nodeは以下の様に定義します。

pub struct Node<'a, T> {
    data: Option<T>,
    next: Option<&'a RefCell<Node<'a, T>>>,
    prev: Option<&'a RefCell<Node<'a, T>>>,
}

初期化の実装

まず、初期化は例のごとくnewを実装します。
List内部のデータの生存期間を呼び出し側に合わせたいので、arenaをnewが受け取るようにします。

impl<'a, T> List<'a, T> {
    pub fn new(arena: &'a Arena<RefCell<Node<'a, T>>>) -> Self {
        List {
            arena,
            head: None,
            tail: None,
            length: 0,
        }
    }
}

以下の様な実装にした場合、arenaの生存期間がListの内部のみになるため、mutな参照を同一スコープで複数回取ることができなくなります。

pub fn new() -> Self {
    List {
        arena: &'a Arena<RefCell<Node<'a, T>>>,
        head: None,
        tail: None,
        length: 0,
    }
}

そのため、複数回pushやpopするとcannot borrow as mutable more than onceエラーが出てしまいます。これを避けるために、arenaを外で初期化して渡すような実装にする必要があります。

メソッドの実装

それではListにメソッドを追加していきましょう。
双方向リンクリストのメソッドといえばなんでしょうか？そう、pushとpopですね。

pushは以下の様になります。push時にarena.allocしてることがわかると思います。
それ以外は普通の双方向リンクリストと同じです。

pub fn push_back(&mut self, data: T) {
    let new_node = self.arena.alloc(RefCell::new(Node {
        data: Some(data),
        next: None,
        prev: None,
    }));
    match self.tail {
        Some(tail) => {
            tail.borrow_mut().next = Some(new_node);
            new_node.borrow_mut().prev = Some(tail);
            self.tail = Some(new_node);
        }
        None => {
            self.head = Some(new_node);
            self.tail = Some(new_node);
        }
    }
    self.length += 1;
}

pub fn push_front(&mut self, data: T) {
    let new_node = self.arena.alloc(RefCell::new(Node {
        data: Some(data),
        next: None,
        prev: None,
    }));
    match self.head {
        Some(head) => {
            head.borrow_mut().prev = Some(new_node);
            new_node.borrow_mut().next = Some(head);
            self.head = Some(new_node);
        }
        None => {
            self.head = Some(new_node);
            self.tail = Some(new_node);
        }
    }
    self.length += 1;
}

次に、popを実装します。dataをOption型にしているので、データ解放時にtake()します。

pub fn pop_back(&mut self) -> Option<T> {
    self.tail.and_then(|tail| {
        let mut tail_node = tail.borrow_mut();
        self.tail = tail_node.prev;
        if let Some(prev) = self.tail {
            prev.borrow_mut().next = None;
        } else {
            self.head = None;
        }
        tail_node.data.take()
    })
}

pub fn pop_front(&mut self) -> Option<T> {
    self.head.and_then(|head| {
        let mut head_node = head.borrow_mut();
        self.head = head_node.next;
        if let Some(next) = self.head {
            next.borrow_mut().prev = None;
        } else {
            self.tail = None;
        }
        self.length -= 1;
        head_node.data.take()
    })
}

これで、基本的な機能を実装できました。

使ってみる

実際に使ってみましょう。呼び出し側の実装はこのような感じになります。

fn main() {
    let arena = Arena::new();  // arenaに紐づくデータの生存期間をmain()に合わせるため、arenaは外に出す
    let mut list: List<i32> = List::new(&arena);

    list.push_back(1);
    list.push_back(2);
    list.push_back(3);
    list.push_front(0);

    // 要素の削除
    while let Some(item) = list.pop_back() {
        println!("Popped: {}", item);
    }
}

実行結果は以下の通りです。

$ cargo run
Popped: 3
Popped: 2
Popped: 1
Popped: 0

全体の実装

改めて全体の実装を載せておきます。

use std::cell::RefCell;
use typed_arena::Arena;

pub struct List<'a, T> {
    arena: &'a Arena<RefCell<Node<'a, T>>>,
    head: Option<&'a RefCell<Node<'a, T>>>,
    tail: Option<&'a RefCell<Node<'a, T>>>,
    length: usize,
}

pub struct Node<'a, T> {
    data: Option<T>,
    next: Option<&'a RefCell<Node<'a, T>>>,
    prev: Option<&'a RefCell<Node<'a, T>>>,
}

impl<'a, T> List<'a, T> {
    pub fn new(arena: &'a Arena<RefCell<Node<'a, T>>>) -> Self {
        List {
            arena,
            head: None,
            tail: None,
            length: 0,
        }
    }

    pub fn push_back(&mut self, data: T) {
        let new_node = self.arena.alloc(RefCell::new(Node {
            data: Some(data),
            next: None,
            prev: None,
        }));
        match self.tail {
            Some(tail) => {
                tail.borrow_mut().next = Some(new_node);
                new_node.borrow_mut().prev = Some(tail);
                self.tail = Some(new_node);
            }
            None => {
                self.head = Some(new_node);
                self.tail = Some(new_node);
            }
        }
        self.length += 1;
    }

    pub fn push_front(&mut self, data: T) {
        let new_node = self.arena.alloc(RefCell::new(Node {
            data: Some(data),
            next: None,
            prev: None,
        }));
        match self.head {
            Some(head) => {
                head.borrow_mut().prev = Some(new_node);
                new_node.borrow_mut().next = Some(head);
                self.head = Some(new_node);
            }
            None => {
                self.head = Some(new_node);
                self.tail = Some(new_node);
            }
        }
        self.length += 1;
    }

    pub fn pop_back(&mut self) -> Option<T> {
        self.tail.and_then(|tail| {
            let mut tail_node = tail.borrow_mut();
            self.tail = tail_node.prev;
            if let Some(prev) = self.tail {
                prev.borrow_mut().next = None;
            } else {
                self.head = None;
            }
            tail_node.data.take()
        })
    }

    pub fn pop_front(&mut self) -> Option<T> {
        self.head.and_then(|head| {
            let mut head_node = head.borrow_mut();
            self.head = head_node.next;
            if let Some(next) = self.head {
                next.borrow_mut().prev = None;
            } else {
                self.tail = None;
            }
            self.length -= 1;
            head_node.data.take()
        })
    }

    pub fn len(&self) -> usize {
        self.length
    }

    pub fn is_empty(&self) -> bool {
        self.length == 0
    }
}

// 使用例
fn main() {
    let arena = Arena::new();
    let mut list: List<i32> = List::new(&arena);

    list.push_back(1);
    list.push_back(2);
    list.push_back(3);
    list.push_front(0);

    // 要素の削除
    while let Some(item) = list.pop_back() {
        println!("Popped: {}", item);
    }
}

まとめ

今回の検証で感じた使い勝手を個人的にまとめておきます。

• typed-arenaを使うと内部可変性パターンよりシンプルかつunsafeより安全に書ける
• 呼び出し側でtyped-arenaを初期化する必要があるので、使い勝手は少し悪いかもしれない。
• 双方向参照を内部に持っているグラフ構造データでは使えそう。

さいごに

peekやIterator,テストコードの実装までやりたかったけど時間足りなかった…
あと生成AIにめっちゃ助けてもらいました。ありがとうClaude3。

参考URL

• https://rust-unofficial.github.io/too-many-lists/fifth-final.html
• https://docs.rs/typed-arena/latest/typed_arena/
• https://doc.rust-jp.rs/book-ja/ch15-05-interior-mutability.html
• https://qnighy.hatenablog.com/entry/2017/04/28/070000