Implement simple allocator using enif_{alloc,free}

Cargo.toml

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

rustler/Cargo.toml

@@ -13,6 +13,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,44 @@
+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
+            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) {
+        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);
+    }
+}

rustler/src/lib.rs

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

rustler_tests/native/deprecated_macros/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/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

@@ -11,4 +11,4 @@ crate-type = ["cdylib"]
 
 [dependencies]
 lazy_static = "1.4"
-rustler = { path = "../../../rustler" }
+rustler = { path = "../../../rustler", features = ["allocator"] }

rustler_tests/native/rustler_test/Cargo.toml

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