Skip to main content

inkwell/
memory_manager.rs

1use llvm_sys::prelude::LLVMBool;
2
3/// A trait for user-defined memory management in MCJIT.
4///
5/// Implementors can override how LLVM's MCJIT engine allocates memory for code
6/// and data sections. This is sometimes needed for:
7/// - custom allocators,
8/// - sandboxed or restricted environments,
9/// - capturing stack map sections (e.g., for garbage collection),
10/// - or other specialized JIT memory management requirements.
11///
12/// # StackMap and GC Integration
13///
14/// By examining the `section_name` argument in [`McjitMemoryManager::allocate_data_section`],
15/// you can detect sections such as `.llvm_stackmaps` (on ELF) or `__llvm_stackmaps`
16/// (on Mach-O). Recording the location of these sections may be useful for
17/// custom garbage collectors. For more information, refer to the [LLVM
18/// StackMaps documentation](https://llvm.org/docs/StackMaps.html#stack-map-section).
19///
20/// Typically, on Darwin (Mach-O), the stack map section name is `__llvm_stackmaps`,
21/// and on Linux (ELF), it is `.llvm_stackmaps`.
22pub trait McjitMemoryManager: std::fmt::Debug {
23    /// Allocates a block of memory for a code section.
24    ///
25    /// # Parameters
26    ///
27    /// * `size` - The size in bytes for the code section.
28    /// * `alignment` - The required alignment in bytes.
29    /// * `section_id` - A numeric ID that LLVM uses to identify this section.
30    /// * `section_name` - A name for this section, if provided by LLVM.
31    ///
32    /// # Returns
33    ///
34    /// Returns a pointer to the allocated memory. Implementors must ensure it is
35    /// at least `size` bytes long and meets `alignment` requirements.
36    fn allocate_code_section(
37        &mut self,
38        size: libc::uintptr_t,
39        alignment: libc::c_uint,
40        section_id: libc::c_uint,
41        section_name: &str,
42    ) -> *mut u8;
43
44    /// Allocates a block of memory for a data section.
45    ///
46    /// # Parameters
47    ///
48    /// * `size` - The size in bytes for the data section.
49    /// * `alignment` - The required alignment in bytes.
50    /// * `section_id` - A numeric ID that LLVM uses to identify this section.
51    /// * `section_name` - A name for this section, if provided by LLVM.
52    /// * `is_read_only` - Whether this data section should be read-only.
53    ///
54    /// # Returns
55    ///
56    /// Returns a pointer to the allocated memory. Implementors must ensure it is
57    /// at least `size` bytes long and meets `alignment` requirements.
58    fn allocate_data_section(
59        &mut self,
60        size: libc::uintptr_t,
61        alignment: libc::c_uint,
62        section_id: libc::c_uint,
63        section_name: &str,
64        is_read_only: bool,
65    ) -> *mut u8;
66
67    /// Finalizes memory permissions for all allocated sections.
68    ///
69    /// This is called once all sections have been allocated. Implementors can set
70    /// permissions such as making code sections executable or data sections
71    /// read-only.
72    ///
73    /// # Errors
74    ///
75    /// If any error occurs (for example, failing to set page permissions),
76    /// return an `Err(String)`. This error is reported back to LLVM as a C string.
77    fn finalize_memory(&mut self) -> Result<(), String>;
78
79    /// Cleans up or deallocates resources before the memory manager is destroyed.
80    ///
81    /// This is called when LLVM has finished using the memory manager. Any
82    /// additional allocations or references should be released here if needed.
83    fn destroy(&mut self);
84}
85
86/// Holds a boxed `McjitMemoryManager` and passes it to LLVM as an opaque pointer.
87///
88/// LLVM calls into the adapter using the extern "C" function pointers defined below.
89#[derive(Debug)]
90pub struct MemoryManagerAdapter {
91    pub memory_manager: Box<dyn McjitMemoryManager>,
92}
93
94// ------ Extern "C" Adapters ------
95
96/// Adapter for `allocate_code_section`.
97///
98/// Called by LLVM with a raw pointer (`opaque`). Casts back to `MemoryManagerAdapter`
99/// and delegates to `allocate_code_section`.
100pub(crate) extern "C" fn allocate_code_section_adapter(
101    opaque: *mut libc::c_void,
102    size: libc::uintptr_t,
103    alignment: libc::c_uint,
104    section_id: libc::c_uint,
105    section_name: *const libc::c_char,
106) -> *mut u8 {
107    let adapter = unsafe { &mut *(opaque as *mut MemoryManagerAdapter) };
108    let sname = unsafe { c_str_to_str(section_name) };
109    adapter
110        .memory_manager
111        .allocate_code_section(size, alignment, section_id, sname)
112}
113
114/// Adapter for `allocate_data_section`.
115///
116/// Note that `LLVMBool` is `0` for false, and `1` for true. We check `!= 0` to
117/// interpret it as a bool.
118pub(crate) extern "C" fn allocate_data_section_adapter(
119    opaque: *mut libc::c_void,
120    size: libc::uintptr_t,
121    alignment: libc::c_uint,
122    section_id: libc::c_uint,
123    section_name: *const libc::c_char,
124    is_read_only: LLVMBool,
125) -> *mut u8 {
126    let adapter = unsafe { &mut *(opaque as *mut MemoryManagerAdapter) };
127    let sname = unsafe { c_str_to_str(section_name) };
128    adapter
129        .memory_manager
130        .allocate_data_section(size, alignment, section_id, sname, is_read_only != 0)
131}
132
133/// Adapter for `finalize_memory`.
134///
135/// If an error is returned, the message is converted into a C string and set in `err_msg_out`.
136pub(crate) extern "C" fn finalize_memory_adapter(
137    opaque: *mut libc::c_void,
138    err_msg_out: *mut *mut libc::c_char,
139) -> libc::c_int {
140    let adapter = unsafe { &mut *(opaque as *mut MemoryManagerAdapter) };
141    match adapter.memory_manager.finalize_memory() {
142        Ok(()) => 0,
143        Err(e) => {
144            let cstring = std::ffi::CString::new(e).unwrap_or_default();
145            unsafe {
146                *err_msg_out = cstring.into_raw();
147            }
148            1
149        },
150    }
151}
152
153/// Adapter for `destroy`.
154///
155/// Called when LLVM is done with the memory manager. Calls `destroy` and drops
156/// the adapter to free resources.
157pub(crate) extern "C" fn destroy_adapter(opaque: *mut libc::c_void) {
158    // Re-box to drop the adapter and its contents.
159    // SAFETY: `opaque` must have been allocated by Box<MemoryManagerAdapter>.
160    let mut adapter = unsafe { Box::from_raw(opaque as *mut MemoryManagerAdapter) };
161
162    // Clean up user-defined resources
163    adapter.memory_manager.destroy();
164
165    // Dropping `adapter` automatically frees the memory
166}
167
168/// Converts a raw C string pointer to a Rust `&str`.
169///
170/// # Safety
171///
172/// The caller must ensure `ptr` points to a valid, null-terminated UTF-8 string.
173/// If the string is invalid UTF-8 or `ptr` is null, an empty string is returned.
174unsafe fn c_str_to_str<'a>(ptr: *const libc::c_char) -> &'a str {
175    if ptr.is_null() {
176        ""
177    } else {
178        unsafe { std::ffi::CStr::from_ptr(ptr) }.to_str().unwrap_or("")
179    }
180}