Rustでページ境界に合わせたメモリアロケーションをするには

この記事は [Rustその2 Advent Calendar 2020](https://qiita.com/advent-calendar/2020/rust2) 12日目 の記事です。


# はじめに

Rust で低レイヤ[^1]のプログラミングを楽しんでいるのですが、メモリブロックを他のプロセスやデバイスと共有しようと思うと、ページ境界に合ったメモリブロックが必要になります。そこで、Rust でページ境界に合わせた大きなメモリブロックを動的に割り当てるにはどうしたらいいの、という観点で勉強してみました。「Rustで低レイヤ入門」以外の方は、さらっと読み飛ばしていただければ。

[^1]: コンパイラ、仮想化、オペレーティングシステム、組込み用プログラムなどのことです。

今回の環境はこんな感じです。ひさびさに stable です。

+ macOS 10.15.7 (19H15)
+ rustc 1.48.0 (7eac88abb 2020-11-16)
+ cargo 1.48.0 (65cbdd2dc 2020-10-14)

まだ Big Sur には馴染めず、セキュリティパッチが出る限りは Catalina で頑張るつもりです。でも、M1 Mac が届いたら、Big Sur になってしまうなぁ...


# 配列

配列はコンパイル時に大きさが決まり、スタック上に確保されますので、目的には合わないのですが、試してみましょう。

```Rust
fn main () {
    const ONEGB:usize = 1024*1024*1024;
    let mem = [0u8;ONEGB];
    ...
}
```
```Shell
$ cargo run --release
...
thread 'main' has overflowed its stack
fatal runtime error: stack overflow
Abort trap: 6
```

さすがに 1GB は無理でした。なお、"..." は適宜省略の意味です。

スタックサイズは Rust のスレッドごとのスタックサイズの上限（環境変数 `RUST_MIN_STACK`）とプロセス全体のスタックサイズの上限（`ulimit -s`）の２つに縛られていて、今回の環境でのデフォルトは、それぞれ 2MB、8MB のようです。貴重なスタック資源は大切に使いましょう。


# ベクタ

ベクタは大きさは動的で、ヒープ上に確保されます。`with_capacity(size)` で `size` の大きさの領域が確保されますが、`len()` はゼロのままになります。`set_len(size)` とすることで、初期化せずにベクタの長さが `size` になります。"初期化せずに" なので、`unsafe` となります。では、やってみましょう。

```Rust
fn print_info(mem: &[u8]) {
    let addr = (&mem[0] as *const u8) as u64;
    let mut bound: u64 = 1;
    while addr & bound == 0 { bound <<= 1; }
    println!("size: {:>10}  addr: 0x{:>012x}  bound: {:>7}", mem.len(), addr, bound);
}

fn main () {
    const ONEGB:usize = 1024*1024*1024;
    let mut size:usize = 2;
    while size <= ONEGB {
        let mut mem0: Vec<u8> = Vec::with_capacity(size);
        unsafe { mem0.set_len(size); }
        print_info(&mem0);
        let mut mem1: Vec<u8> = Vec::with_capacity(size);
        unsafe { mem1.set_len(size); }
        print_info(&mem1);
        size *= 2;
    }
}
```
```Shell
$ cargo run --release
...
size:          2  addr: 0x7fc528405b90  bound:      16
size:          2  addr: 0x7fc528405c00  bound:    1024
size:          4  addr: 0x7fc528405b90  bound:      16
size:          4  addr: 0x7fc528405c00  bound:    1024
...
size:       1024  addr: 0x7fc528808c00  bound:    1024
size:       1024  addr: 0x7fc528809000  bound:    4096
...
size:      65536  addr: 0x7fc528500000  bound: 1048576
size:      65536  addr: 0x7fc528511000  bound:    4096
...
size: 1073741824  addr: 0x7fc4c8400000  bound: 4194304
size: 1073741824  addr: 0x7fc488400000  bound: 4194304

```

bound はメモリブロックの先頭アドレスが何バイト境界に乗っているかを示しています。

`u8` のベクタの場合、小さいサイズでは16B境界、1KB以上では1KB境界、64KB以上では4KB境界、.... などとなっています。ここでは、具体的なアロケーションの仕組みには深入りしないことにしますが、4KB以上でもページ境界（macOSだと4KB）にならないこともあるのですね。


# ボックス

ボックスはヒープ上に確保されますが、コンパイル時にサイズが決まっている必要がありますし、初期化式をスタック上で評価してからヒープ上に書き込まれます。unstable な box_syntax を使えば、スタックを経由せずにダイレクトにヒープ上に初期化の内容が書き込むことができますが、今回の目的には合いません。

さらっと、次、いきましょう！


# `std::alloc` + `std::slice`

というわけで、やっと本題です。

`std::alloc::alloc()` はアロケーション時のレイアウトを指定することができます。レイアウトは `struct std::alloc::Layout` で指定します。まずは使ってみましょう。

```Rust
use std::alloc::{alloc, Layout};
use std::slice;

fn aligned_alloc(size: usize) -> &'static mut [u8] {
    unsafe {
        let layout = Layout::from_size_align(size, 4096).unwrap();
        let raw_mem = alloc(layout);
        slice::from_raw_parts_mut(raw_mem, size)
    }
}

fn main () {
    const ONEGB:usize = 1024*1024*1024;
    let mut size:usize = 4;
    while size <= ONEGB {
        let mem0 = aligned_alloc(size);
        print_info(&mem0);
        let mem1 = aligned_alloc(size);
        mem1[0] = 0;
        print_info(&mem1);
        size *= 4;
    }
}
```
```Shell
$ cargo run --release --bin std-alloc
...
size:          4  addr: 0x7ffeb5009000  bound:    4096
size:          4  addr: 0x7ffeb500a000  bound:    8192
size:         16  addr: 0x7ffeb500b000  bound:    4096
size:         16  addr: 0x7ffeb500d000  bound:    4096
...
size: 1073741824  addr: 0x000111243000  bound:    4096
size: 1073741824  addr: 0x000151243000  bound:    4096
```

途中省略しましたが、4KB境界に乗っています。drop する手段が提供されず、ライフタイムが static なので、次々と新しい領域がアロケートされています[^2]。

[^2]: 関数ではなく、マクロで実装した方が良かったかもしれません。

# 構造体にしてみる

`std::alloc` と `std::slice` で上手くいくことが分かりましたので、もう少し、使いやすくしてみましょう。Drop トレイトを実装することで、スコープを外れた時に自動的にデアロケートできます。また、Deref, DerefMut トレイトを実装すると、自動的に `u8` のスライスとして扱えるようになりますので、さらに便利です[^3]。デバッグ用に `#[derive(Debug)]` しましょう。こんな感じでいかがでしょうか。

[^3]: あとは、必要に応じて `AsRef`, `AsMut` などを実装すれば良いですね。

```Rust
use std::alloc::{alloc, dealloc, Layout, LayoutErr};
use std::ops::{Deref, DerefMut};
use std::slice;

#[derive(Debug)]
struct AlignedAlloc {
    ptr: *mut u8,
    layout: Layout
}

impl AlignedAlloc {
    fn new(len: usize) -> Result<AlignedAlloc, LayoutErr> {
        let layout = match Layout::from_size_align(len, 4096) {
            Ok(l) => l, Err(e) => return Err(e)
        };
        unsafe {
            let ptr = alloc(layout);
            Ok(AlignedAlloc{ ptr, layout })
        }
    }
}

impl Drop for AlignedAlloc {
    fn drop(&mut self) {
        unsafe { dealloc(self.ptr, self.layout) }
    }
}

impl Deref for AlignedAlloc {
    type Target = [u8];
    fn deref(&self) -> &[u8] {
        unsafe { slice::from_raw_parts(self.ptr, self.layout.size()) }
    }
}

impl DerefMut for AlignedAlloc {
    fn deref_mut(&mut self) -> &mut [u8] {
        unsafe { slice::from_raw_parts_mut(self.ptr, self.layout.size()) }
    }
}

fn main() {
    const ONEGB:usize = 1024*1024*1024;
    let mut size:usize = 4;
    while size <= ONEGB {
        let mem0 = match AlignedAlloc::new(size) {
            Ok(m) => m,
            Err(e) => { panic!("AlignedAlloc fail {:?}", e); }
        };
        print_info(&mem0);
        let mut mem1 = match AlignedAlloc::new(size) {
            Ok(m) => m,
            Err(e) => { panic!("AlignedAlloc fail {:?}", e); }
        };
        mem1[0] = 0;
        print_info(&mem1);
        size *= 4;
    }
}
```
```Shell
$ cargo run --release
...
size:          4  addr: 0x7ffbad009000  bound:    4096
size:          4  addr: 0x7ffbad00a000  bound:    8192
size:         16  addr: 0x7ffbad00a000  bound:    8192
size:         16  addr: 0x7ffbad009000  bound:    4096
...
size: 1073741824  addr: 0x7ffb6cc00000  bound: 4194304
size: 1073741824  addr: 0x7ffb2cc00000  bound: 4194304
```

`while` ループを回るたびに `mem0` と `mem1` がスコープからを外れるので、領域が開放されて、再利用されていることが分かります。

## 少し補足

`struct AlignedAlloc` を書き始めた頃は次のようになっていました。

```rust
struct AlignedAlloc<'a> {
    mem: &'a mut [u8]
}
```

`new()` はこれで実装できるのですが、`drop()` を書き始め、悩みました。`dealloc()` するためには、`alloc()` に使った `Layout` と `*mut u8` のポインタが必要になります。

```rust
struct AlignedAlloc<'a> {
    mem: &'a mut [u8],
    ptr: *mut u8,
    layout: Layout
}
```

ちょっと気持ち悪いです。`[u8]` を裏から見たものが `*u8` と `Layout` なので、単に抽象化の仕方が違うだけです。表現として冗長ですし、使い勝手も良くありません。

```rust
let a = AlignedAlloc::new(4096).unwrap();
a.mem[0] = 0;
...
``` 

というわけで、構造体としては `*u8` と `Layout` を持たせ、参照を解決する時にスライスを作って返せば、コード量はちょっと増えますが、きれいに実装できるし、利用する場合も楽になるという訳です。


# おわりに

ということで、メモリアロケーションについて勉強したことをまとめてみましたが、今回はイントロです。次回は、[あなたの知らない（？）ディープな世界へと続く](https://qiita.com/moriai/items/4e2ec2d9c3b352394ef3)...?

なお、κeenさんが[奥ゆかしいコメント記事](https://qiita.com/blackenedgold/items/823ab427477e37995ee6)を書いてくださったので、そちらもあわせてお読み下さい。（2020年12月14日追記）