diff --git a/rustler/Cargo.toml b/rustler/Cargo.toml index e5c53a99..0b95d1b3 100644 --- a/rustler/Cargo.toml +++ b/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"] diff --git a/rustler/src/alloc.rs b/rustler/src/alloc.rs new file mode 100644 index 00000000..1ee48825 --- /dev/null +++ b/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); + } +} diff --git a/rustler/src/lib.rs b/rustler/src/lib.rs index ce700ef6..82b52517 100644 --- a/rustler/src/lib.rs +++ b/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]