Implement simple allocator using enif_{alloc,free}

rustler/Cargo.toml

@@ -9,9 +9,10 @@ readme = "../README.md"
 edition = "2021"
 
 [features]
-default = ["derive", "nif_version_2_15"]
+default = ["derive", "nif_version_2_15", "allocator"]
 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,16 @@
+use std::alloc::{GlobalAlloc, Layout};
+
+/// 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 {
+        // TODO: Check the requested alignment.
+        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);
+    }
+}

rustler/src/lib.rs

@@ -29,6 +29,13 @@ pub mod wrapper;
 #[doc(hidden)]
 pub mod codegen_runtime;
 
+mod alloc;
+pub use alloc::EnifAllocator;
+
+#[cfg(feature = "allocator")]
+#[global_allocator]
+static ALLOCATOR: EnifAllocator = EnifAllocator;
+
 pub use lazy_static;
 
 #[macro_use]