diff --git a/Cargo.toml b/Cargo.toml index 7f43e19e..6b494faa 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,3 +12,8 @@ members = [ "rustler_tests/native/rustler_compile_tests", "rustler_benchmarks/native/benchmark", ] +default-members = [ + "rustler", + "rustler_codegen", + "rustler_sys", +] diff --git a/rustler/Cargo.toml b/rustler/Cargo.toml index 15567975..f25bcb6f 100644 --- a/rustler/Cargo.toml +++ b/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"] diff --git a/rustler/src/alloc.rs b/rustler/src/alloc.rs new file mode 100644 index 00000000..799841ac --- /dev/null +++ b/rustler/src/alloc.rs @@ -0,0 +1,46 @@ +use std::alloc::{GlobalAlloc, Layout}; + +const SIZEOF_USIZE: usize = std::mem::size_of::(); +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 + 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 { + // 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); + } +} diff --git a/rustler/src/lib.rs b/rustler/src/lib.rs index fb0e1c91..46a3ffc0 100644 --- a/rustler/src/lib.rs +++ b/rustler/src/lib.rs @@ -29,6 +29,8 @@ pub mod wrapper; #[doc(hidden)] pub mod codegen_runtime; +mod alloc; + #[macro_use] pub mod types; diff --git a/rustler_tests/native/dynamic_load/Cargo.toml b/rustler_tests/native/dynamic_load/Cargo.toml index 4cc7c5d1..126263b3 100644 --- a/rustler_tests/native/dynamic_load/Cargo.toml +++ b/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"] } diff --git a/rustler_tests/native/rustler_bigint_test/Cargo.toml b/rustler_tests/native/rustler_bigint_test/Cargo.toml index 9f64a529..47374e85 100644 --- a/rustler_tests/native/rustler_bigint_test/Cargo.toml +++ b/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"] } diff --git a/rustler_tests/native/rustler_compile_tests/Cargo.toml b/rustler_tests/native/rustler_compile_tests/Cargo.toml index f4b70a6a..9c3ca9fe 100644 --- a/rustler_tests/native/rustler_compile_tests/Cargo.toml +++ b/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"] } diff --git a/rustler_tests/native/rustler_test/Cargo.toml b/rustler_tests/native/rustler_test/Cargo.toml index c2418189..acc27453 100644 --- a/rustler_tests/native/rustler_test/Cargo.toml +++ b/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"] }