Rustにおけるバイト列の効率的な扱い方

aptpod Advent Calendar 2023 の12月5日を担当するintdashグループの大久保です。

弊社の製品のうち、エッジ上で動くものはRustを採用しているものが複数あります。例えば、intdashに接続可能なデバイスを簡単に開発するためデバイス開発キットの基盤であるDevice Connector Framework においてもRustを採用しています。

これらのソフトウェアでは大量のバイナリデータを高速に捌く必要があるため、その実装にはRustが適していると考え採用しています。とはいえ、Rustを使えば無条件で効率的な実装になるとは限りません。Rustでバイト列を表現する方法は複数あり、最適化のためにはそれらを使い分ける必要があります。今回は、Rustでのバイト列を扱う方法をまとめてみようと思います。

バイト列を表現する型

まず、一言でバイト列といってもRustでは様々な型を使って表現します。それらの型をここで紹介します。

Vec<u8>

Vec<u8>はヒープ領域上に確保されたバイト列です。mutであれば中身を変更することができ、ライフタイムもついてこないので扱いやすい型です。ただし、clone()を行う場合は新たにヒープ領域への確保と中身のコピーが行われます。中身が同じバイト列をプログラムのあちこちで扱う場合に、cloneを行うたび高コストの操作が行われてしまいます。

&[u8]

別のオブジェクトが所有するバイト列を参照するスライスです。Vec<u8>と異なり、cloneは低コストです。しかしライフタイムが付いてくるため、'staticなライフタイムが必要な場合など使えない場面も多いです。

Cow<[u8]>

Cowは、所有権を持つ型と参照のenumであり、書き換えが必要になったら所有権を持つ型にcloneするという動作をします。バイト列を扱う場合にも使えます。

Bytes

Bytesは、bytesクレートが提供するバイト列を扱うための型です。参照カウンタを用いてメモリを管理するため、効率的なclone()が可能です。

SmallVec

SmallVecは、smallvecクレートが提供する、与えたサイズを超えるデータを保持するときだけヒープ領域を使用するVecに似た型です。その性質上、たいていは短いデータしか扱わないけれども、時々長いデータがやってくる可能性がある場合に効果があります。短いデータはスタックに置くことができ、clone()でヒープ領域の確保が起こらなくなるため高速になります。スタックに収まらないデータが大半ならば他の型を使うべきでしょう。

ArrayVec

ArrayVecは、arrayvecクレートが提供する、ヒープ領域を使わずデータを保持するVecに似た型です。その容量はコンパイル時に決め、容量を超えるデータをpushしようとするとpanicになります。あらかじめ決まった長さ以下の短いデータしか扱わない場合に使えます。

TinyVec

TinyVecは、tinyvecクレートが提供するSmallVecとほぼ同様の型です。このクレートはArrayVecと同等の型も提供します。異なる点は、このクレートはunsafeなコードを含まないためより安全であることが期待できること、その代わり初期化のため配列に入れる型が Defaultを実装している必要があり、パフォーマンスは若干劣ると思われることです。

Bytesの使い方

以前、Rustでバイナリを読み書きするのに必要なクレート3選という記事でbytesクレートについては少し取り上げています。ここでは、バイト列の効率的な表現方法として、Byte型に注目します。

作成とクローン

use bytes::Bytes;

let a = Bytes::copy_from_slice(&[1, 2, 3]);
let b = a.clone();
let c = a.slice(..2);

assert_eq!(&*a, &[1, 2, 3]);
assert_eq!(&*b, &[1, 2, 3]);
assert_eq!(&*c, &[1, 2]);

この例では[1, 2, 3]が中身となるBytesを作り出し、cloneとsliceを呼び出しています。Bytesのcloneは参照カウンタを増やすだけで、Vec<u8>のように新規のヒープ領域確保と中身のコピーが行われません。そのため低コストで呼び出せます。また、slice(..2)はバイト列の一部分を取り出しますが、これは参照カウンタを増やして元のバイト列の一部を指すBytesを返します。配列のスライス&a[..2]と異なり、ライフタイムを気にする必要が無くなります。

