Implement simple allocator using enif_{alloc,free}

CHANGELOG.md

@@ -11,6 +11,10 @@ versions.
 ## [unreleased]
 
 ### Added
+
+- Add optional support for using Erlang's allocator as Rust's global allocator
+  (#580).
+
 ### Fixed
 ### Changed
 

Cargo.toml

@@ -12,3 +12,8 @@ members = [
   "rustler_tests/native/rustler_compile_tests",
   "rustler_benchmarks/native/benchmark",
 ]
+default-members = [
+  "rustler",
+  "rustler_codegen",
+  "rustler_sys",
+]

rustler/Cargo.toml

@@ -14,6 +14,7 @@ big_integer = ["dep:num-bigint"]
 default = ["derive", "nif_version_2_15"]
 derive = ["rustler_codegen"]
 alternative_nif_init_name = []
+allocator = []
 nif_version_2_14 = ["rustler_sys/nif_version_2_14"]
 nif_version_2_15 = ["nif_version_2_14", "rustler_sys/nif_version_2_15"]
 nif_version_2_16 = ["nif_version_2_15", "rustler_sys/nif_version_2_16"]

rustler/src/alloc.rs

@@ -0,0 +1,55 @@
+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;
+
+/// Allocator implementation that forwards all allocation calls to Erlang's allocator. Allows the
+/// memory usage to be tracked by the BEAM.
+pub struct EnifAllocator;
+
+unsafe impl GlobalAlloc for EnifAllocator {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        if layout.align() > MAX_ALIGN {
+            // Overallocate and store the original pointer in memory immediately before the aligned
+            // section.
+            //
+            // The requested size is chosen such that we can always get an aligned buffer of size
+            // `layout.size()`: Ignoring `SIZEOF_USIZE`, there must always be an aligned pointer in
+            // the interval `[ptr, layout.align())`, so in the worst case, we have to pad with
+            // `layout.align() - 1`. The requirement for an additional `usize` just shifts the
+            // problem without changing the padding requirement.
+            let total_size = SIZEOF_USIZE + layout.size() + layout.align() - 1;
+            let ptr = rustler_sys::enif_alloc(total_size) as *mut u8;
+
+            // Shift the returned pointer to make space for the original pointer
+            let ptr1 = ptr.wrapping_add(SIZEOF_USIZE);
+
+            // Align the result to the requested alignment
+            let aligned_ptr = ptr1.wrapping_add(ptr1.align_offset(layout.align()));
+
+            // Write the original pointer immediately in front of the aligned pointer
+            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) {
+        let ptr = if layout.align() > MAX_ALIGN {
+            // Retrieve the original pointer
+            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);
+    }
+}

rustler/src/lib.rs

@@ -29,6 +29,8 @@ pub mod wrapper;
 #[doc(hidden)]
 pub mod codegen_runtime;
 
+mod alloc;
+
 #[macro_use]
 pub mod types;
 

rustler_tests/native/dynamic_load/Cargo.toml

@@ -9,4 +9,4 @@ path = "src/lib.rs"
 crate-type = ["cdylib"]
 
 [dependencies]
-rustler = { path = "../../../rustler", features = ["big_integer"] }
+rustler = { path = "../../../rustler", features = ["big_integer", "allocator"] }

rustler_tests/native/rustler_bigint_test/Cargo.toml

@@ -9,4 +9,4 @@ path = "src/lib.rs"
 crate-type = ["cdylib"]
 
 [dependencies]
-rustler = { path = "../../../rustler", features = ["big_integer"] }
+rustler = { path = "../../../rustler", features = ["big_integer", "allocator"] }

rustler_tests/native/rustler_compile_tests/Cargo.toml

@@ -10,4 +10,4 @@ path = "src/lib.rs"
 crate-type = ["cdylib"]
 
 [dependencies]
-rustler = { path = "../../../rustler" }
+rustler = { path = "../../../rustler", features = ["allocator"] }

rustler_tests/native/rustler_test/Cargo.toml

@@ -20,4 +20,4 @@ nif_version_2_16 = ["nif_version_2_15", "rustler/nif_version_2_16"]
 nif_version_2_17 = ["nif_version_2_16", "rustler/nif_version_2_17"]
 
 [dependencies]
-rustler = { path = "../../../rustler" }
+rustler = { path = "../../../rustler", features = ["allocator"] }