From dc7ddeed1bc75f3e3b610e0875d134e9208e80ed Mon Sep 17 00:00:00 2001
From: Benedikt Reinartz <filmor@gmail.com>
Date: Thu, 2 May 2024 21:37:06 +0200
Subject: [PATCH] Implement proper alignment

---
 rustler/src/alloc.rs | 32 ++++++++++++++++++++++++++++----
 1 file changed, 28 insertions(+), 4 deletions(-)

diff --git a/rustler/src/alloc.rs b/rustler/src/alloc.rs
index a3be8f4e..c62d1593 100644
--- 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);
     }
 }