reifydb_flow_operator_sdk/ffi/
arena.rs

1//! Memory arena for FFI allocations
2//!
3//! Provides a simple arena allocator that is automatically cleaned up
4//! after each FFI operator invocation.
5
6use std::{
7	alloc::{Layout, alloc, dealloc},
8	ptr,
9};
10
11/// Memory arena for FFI allocations
12///
13/// All memory allocated from this arena is freed when the arena is dropped.
14pub struct Arena {
15	allocations: Vec<Allocation>,
16}
17
18struct Allocation {
19	ptr: *mut u8,
20	layout: Layout,
21}
22
23// SAFETY: Arena manages raw pointers but ensures proper cleanup
24unsafe impl Send for Arena {}
25unsafe impl Sync for Arena {}
26
27impl Arena {
28	/// Create a new empty arena
29	pub fn new() -> Self {
30		Self {
31			allocations: Vec::new(),
32		}
33	}
34
35	/// Allocate memory from the arena
36	///
37	/// Returns null on allocation failure
38	pub fn alloc(&mut self, size: usize) -> *mut u8 {
39		if size == 0 {
40			return ptr::null_mut();
41		}
42
43		// Align to 8 bytes for safety
44		let layout = match Layout::from_size_align(size, 8) {
45			Ok(layout) => layout,
46			Err(_) => return ptr::null_mut(),
47		};
48
49		// SAFETY: We're allocating with a valid layout
50		let ptr = unsafe { alloc(layout) };
51
52		if ptr.is_null() {
53			return ptr::null_mut();
54		}
55
56		self.allocations.push(Allocation {
57			ptr,
58			layout,
59		});
60		ptr
61	}
62
63	/// Allocate and zero-initialize memory from the arena
64	pub fn alloc_zeroed(&mut self, size: usize) -> *mut u8 {
65		let ptr = self.alloc(size);
66		if !ptr.is_null() {
67			// SAFETY: We just allocated this memory
68			unsafe {
69				ptr::write_bytes(ptr, 0, size);
70			}
71		}
72		ptr
73	}
74
75	/// Allocate memory for a type T
76	pub fn alloc_type<T>(&mut self) -> *mut T {
77		self.alloc(std::mem::size_of::<T>()) as *mut T
78	}
79
80	/// Allocate and copy bytes into the arena
81	pub fn copy_bytes(&mut self, bytes: &[u8]) -> *mut u8 {
82		if bytes.is_empty() {
83			return ptr::null_mut();
84		}
85
86		let ptr = self.alloc(bytes.len());
87		if !ptr.is_null() {
88			// SAFETY: We just allocated sufficient memory
89			unsafe {
90				ptr::copy_nonoverlapping(bytes.as_ptr(), ptr, bytes.len());
91			}
92		}
93		ptr
94	}
95
96	/// Clear all allocations
97	pub fn clear(&mut self) {
98		// Free all allocations
99		for allocation in self.allocations.drain(..) {
100			// SAFETY: We allocated this memory with the same layout
101			unsafe {
102				dealloc(allocation.ptr, allocation.layout);
103			}
104		}
105	}
106
107	/// Get the number of allocations
108	pub fn allocation_count(&self) -> usize {
109		self.allocations.len()
110	}
111
112	/// Get the total allocated size
113	pub fn total_size(&self) -> usize {
114		self.allocations.iter().map(|a| a.layout.size()).sum()
115	}
116}
117
118impl Default for Arena {
119	fn default() -> Self {
120		Self::new()
121	}
122}
123
124impl Drop for Arena {
125	fn drop(&mut self) {
126		self.clear();
127	}
128}
129
130#[cfg(test)]
131mod tests {
132	use super::*;
133
134	#[test]
135	fn test_arena_basic() {
136		let mut arena = Arena::new();
137
138		let ptr1 = arena.alloc(100);
139		assert!(!ptr1.is_null());
140
141		let ptr2 = arena.alloc(200);
142		assert!(!ptr2.is_null());
143
144		assert_eq!(arena.allocation_count(), 2);
145		assert_eq!(arena.total_size(), 300);
146	}
147
148	#[test]
149	fn test_arena_zeroed() {
150		let mut arena = Arena::new();
151
152		let ptr = arena.alloc_zeroed(100);
153		assert!(!ptr.is_null());
154
155		unsafe {
156			for i in 0..100 {
157				assert_eq!(*ptr.add(i), 0);
158			}
159		}
160	}
161
162	#[test]
163	fn test_arena_copy_bytes() {
164		let mut arena = Arena::new();
165
166		let data = vec![1u8, 2, 3, 4, 5];
167		let ptr = arena.copy_bytes(&data);
168		assert!(!ptr.is_null());
169
170		unsafe {
171			for i in 0..5 {
172				assert_eq!(*ptr.add(i), data[i]);
173			}
174		}
175	}
176
177	#[test]
178	fn test_arena_clear() {
179		let mut arena = Arena::new();
180
181		arena.alloc(100);
182		arena.alloc(200);
183		assert_eq!(arena.allocation_count(), 2);
184
185		arena.clear();
186		assert_eq!(arena.allocation_count(), 0);
187	}
188}
reifydb_flow_operator_sdk/ffi/arena.rs

reifydb_flow_operator_sdk/ffi/
arena.rs