BytesMutとの組み合わせ

Bytesの中身は編集できないので、これだけだと使い勝手が悪いです。中身の変更を伴う処理にはBytesMutを使います。

use bytes::{BufMut, BytesMut};

let mut bytes = BytesMut::new();

bytes.extend_from_slice(&[1, 2, 3]);
let a = bytes.split().freeze(); // 先頭を`split`で切り出し、`freeze`で`Bytes`に変換する

bytes.put_u8(4); // `BufMut`トレイトで定義されたメソッドで書き込み
let b = bytes.split().freeze();

assert_eq!(&*a, &[1, 2, 3]);
assert_eq!(&*b, &[4]);
assert_eq!(&*bytes, &[]);

BytesMutは、BufMutトレイトが提供するメソッドで編集することができ、その結果をsplit系のメソッドで切り分け、それらをfreezeでBytes型に変換できます。この切り分けは低コストでできるため、ストリームで流れてきたデータをパースして分割する場合などに効果を発揮できるでしょう。

並行処理

bytesクレートはtokioの下で開発されており、それゆえ並行処理に適するよう実装されています。BytesはSendを実装しているため、スレッド間での共有が容易であり、その共有も参照カウンタによるものなので低コストです。tokioのAPIでもbytesクレートの型は多用されているため、tokioで並行処理を行う場合はこれらの型を使いこなせると便利です。

構造体でのバイト列の持たせ方

構造体にバイト列を持たせる場合、Vec<u8>を使うことが一般的でしょう。

struct MyData {
    id: u32,
    data: Vec<u8>,
}

impl MyData {
    // `data`を使うなにかのメソッド
    fn do_something(&self) {
        let data: &[u8] = self.data.as_ref();
    }
}

let data = &[1, 2, 3];
let my_data = MyData {
    id: 0,
    data: data.to_vec(),
};
my_data.do_something();

しかし、構造体のメソッドが全てVec<u8>が必要とは限らず、参照&[u8]で良い場合も多いでしょう。そのような場合、構造体に定義されたメソッドを使うために、Vec<u8>を作るためのヒープ領域の確保が行われてしまいます。

そのような場合、バイト列をジェネリクスで扱うと良いでしょう。

struct MyData<T> {
    id: u32,
    data: T,
}

impl<T: AsRef<[u8]>> MyData<T> {
    // `data`を使うなにかのメソッド
    fn do_something(&self) {
        let data: &[u8] = self.data.as_ref();
    }
}

let data = &[1, 2, 3];
let my_data = MyData { id: 0, data };
my_data.do_something();

文字列の扱い

バイト列と同様によく扱われるデータとして文字列(UTF-8)があります。これまで紹介したバイト列を表現する型に対応する、文字列の型も存在するのでこちらで紹介します。

バイト列	文字列
Vec<u8>	String
&[u8]	&str
Cow<[u8]>	Cow<str>
Bytes	ByteString
SmallVec<u8>	SmallString
ArrayVec<u8>	ArrayString

その他、SmartStringもここであげておきます。これは、String型と同じサイズで、そのサイズ内に収まる文字列(64bit環境なら23バイトまで)ならヒープ領域を確保せずその場に記憶してくれます。短い文字列が多数で時々長い文字列を格納するときに効果があり、定義されているメソッドやトレイトが通常のStringと同様のものにしてあるので置き換えも容易でしょう。

最後に

以上、Rustにおける様々なバイト列の扱いを紹介しました。バイト列1つでこれだけ選択肢があるのはややこしいですが、うまく使えば安全で高速な実装を行うことができます。これらの知識は実際の製品開発にも活かすようにしています。ここまで読んでくださった方にも本記事の内容がご参考になれば幸いです。