RustのAllocatorを自作してみる

C++やRustのコレクションのコードを読んでいるときに、よく `Allocator` というものを目にしていたのですが、これは何なんだろうって思っていたので、調べつつsampleのアロケータを書いて動かしてみました。


# アロケーターとは

RustやC++などには、VectorやHashMap(c++ならmapなど)のコレクションが存在します。
使い方としてはこんな感じ

```rust
let mut vec = Vec::new();
for i in (0..100) {
    vec.push(i * i);
}
```

C言語だと、メモリをプログラマが確保/開放する必要あるわけですが、コレクションでは自分でメモリを確保したりはしません。
とはいえ、プログラムなのでメモリを管理する必要があります。

そこでコレクションのメモリを管理してくれるのがアロケータ（という理解）です。

Vectorに限定して話を進めると、Vectorは内部的には先頭へのポインタとvectorのサイズ,vectorのキャパシティを持っています。
キャパシティというのは、要素が追加されるたびにメモリのアロケーションをすると効率が悪いため、まとめて確保している要素数のことです。

要素が追加される場合は以下のような流れになります

```mermaid
flowchart
A[Push] --> B{キャパシティが\n空いている}
B --> |Yes| C[サイズを増やす]
C --> End
B --> |No| D["キャパシティを増やす\n(メモリ確保)"]
D --> E["新しく確保した領域に\n値をコピーする"]
E --> F["古い領域を開放する"]
F --> End
```

そして、Vectorのスコープが終わるときに、自動的にVectorが保持しているメモリが開放されます。


つまりアロケータは最低限の機能として指定サイズのメモリ確保と、その領域の開放が必要になります。
それを定義したTraitが `GlobalAlloc` です。

```rust
pub unsafe trait GlobalAlloc {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8;
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout);

    // allocとdeallocを使ったデフォルト実装がある
    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8;
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8;
}
```


# サンプルのAllocatorを実装してみる

https://doc.rust-lang.org/std/alloc/struct.System.html

公式ドキュメントに記載されている Counter アロケータを作ってみます。


```rust
struct Counter;

static ALLOCATED: AtomicUsize = AtomicUsize::new(0);

unsafe impl GlobalAlloc for Counter {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let ret = System.alloc(layout);
        if !ret.is_null() {
            ALLOCATED.fetch_add(layout.size(), Ordering::Relaxed);
        }
        ret
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        System.dealloc(ptr, layout);
        ALLOCATED.fetch_sub(layout.size(), Ordering::Relaxed);
    }
}

#[global_allocator]
static A: Counter = Counter;
```

`#[global_allocator]` で、グローバルアロケータに登録しています。

このCounterアロケータは、`alloc`や`dealloc`が呼ばれたときに、要求されたサイズを、`ALLOCATED`変数に加算/減算して、現在のメモリ確保量を記録しています。

どれだけメモリが確保されているのかを表示してみます。

```rust
fn reset_counter() {
    ALLOCATED.store(0, Ordering::Relaxed);
}

fn print_alloc() {
    println!("allocated: {}", ALLOCATED.load(Ordering::Relaxed),);
}

fn main() {
    reset_counter();
    let vec = (1..1000).map(|i| i).collect::<Vec<_>>();
    print_alloc();

    reset_counter();
    let mut vec = Vec::new();
    for i in 1..1000 {
        vec.push(i);
    }
    print_alloc();

    reset_counter();
    let mut vec = Vec::with_capacity(1000);
    for i in 1..1000 {
        vec.push(i);
    }
    print_alloc();
}

```

```bash
$ cargo run
allocated: 3996
allocated: 4096
allocated: 4000
```

1. mapで生成する
1. capacityを指定せずに順次pushする
1. capacityを指定して順次pushする

若干の違いが現れました。

Rust Playgroundを貼っておきます。

https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=9a09653379d211a986ce8a8491ea2de5

## ちょっと機能を足してみる

これだけだと味気ないので、ちょっと機能を足してみます。
合計アロケートサイズだけじゃなく、アロケート回数もカウントするようにしてみます。

```diff
diff --git a/src/main.rs b/src/main.rs
index 4528bd8..c731b65 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -73,12 +73,14 @@ unsafe impl Allocator for FixedSizeAllocator {
 struct Counter;

 static ALLOCATED: AtomicUsize = AtomicUsize::new(0);
+static ALLOCATE_COUNT: AtomicUsize = AtomicUsize::new(0);

 unsafe impl GlobalAlloc for Counter {
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
         let ret = System.alloc(layout);
         if !ret.is_null() {
             ALLOCATED.fetch_add(layout.size(), Ordering::Relaxed);
+            ALLOCATE_COUNT.fetch_add(1, Ordering::Relaxed);
         }
         ret
     }
@@ -94,10 +96,15 @@ static A: Counter = Counter;

 fn reset_counter() {
     ALLOCATED.store(0, Ordering::Relaxed);
+    ALLOCATE_COUNT.store(0, Ordering::Relaxed);
 }

 fn print_alloc() {
-    println!("allocated: {}", ALLOCATED.load(Ordering::Relaxed),);
+    println!(
+        "allocate count: {}, allocated: {}",
+        ALLOCATE_COUNT.load(Ordering::Relaxed),
+        ALLOCATED.load(Ordering::Relaxed)
+    );
 }

 fn main() {
```

ALLOCATE_COUNTの変数を追加し、allocのタイミングで1足すようにしてみました。

ついでに要素数を 100000 に増やして実行してみました。

https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=09d1c241bba58b008054ac76c01415db

```bash
$ cargo run
allocate count: 1, allocated: 799992
allocate count: 16, allocated: 1048576
allocate count: 1, allocated: 800000
```

`map`で`collect`するパターンと、最初に`with_capacity`で確保しておくパターンでは1回しかallocは呼ばれていません。
確保されたメモリ量も明らかにまとめて取得しておくほうが少ないです。
なのでサイズが事前にわかっている場合は `with_capacity` で確保しておくと良さそうです。


このようにRustではグローバルアロケータの挙動を自分で作ることができるため、例えばOSが無いような環境でプログラムが使えるメモリ領域が固定アドレスで決まっていてもそこにコレクションを使用することができます。


# まとめ

アロケータとは、メモリ管理を抽象化してくれるものでした。
コレクションが内部でメモリを要求した際に `alloc`が呼ばれ、メモリを開放する際に`dealloc`が呼ばれることがわかりました。

次回はコレクションごとにアロケータを差し替えてみようと思います。