Search code examples
rustalignmentheap-memory

How can I align memory on the heap (in a box) in a convenient way?


I'd like to align my heap memory to a specific alignment boundary. This boundary is known at compilation time. Anyhow, the Box abstraction doesn't allow me to specify a certain alignment. Writing my own Box-abstraction doesn't feel right too, because all of the Rust ecosystem uses Box already. What is a convenient way to achieve alignment for heap allocations?

PS: In my specific case I need page alignments.


Solution

  • If nightly is acceptable, the Allocator api provides a fairly convenient way to do this:

    #![feature(allocator_api)]
    
    use std::alloc::*;
    use std::ptr::NonNull;
    struct AlignedAlloc<const N: usize>;
    
    
    unsafe impl<const N: usize> Allocator for AlignedAlloc<N> {
        fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
            Global.allocate(layout.align_to(N).unwrap())
        }
        unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
            Global.deallocate(ptr, layout.align_to(N).unwrap())
        }
    }
    
    fn main() {
        let val = Box::new_in(3, AlignedAlloc::<4096>);
        let ptr: *const u8 = &*val;
        println!(
            "val:{}, alignment:{}",
            val,
            1 << (ptr as usize).trailing_zeros()
        );
    }
    

    Playground

    If you wanted you could also add support for using other allocators or choosing the value dynamically.

    Edit: Come to think of it, this approach can also be used in stable via the GlobalAlloc trait and the #[global_allocator] attribute, but setting an allocator that way forces all allocations to be aligned to that same alignment, so it would be fine if you needed to ensure a relatively small alignment, but it is probably not a good idea for a page boundary alignment.

    Edit 2: switched from using the System allocator to the Global allocator, because that's bundled in with the allocater_api feature, and it is a more sensible and predictable default.