NonNull安定化記念にInternerを書いてみる

概要: NonNullが安定化され、1.25.0から使えるようになる。そこでNonNullの利用例としてInternerを実装してみた。

NonNull とは

NonNull はRustにある生ポインタ型のひとつです。元々 Unique, Shared という2つの生ポインタ型でしたが、安定化を機に統合・仕様変更が行われ、 NonNull という名前になりました。 (Unique は libcore 内部には残っていますが、安定化の予定はなくなりました。) 予定通りに進めば、 Rust 1.25.0 から使えるようになるはずです。

Rustのポインタ関連型は仕様の微妙な違いを意識して使い分ける必要があります。以下にそれを列挙しました。

	分類	エイリアス	変性	所有	非0	Send	Sync
&T	参照	なし※1	共変	×	○	継承※2	継承
&mut T	参照	なし	非変	×	○	継承	継承
&Cell<T>	参照	あり	非変	×	○	×	×
Box<T>	スマート ポインタ	なし	共変	○	○	継承	継承
*const T	生ポインタ	－	共変	×	×	×	×
*mut T	生ポインタ	－	非変	×	×	×	×
NonNull<T>	生ポインタ	－	共変	×	○	×	×

• ※1 &T は T が UnsafeCell を含まないときだけエイリアスなしと仮定される
• ※2 &T: Send は T: Sync のときに成立

特に参照型はエイリアスに関する仮定があることに注意してください。ここでの「エイリアス: なし」の意味は、その参照の指すメモリには他からの干渉はないとして最適化してよいという意味で、関数の引数が参照だった場合に適用されます。 Box<T> も特別に、引数や戻り値で渡されたときにエイリアスなしと仮定されます。

NonNull<T> はどんなときに便利か

Rustの型システムでうまく表現できない種類の参照を表す必要がある場合に、自分で不変条件を保ちながら unsafe なコードを書くときは、「ライフタイムがついていないけど、真正な参照」というのを表現する必要がある場合があります。この手のことをするときにはNULLかもしれない *const T よりも NonNull<T> を使うほうが適切です。

Rustの型システムでうまく表現できない参照としては、双方向連結リストがよく知られています。ここでは別の例として、internerと呼ばれるデータ構造を書いてみます。これは以下のようなAPIを備えたデータ構造です。

struct Interner { .. }
impl Interner {
    pub fn new() -> Self { .. }
    // 文字列に対して一意な整数を0から順に振って返す。
    pub fn intern(&mut self, s: &str) -> usize { .. }
    // 振られた番号から文字列を復元する。なかったらpanic。
    pub fn get(&self, idx: usize) -> &str { .. }
}

Internerの設計

Internerは正引きと逆引きのためのデータ構造が必要なので、素朴に書くと次のようになります。

use std::collections::HashMap;
pub struct Interner {
    strings: Vec<String>,
    rev_map: HashMap<String, usize>,
}

しかしこの場合同じ文字列を二重に格納するので無駄があります。そこで今回は危険を承知で、 strings 側で所有しているポインタを rev_map で使い回すことにします。概念的には次のようなものを考えます。

use std::collections::HashMap;
pub struct Interner {
    strings: Vec<String>,
    rev_map: HashMap<&'strings str, usize>,
}

ちなみにこのコードは以前の記事で説明した rental crateのマクロに通すとコンパイルはできますが、 strings に対する変更ができないので意図した通りにはなりません。

ライフタイム問題を回避する

そういうわけなので rev_map のキーを生ポインタで保持することを考えます。つまり以下のようにすることを考えます。

use std::collections::HashMap;
pub struct Interner {
    strings: Vec<String>,
    rev_map: HashMap<*const str, usize>,
}

その上で、このデータ構造は以下の不変条件を満たすようにします。

• rev_map のキーに使われているポインタは常に strings の要素である。

非ゼロ制約をつける

HashMap がどのような構造になっているかわからないので、パフォーマンス上の利点があると断言はできませんが、非ゼロであるとわかっているものは非ゼロにしたほうがいいはずです。そこで *const str ではなく NonNull を使います。

use std::collections::HashMap;
use std::ptr::NonNull;
pub struct Interner {
    strings: Vec<String>,
    rev_map: HashMap<NonNull<str>, usize>,
}

変性を考える

今回作っている Interner は型パラメータを持たないので、変性については特に考える必要はありません。一般論としては、 NonNull<T> は共変なので注意が必要ですが、Cellのように特殊なことをやっていなければ共変で問題ないことが多いです。

所有を考える

今回作っている Interner は型パラメータを持たないので、所有については特に考える必要はありません。一般論としては、自作の Drop::drop 中で drop_in_place を呼ぶ場合には、 #[may_dangle] と PhantomData<T> を組み合わせて使う必要があるかもしれません。これらはdrop checkerを宥めるためのものなので、特に困っていなければ保守的に「 #[may_dangle] をつけない」としてもたいてい安全です。

Send, Sync を実装する

生ポインタはデフォルトでは Send でも Sync でもありません。しかし、今回は他のスレッドに転送したり、複数スレッドから共有参照しても特に問題なさそうなので、 Send と Sync を明示的に実装します。

