Implement proper alignment

rusterlium · May 2, 2024 · dc7ddee · dc7ddee
1 parent 5a9b9e7
commit dc7ddee
Showing 1 changed file with 28 additions and 4 deletions.
diff --git a/rustler/src/alloc.rs b/rustler/src/alloc.rs
@@ -1,5 +1,8 @@
 use std::alloc::{GlobalAlloc, Layout};
 
+const SIZEOF_USIZE: usize = std::mem::size_of::<usize>();
+const MAX_ALIGN: usize = 8;
+
 #[cfg(feature = "allocator")]
 #[global_allocator]
 static ALLOCATOR: EnifAllocator = EnifAllocator;
@@ -10,11 +13,32 @@ pub struct EnifAllocator;
 
 unsafe impl GlobalAlloc for EnifAllocator {
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
-        // TODO: Check enif_alloc's real alignment
-        rustler_sys::enif_alloc(layout.size().max(layout.align())) as *mut u8
+        if layout.align() > MAX_ALIGN {
+            // overallocate
+            let padded = layout.pad_to_align();
+            let total_size = SIZEOF_USIZE + padded.size();
+            let ptr = rustler_sys::enif_alloc(total_size) as *mut u8;
+
+            let ptr1 = ptr.wrapping_add(SIZEOF_USIZE);
+            let aligned_ptr = ptr1.wrapping_add(ptr1.align_offset(layout.align()));
+
+            let header = aligned_ptr.wrapping_sub(SIZEOF_USIZE);
+            *(header as *mut usize) = ptr as usize;
+
+            aligned_ptr
+        } else {
+            rustler_sys::enif_alloc(layout.size()) as *mut u8
+        }
     }
 
-    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
-        rustler_sys::enif_free(ptr as *mut rustler_sys::c_void);
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        let ptr = if layout.align() > MAX_ALIGN {
+            let header = ptr.wrapping_sub(SIZEOF_USIZE);
+            let ptr = *(header as *mut usize);
+            ptr as *mut rustler_sys::c_void
+        } else {
+            ptr as *mut rustler_sys::c_void
+        };
+        rustler_sys::enif_free(ptr);
     }
 }