Zenn
Closed5

Rustのメモリアロケーション

ohkeohke

https://tatsu-zine.com/books/gcbook を読みながら、Rustのメモリアロケーションで以下の点が気になったので調べます。

  • VecやBoxではどうやってヒープ領域を確保しているのか
  • 生のメモリ空間を操作したい場合にどうやって実装すれば良いのか
ohkeohke

まずは "Rust メモリアロケーション" で雑に検索して、こちらの記事にヒット。これが答えでした。
https://qiita.com/moriai/items/4e2ec2d9c3b352394ef3

自分でも1.51のソースコードを追いながら確認。

Vecは以下の要素を持ち、インスタンス化時にAllocator traitを実装した構造体Globalが渡されます。

pub struct Vec<T, A: Allocator = Global> {
    buf: RawVec<T, A>,
    len: usize,
}

RawVec構造体はアロケーションはallocを通して行われます。

pub struct RawVec<T, A: Allocator = Global> {
    ptr: Unique<T>,
    cap: usize,
    alloc: A,
}
// ...
impl<T, A: Allocator> RawVec<T, A> {
    fn allocate_in(capacity: usize, init: AllocInit, alloc: A) -> Self {
        if mem::size_of::<T>() == 0 {
            Self::new_in(alloc)
        } else {
            // ...
            let result = match init {
                AllocInit::Uninitialized => alloc.allocate(layout),
                AllocInit::Zeroed => alloc.allocate_zeroed(layout),
            };
            // ...
            Self {
                ptr: unsafe { Unique::new_unchecked(ptr.cast().as_ptr()) },
                cap: Self::capacity_from_bytes(ptr.len()),
                alloc,
            }
        }
    }
}
// ...
unsafe impl<#[may_dangle] T, A: Allocator> Drop for RawVec<T, A> {
    fn drop(&mut self) {
        if let Some((ptr, layout)) = self.current_memory() {
            unsafe { self.alloc.deallocate(ptr, layout) }
        }
    }
}

Allocator traitは少なくともallocate()とdeallocate()の2つを実装する必要があります。

pub unsafe trait Allocator {
    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError>;
    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout);
}

ではGlobalでは何が実装されているのかというのをallocateに絞って追ってみると、__rust_alloc()というマジックシンボルに行き着きます。いずれかのアロケーション実装が呼びされるようになってます。

  • #[global_allocator]属性が付与された実装が存在する場合は、それが使われる
  • 無ければ、libstdが用いられる (デフォルト)
pub struct Global;

unsafe impl Allocator for Global {
    #[inline]
    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
        self.alloc_impl(layout, false)
    }
    // ...
}

impl Global {
    #[inline]
    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()) }
}

extern "Rust" {
    #[rustc_allocator]
    #[rustc_allocator_nounwind]
    fn __rust_alloc(size: usize, align: usize) -> *mut u8;
    // ...
}

なので、GlobalAllocを実装して#[global_allocator]を付与することで、自前のアロケーションが実現できます。
https://doc.rust-lang.org/std/alloc/index.html

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

デフォルトでは System が用いられてます。
https://doc.rust-lang.org/std/alloc/struct.System.html

なので、デフォルトでは以下の定義がされているように振る舞います。

#[global_allocator]
static A: System = System;

Systemの実装は以下のようになっています。GlobalAllocがOSごとに異なります。

pub struct System;

unsafe impl Allocator for System {
    #[inline]
    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
        self.alloc_impl(layout, false)
    }
    #[inline]
    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
        if layout.size() != 0 {
            unsafe { GlobalAlloc::dealloc(self, ptr.as_ptr(), layout) }
        }
    }
    // ...
}

impl System {
    #[inline]
    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 {
                    GlobalAlloc::alloc_zeroed(self, layout)
                } else {
                    GlobalAlloc::alloc(self, layout)
                };
                let ptr = NonNull::new(raw_ptr).ok_or(AllocError)?;
                Ok(NonNull::slice_from_raw_parts(ptr, size))
            },
        }
    }
    // ...
}
ohkeohke

Unixの実装を見ると、SystemにGlobalAlloc traitを実装していて、libc::malloc()とlibc::free()で確保・解放が行われています。
https://github.com/rust-lang/rust/blob/1.51.0/library/std/src/sys/unix/alloc.rs
(windowsであれば HeapAlloc(), HeapFree()が使われます。)

unsafe impl GlobalAlloc for System {
    #[inline]
    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)
        }
    }
    #[inline]
    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
        libc::free(ptr as *mut libc::c_void)
    }
    // ...
}
ohkeohke

同じ方の https://qiita.com/moriai/items/67761b3c0d83da3b6bb5 を読みながら、もうちょっと。

システムごとに異なる__rust_alloc()や__rust_dealloc()に直接アクセスするために、std::allocでalloc()やdealloc()といった関数が定義されています。
これらの関数を使えば、メモリレイアウトもコントロールできます。(もちろんunsafeですが)

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

Layoutでサイズやアライメントを指定できます。

unsafe {
    let heap1 = alloc(Layout::from_size_align(1024, 256).unwrap());
    let ptr1 = slice::from_raw_parts(heap1, 256);
    println!("address: {:p}, len: {}", ptr1, ptr1.len());
    // address: 0x7ff4e8008800, len: 256

    let heap2 = alloc(Layout::from_size_align(1024, 256).unwrap());
    let ptr2 = slice::from_raw_parts(heap2, 256);
    println!("address: {:p}, len: {}", ptr2, ptr2.len());
    // address: 0x7f8e44809000, len: 256
}
このスクラップは2021/09/20にクローズされました