Rustのメモリアロケーションをちょっとだけ掘ってみた

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

# はじめに

[前回の記事](https://qiita.com/moriai/items/67761b3c0d83da3b6bb5)を書きながら、Rust のベクタ `Vec<u8>` で領域を確保すると、小さいサイズでは16B境界、1KB以上では512B境界、64KB以上では4KB境界、.... のようになるのはなぜかなと気になっていたので、調べてみました。ちょうどカレンダーが空いていたので、まとめてみることにします。

低レイヤ人の性で、言語ランタイムが macOS や Linux のどんなライブラリコールやシステムコールを使って実現しているのか、気になってしまいます。今回はそんな感じで掘ってみたお話です。

ちなみに、環境は前回と概ね同じです。

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


# 気になったこと

気になったのは、次のプログラムの挙動です。プログラムを少しシンプルにしました。

```rust:src/bin/vec-u8-simple.rs
pub fn print_info(mem: &[u8]) {
    let addr = mem as *const _ 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 BIGSIZE:usize = 1024*1024;
    let mut sz:usize = 2;
    while sz <= BIGSIZE {
        unsafe {
            let len = sz + 1;
            let mut mem: Vec<u8> = Vec::with_capacity(len);
            mem.set_len(len);
            print_info(&mem);
        }
        sz *= 2;
    }
}
```

これを実行すると、こんな感じになります。

```shell
$ cargo r --bin vec-u8-simple
...
size:          3  addr: 0x7f9084405bb0  bound:      16
size:          5  addr: 0x7f9084405bb0  bound:      16
size:          9  addr: 0x7f9084405bb0  bound:      16
size:         17  addr: 0x7f9084405c20  bound:      32
size:         33  addr: 0x7f9084405c40  bound:      64
size:         65  addr: 0x7f9084405c70  bound:      16
size:        129  addr: 0x7f9084405cc0  bound:      64
size:        257  addr: 0x7f9084405d50  bound:      16
size:        513  addr: 0x7f9084405e60  bound:      32
size:       1025  addr: 0x7f9084808c00  bound:    1024
size:       2049  addr: 0x7f9084809200  bound:     512
size:       4097  addr: 0x7f9084809c00  bound:    1024
size:       8193  addr: 0x7f908480ae00  bound:     512
size:      16385  addr: 0x7f908480d000  bound:    4096
size:      32769  addr: 0x7f9084500000  bound: 1048576
size:      65537  addr: 0x7f908450a000  bound:    8192
size:     131073  addr: 0x7f908451b000  bound:    4096
size:     262145  addr: 0x7f908453c000  bound:   16384
size:     524289  addr: 0x7f908457d000  bound:    4096
size:    1048577  addr: 0x7f90845fe000  bound:    8192
```

（途中、`...`で省略しています）

単純化したので、アロケートするサイズとその境界（アラインメント）は不明瞭になっていますが、誰がこういう境界にメモリ領域を置いたのか、そいつを追求したい、という訳です。


# ソースコードを読んでみる

Rust のソースコード＝ドキュメンテーションを上から順に掘っていきます。ソースコードのトップディレクトリは https://github.com/rust-lang/rust/tree/1.48.0/library です。同じものは、~/.rustup/toolchains/stable-{triple}/lib/rustlib/src/rust/library にありますので、そちらを参照しても良いでしょう。

では、`Vec::with_capacity()` からいきます。

```rust:alloc/src/vec.rs
impl<T> Vec<T> {
    ...
    pub fn with_capacity(capacity: usize) -> Vec<T> {
        Vec { buf: RawVec::with_capacity(capacity), len: 0 }
    }
    ...
}
```
```rust:alloc/src/raw_vec.rs
impl<T> RawVec<T, Global> {
    ...
    pub fn with_capacity(capacity: usize) -> Self {
        Self::with_capacity_in(capacity, Global)
    }
    ...
}

impl<T, A: AllocRef> RawVec<T, A> {
    ...
    pub fn with_capacity_in(capacity: usize, alloc: A) -> Self {
        Self::allocate_in(capacity, AllocInit::Uninitialized, alloc)
    }

    fn allocate_in(capacity: usize, init: AllocInit, alloc: A) -> Self {
        ...
        let result = match init {
            AllocInit::Uninitialized => alloc.alloc(layout),
            AllocInit::Zeroed => alloc.alloc_zeroed(layout),
        };
        ...
    }
    ...
}
```

つまり、`Vec::with_capacity(capacity)` を呼ぶと、`alloc.alloc(layout)` が呼ばれます。ここでは `alloc` は `Global` なので、`Global.alloc(layout)` が呼ばれます。

では、`Global` の周りをみてみましょう。

```rust:alloc/src/alloc.rs
pub struct Global;

unsafe impl AllocRef for Global {
    ...
    fn alloc(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
        self.alloc_impl(layout, false)
    }
    ...
}

impl Global {
    ...
    fn alloc_impl(&self, layout: Layout, zeroed: bool) -> Result<NonNull<[u8]>, AllocError> {
        match layout.size() {
            0 => Ok(NonNull::slice_from_raw_parts(layout.dangling(), 0)),
            size => unsafe {
                let raw_ptr = if zeroed { alloc_zeroed(layout) } else { alloc(layout) };
                let ptr = NonNull::new(raw_ptr).ok_or(AllocError)?;
                Ok(NonNull::slice_from_raw_parts(ptr, size))
            },
        }
    }
    ...
}

pub unsafe fn alloc(layout: Layout) -> *mut u8 {
    unsafe { __rust_alloc(layout.size(), layout.align()) }
}
```

`struct Global` はグローバルメモリアロケータで、`AllocRef` トレイトの実装です。`global_allocator` 属性で登録されたアロケータがある場合はそのアロケータへ、なければ `std` にあるデフォルトのメモリアロケータへ、呼び出しをフォワードするためのものです。

この中で呼び出される `alloc()` が `global_allocator` 属性で登録されたアロケータがあればそれの、なければ `std` の、`GlobalAlloc::alloc` メソッドを呼び出します。この呼び出し分けはマジックシンボル `__rust_alloc` によって実現されており、コンパイラが適切なメソッドが呼び出されるようにマジックシンボルを生成します。

`GlobalAlloc` トレイトの実装はシステムの種類ごとに用意されています。Unix系の実装をみてみましょう。

```rust:std/src/sys/unix/alloc.rs
...
unsafe impl GlobalAlloc for System {
    ...
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        if layout.align() <= MIN_ALIGN && layout.align() <= layout.size() {
            libc::malloc(layout.size()) as *mut u8
        } else {
            #[cfg(target_os = "macos")]
            {
                if layout.align() > (1 << 31) {
                    return ptr::null_mut();
                }
            }
            aligned_malloc(&layout)
        }
    }
    ...
    unsafe fn aligned_malloc(layout: &Layout) -> *mut u8 {
        let mut out = ptr::null_mut();
        // posix_memalign requires that the alignment be a multiple of `sizeof(void*)`.
        // Since these are all powers of 2, we can just use max.
        let align = layout.align().max(crate::mem::size_of::<usize>());
        let ret = libc::posix_memalign(&mut out, align, layout.size());
        if ret != 0 { ptr::null_mut() } else { out as *mut u8 }
    }
    ...
```

アラインメントが「`MIN_ALIGN` 以下、かつ、アロケートするサイズ以下」であれば、`libc::malloc` をそのまま呼び出し、そうでなければ、`aligned_malloc` を呼び出します。`aligned_malloc` は複雑に `cfg` がついていて、ここでは Linux や macOS などの部分だけを抜き出してみましたが、最終的には `libc::posix_memalign` 呼び出しています。なお、`MIN_ALIGN` は x86_64 では 16 となっています。


# `libc::malloc` を直接呼んでみよう

グローバルメモリアロケータは簡単に入れ替えることができるので、`libc::malloc` を直接呼んでみて、動作を確かめてみましょう。

```rust:src/bin/vec-u8-malloc.rs
use std::alloc::{GlobalAlloc, Layout};

struct MyAllocator;

unsafe impl GlobalAlloc for MyAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        libc::malloc(layout.size()) as *mut u8
    }

    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
        libc::free(ptr as *mut libc::c_void);
    }
}

#[global_allocator]
static GLOBAL: MyAllocator = MyAllocator;

fn main () {
    // 中身は vec-u8-simple.rs と同じ
}
```
```shell
$ cargo r --bin vec-u8-malloc
...
size:          3  addr: 0x7fe3a9c059b0  bound:      16
size:          5  addr: 0x7fe3a9c059b0  bound:      16
size:          9  addr: 0x7fe3a9c059b0  bound:      16
size:         17  addr: 0x7fe3a9c05c20  bound:      32
size:         33  addr: 0x7fe3a9c05c40  bound:      64
size:         65  addr: 0x7fe3a9c05c70  bound:      16
size:        129  addr: 0x7fe3a9c05cc0  bound:      64
size:        257  addr: 0x7fe3a9c05d50  bound:      16
size:        513  addr: 0x7fe3a9c05e60  bound:      32
size:       1025  addr: 0x7fe3aa008c00  bound:    1024
size:       2049  addr: 0x7fe3aa009200  bound:     512
size:       4097  addr: 0x7fe3aa009c00  bound:    1024
size:       8193  addr: 0x7fe3aa00ae00  bound:     512
size:      16385  addr: 0x7fe3aa00d000  bound:    4096
size:      32769  addr: 0x7fe3a9d00000  bound: 1048576
size:      65537  addr: 0x7fe3a9d0a000  bound:    8192
size:     131073  addr: 0x7fe3a9d1b000  bound:    4096
size:     262145  addr: 0x7fe3a9d3c000  bound:   16384
size:     524289  addr: 0x7fe3a9d7d000  bound:    4096
size:    1048577  addr: 0x7fe3a9dfe000  bound:    8192
```

`Vec<u8>` の場合とアラインメントは同じになりましたね。

ということで、今回みてきた範囲では、Rust のメモリアロケーションは Unix/Linux 系 OS（POSIX）の malloc(3)/memalign(3) [^1] [^2]を抽象化したものであることがわかりました。これ以上は Rust の世界から出てしまいますので、今回はここでやめておくことにします。malloc については良い解説も色々とありますので。

[^1]: Rust 1.32 からシステム標準の malloc(3) が呼ばれるようになったそうです。以前は、[jemalloc](http://jemalloc.net)が使われていたみたいですね。

[^2]: macOS では Apple の独自実装が使われているみたいです。詳細は[こちら](https://opensource.apple.com/source/libmalloc/libmalloc-283.100.6/)をご覧下さい。

なお、あまり使い途はなさそうですが、上記のコードの `MyAllocator` の実装を下記のように書き換えれば、必ずページサイズに境界を合わせることも簡単にできてしまいます。あくまでも、「できてしまう」という話で、実験用以外には価値がないかな。

```rust:src/bin/vec-u8-aligned.rs
...
unsafe impl GlobalAlloc for MyAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let mut out = ptr::null_mut();
        let ret = libc::posix_memalign(&mut out, page_size::get(), layout.size());
        if ret != 0 { ptr::null_mut() } else { out as *mut u8 }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
        libc::free(ptr as *mut libc::c_void);
    }
}
...
```
```shell
$ cargo r --bin vec-u8-aligned
...
size:          3  addr: 0x7fc0bd80f000  bound:    4096
size:          5  addr: 0x7fc0bd813000  bound:    4096
size:          9  addr: 0x7fc0bd80f000  bound:    4096
size:         17  addr: 0x7fc0bd813000  bound:    4096
size:         33  addr: 0x7fc0bd80f000  bound:    4096
size:         65  addr: 0x7fc0bd813000  bound:    4096
size:        129  addr: 0x7fc0bd80f000  bound:    4096
size:        257  addr: 0x7fc0bd813000  bound:    4096
...
```

`MyAllocator` の実装は、次のようにもできそうです。`layout` をページサイズにアラインし直しているだけなので、基本的にエラーは起きないはずなので、`unwrap()` しても大丈夫でしょう。

```rust:src/bin/vec-u8-aligned2.rs
...
use std::alloc::System;
unsafe impl GlobalAlloc for MyAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let layout = layout.align_to(page_size::get()).unwrap();
        System.alloc(layout)
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        System.dealloc(ptr, layout)
    }
}
...
```

# おわりに

今回は Rust のメモリアロケーションをちょっとだけ掘ってみました。途中、かなり端折りましたが、抽象化を重ねて、柔軟な仕組みを提供していることがわかりました。このネタで掘るとすれば、この抽象化の実行時のコストぐらいでしょうか。記事では省略しましたが、メモリアロケーションのメソッドは基本的に `#[inline]` となっているので、最適化で全て埋め込まれてしまうんだろうなあと思っていますが。