unsafe impl Send for Interner {}
unsafe impl Sync for Interner {}

Eq, Hash, Borrow<str> を実装する

おおよそ適切なデータ構造になってきましたが、 HashMap を動作させるにはキーが Eq + Hash を実装している必要があります。また、 &str で検索するには、キーが Borrow<str> を実装している必要があります。

そこで、 NonNull<T> のラッパーを書きます。このラッパーは大変危険なことをしているので、 Interner 内部以外で使わないよう細心の注意を払う必要があります。

pub struct Interner {
    ...
    rev_map: HashMap<StrPtr, usize>,
}

use std::hash::{Hash, Hasher};
use std::borrow::Borrow;

struct StrPtr(NonNull<str>);

impl Borrow<str> for StrPtr {
    fn borrow(&self) -> &str {
        unsafe { self.0.as_ref() }
    }
}
impl Eq for StrPtr {}
impl PartialEq for StrPtr {
    fn eq(&self, rhs: &Self) -> bool {
        unsafe { str::eq(self.0.as_ref(), rhs.0.as_ref()) }
    }
}
impl Hash for StrPtr {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        unsafe { str::hash(self.0.as_ref(), hasher) }
    }
}

実装を書く

最後に intern の実装を書きます。順番に細心の注意を払います。まずハッシュから検索して、あったらそれを返します。そうでなかったら追加する必要があるので &str を String にします。ポインタの値を覚えておいて、 String は正引き用のVecに突っ込んでしまいます。覚えておいたポインタを NonNull で包んで、逆引き用のHashMapに追加します。 Entry を使うと検索と追加の二度手間を短縮できるのですが、 &str で検索しつつ String で追加するという手順では Entry が使えないのでここでは使っていません。

impl Interner {
    pub fn new() -> Self {
        Interner {
            strings: Vec::new(),
            rev_map: HashMap::new(),
        }
    }
    pub fn intern(&mut self, s: &str) -> usize {
        if let Some(&idx) = self.rev_map.get(s) {
            return idx;
        }
        let idx = self.strings.len();
        let s = s.to_string();
        let s_ptr = unsafe { StrPtr(NonNull::new_unchecked(s.as_str() as *const str as *mut str)) };
        self.strings.push(s);
        self.rev_map.insert(s_ptr, idx);
        return idx;
    }
    pub fn get(&self, idx: usize) -> &str {
        &self.strings[idx]
    }
}

完成

これでできました。確かに動いていることを確認しました。

use std::hash::{Hash, Hasher};
use std::collections::HashMap;
use std::ptr::NonNull;
use std::borrow::Borrow;

pub struct Interner {
    strings: Vec<String>,
    rev_map: HashMap<StrPtr, usize>,
}
unsafe impl Send for Interner {}
unsafe impl Sync for Interner {}

impl Interner {
    pub fn new() -> Self {
        Interner {
            strings: Vec::new(),
            rev_map: HashMap::new(),
        }
    }
    pub fn intern(&mut self, s: &str) -> usize {
        if let Some(&idx) = self.rev_map.get(s) {
            return idx;
        }
        let idx = self.strings.len();
        let s = s.to_string();
        let s_ptr = unsafe { StrPtr(NonNull::new_unchecked(s.as_str() as *const str as *mut str)) };
        self.strings.push(s);
        self.rev_map.insert(s_ptr, idx);
        return idx;
    }
    pub fn get(&self, idx: usize) -> &str {
        &self.strings[idx]
    }
}

struct StrPtr(NonNull<str>);

impl Borrow<str> for StrPtr {
    fn borrow(&self) -> &str {
        unsafe { self.0.as_ref() }
    }
}
impl Eq for StrPtr {}
impl PartialEq for StrPtr {
    fn eq(&self, rhs: &Self) -> bool {
        unsafe { str::eq(self.0.as_ref(), rhs.0.as_ref()) }
    }
}
impl Hash for StrPtr {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        unsafe { str::hash(self.0.as_ref(), hasher) }
    }
}

fn main() {
    let mut i = Interner::new();
    for s in ["a", "c", "a", "b", "b", "c", "d", "a"].into_iter() {
        println!("intern({:?}) = {:?}", s, i.intern(s));
    }
}

最後にもう一押し: drop の順番

おそらく今回は問題ないですが、このようなコードを書くときには drop の順番にも気をつかう必要があります。この場合は rev_map が先にdropされたほうが安心なので、 rev_map に ManuallyDrop を使います。

pub struct Interner {
    ...
    rev_map: ManuallyDrop<HashMap<StrPtr, usize>>,
}

impl Drop for Interner {
    fn drop(&mut self) {
        unsafe {
            ManuallyDrop::drop(&mut self.rev_map);
        }
    }
}

まとめ

NonNull が安定化したので、ライフタイムではちょっと痒いところに手が届かない……というときに生ポインタを使ったコードを書きやすくなりましたが、そういったコードを書くにあたっては細心の注意が必要になります。そこで実際に生ポインタを使ったコードを書きながら、どのような点に注意をしないといけないかを復習してみました。(何か他に見落としている点があったらご指摘ください。)

2018/02/08 追記: NonNullのsinceアトリビュートが1.24から1.25.0に修正されたのを記事に反映した。