reifydb_engine/ffi/callbacks/
memory.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2// Copyright (c) 2025 ReifyDB
3
4//! Memory management callbacks for FFI operators
5//!
6//! Provides arena-based memory allocation for FFI operators with fallback to system allocator.
7
8use std::{
9	alloc::{Layout, alloc, dealloc, realloc as system_realloc},
10	cell::RefCell,
11	ptr,
12};
13
14use reifydb_sdk::ffi::arena::Arena;
15
16// Thread-local storage for the current arena
17// All allocations during an FFI operation will use this arena
18thread_local! {
19	static CURRENT_ARENA: RefCell<Option<*mut Arena>> = RefCell::new(None);
20}
21
22/// Set the current arena for this thread
23pub fn set_current_arena(arena: *mut Arena) {
24	CURRENT_ARENA.with(|a| {
25		*a.borrow_mut() = Some(arena);
26	});
27}
28
29/// Clear the current arena for this thread
30pub fn clear_current_arena() {
31	CURRENT_ARENA.with(|a| {
32		*a.borrow_mut() = None;
33	});
34}
35
36/// Allocate memory from the current arena or system allocator
37#[unsafe(no_mangle)]
38pub extern "C" fn host_alloc(size: usize) -> *mut u8 {
39	if size == 0 {
40		return ptr::null_mut();
41	}
42
43	// Try to use the thread-local arena first
44	CURRENT_ARENA.with(|a| {
45		if let Some(arena_ptr) = *a.borrow() {
46			unsafe { (*arena_ptr).alloc(size) }
47		} else {
48			// Fallback to system allocator if no arena set
49			let layout = match Layout::from_size_align(size, 8) {
50				Ok(layout) => layout,
51				Err(_) => return ptr::null_mut(),
52			};
53			unsafe { alloc(layout) }
54		}
55	})
56}
57
58/// Free memory (no-op for arena memory, system free otherwise)
59#[unsafe(no_mangle)]
60pub extern "C" fn host_free(ptr: *mut u8, size: usize) {
61	if ptr.is_null() || size == 0 {
62		return;
63	}
64
65	// Check if this is arena-allocated memory
66	CURRENT_ARENA.with(|a| {
67		if (*a.borrow()).is_some() {
68			// Arena memory is freed automatically - do nothing
69			return;
70		}
71	});
72
73	// Otherwise use system deallocator
74	let layout = match Layout::from_size_align(size, 8) {
75		Ok(layout) => layout,
76		Err(_) => return,
77	};
78	unsafe { dealloc(ptr, layout) }
79}
80
81/// Reallocate memory (allocates new for arena, uses system realloc otherwise)
82#[unsafe(no_mangle)]
83pub extern "C" fn host_realloc(ptr: *mut u8, old_size: usize, new_size: usize) -> *mut u8 {
84	// For arena allocations, we can't realloc in place, so alloc new and copy
85	if ptr.is_null() {
86		return host_alloc(new_size);
87	}
88
89	if new_size == 0 {
90		host_free(ptr, old_size);
91		return ptr::null_mut();
92	}
93
94	// Check if using arena
95	CURRENT_ARENA.with(|a| {
96		if let Some(arena_ptr) = *a.borrow() {
97			// Arena can't realloc, so alloc new and copy
98			let new_ptr = unsafe { (*arena_ptr).alloc(new_size) };
99			if !new_ptr.is_null() {
100				let copy_size = old_size.min(new_size);
101				unsafe {
102					ptr::copy_nonoverlapping(ptr, new_ptr, copy_size);
103				}
104			}
105			// Note: old arena memory will be freed when arena is cleared
106			new_ptr
107		} else {
108			// Use system realloc
109			let old_layout = match Layout::from_size_align(old_size, 8) {
110				Ok(layout) => layout,
111				Err(_) => return ptr::null_mut(),
112			};
113			let new_layout = match Layout::from_size_align(new_size, 8) {
114				Ok(layout) => layout,
115				Err(_) => return ptr::null_mut(),
116			};
117			unsafe { system_realloc(ptr, old_layout, new_layout.size()) }
118		}
119	})
120}
reifydb_engine/ffi/callbacks/memory.rs

reifydb_engine/ffi/callbacks/
memory.rs
