Rust で UTF-8 のコードポイントをちょっと高速に数えてみる

UTF-8のコードポイントはどうやって高速に数えるかという記事を見て、Rust の実装はどうなっているのか気になって調べてみました。

Rust の実装はどうなっているか

Rust で文字数を数える場合は"aα㌁".chars().count()と書きます。(左記は3を返します。"aα㌁".len()はバイト数の6を返します。) では、impl Iterator for Charsを見てみましょう。

#[inline]
fn count(self) -> usize {
    // length in `char` is equal to the number of non-continuation bytes
    let bytes_len = self.iter.len();
    let mut cont_bytes = 0;
    for &byte in self.iter {
        cont_bytes += utf8_is_cont_byte(byte) as usize;
    }
    bytes_len - cont_bytes
}

Ruby のような工夫はしておらず、愚直なコードでした。では Ruby を参考にちょっと高速化してみましょう。

Ruby を参考にした実装

下記の通りです。

trait CharsCount {
    fn chars_count(&self) -> usize;
}

impl CharsCount for str {
    #[inline]
    fn chars_count(&self) -> usize {
        unsafe {
            let mut count = 0;
            let mut p = self.as_ptr();
            let e = self.as_ptr().add(self.len());
            while 8 < e as usize - p as usize {
                let mut u = *(p as *const u8 as *const u64);
                u = u >> 6 | !u >> 7;
                u &= 0x8080_8080_8080_8080 >> 7;
                count += u.count_ones() as usize;
                p = p.add(8);
            }
            while p < e {
                if *p & 0xC0 != 0x80 {
                    count += 1;
                }
                p = p.add(1);
            }
            count
        }
    }
}

これでstr.chars_count()と書くと文字数が返るようになります。u64にcount_ones()というメソッドが実装されていたのには今回始めて気づきました。Rust はいいぞ。
上記コードでベンチマークを取るとstr.chars().count()より 1.5 倍くらい速くなりました。愚直な実装と比較しても、さほどパフォーマンスに差は出ないようです。LLVM による最適化のおかげでしょうか。AVX による実装は今後の課題とします。

メモリアライメントも考慮した実装

メモリアライメントについてコメントがあったので追記します。先の記事の string.c を Rust に愚直に写経しました。

trait CharsCount {
    fn chars_count(&self) -> usize;
}

impl CharsCount for str {
    #[inline]
    fn chars_count(&self) -> usize {
        unsafe {
            let mut count = 0;
            let mut p = self.as_ptr();
            let e = self.as_ptr().add(self.len());
            if mem::size_of::<*const usize>() * 2 < e as usize - p as usize {
                let lowbits = mem::size_of::<*const usize>() - 1;
                let mut s = (!lowbits & p.add(lowbits) as usize) as *const u8;
                let t = (!lowbits & e as usize) as *const u8;
                while p < s {
                    if *p & 0xC0 != 0x80 {
                        count += 1;
                    }
                    p = p.add(1);
                }
                while s < t {
                    let mut u = *(s as *const u8 as *const u64);
                    u = u >> 6 | !u >> 7;
                    u &= 0x8080_8080_8080_8080 >> 7;
                    count += u.count_ones() as usize;
                    s = s.add(8);
                }
                p = s;
            }
            while p < e {
                if *p & 0xC0 != 0x80 {
                    count += 1;
                }
                p = p.add(1);
            }
            count
        }
    }
}

ベンチマークを取るとimpl Iterator for Charsの愚直な実装と比較して 5 倍くらい速くなりました。System Programing Language 面白いですね。メモリにどんな風にデータが格納されているのか、今まで高級言語ばっかり触ってたので何も考えていませんでしたが、こんな風に実感できると素敵です。Rust はいいぞ。