rarena_allocator/
lib.rs

1#![doc = include_str!("../README.md")]
2#![cfg_attr(not(any(feature = "std", test)), no_std)]
3#![cfg_attr(docsrs, feature(doc_cfg))]
4#![cfg_attr(docsrs, allow(unused_attributes))]
5#![deny(missing_docs)]
6
7#[cfg(not(any(feature = "std", feature = "alloc")))]
8compile_error!("`rarena-allocator` requires either the 'std' or 'alloc' feature to be enabled");
9
10#[cfg(not(feature = "std"))]
11extern crate alloc as std;
12
13#[cfg(feature = "std")]
14extern crate std;
15
16mod allocator;
17mod common;
18mod error;
19mod memory;
20mod options;
21mod sealed;
22
23#[cfg(test)]
24#[macro_use]
25mod tests;
26
27use core::mem;
28use dbutils::checksum::{BuildChecksumer, Checksumer};
29
30pub use allocator::Allocator;
31pub use dbutils::checksum;
32pub use either;
33pub use error::*;
34pub use options::*;
35
36const FREELIST_OFFSET: usize = 1;
37const FREELIST_SIZE: usize = mem::size_of::<Freelist>();
38const MAGIC_TEXT: [u8; 2] = *b"al";
39const MAGIC_TEXT_OFFSET: usize = FREELIST_OFFSET + FREELIST_SIZE;
40const MAGIC_TEXT_SIZE: usize = MAGIC_TEXT.len();
41const MAGIC_VERISON_OFFSET: usize = MAGIC_TEXT_OFFSET + MAGIC_TEXT_SIZE;
42const MAGIC_VERISON_SIZE: usize = mem::size_of::<u16>();
43const VERSION_OFFSET: usize = MAGIC_VERISON_OFFSET + MAGIC_VERISON_SIZE;
44const VERSION_SIZE: usize = mem::size_of::<u16>();
45const CURRENT_VERSION: u16 = 0;
46const SENTINEL_SEGMENT_NODE_OFFSET: u32 = u32::MAX;
47const SENTINEL_SEGMENT_NODE_SIZE: u32 = u32::MAX;
48
49#[cfg(all(feature = "std", not(target_family = "wasm")))]
50static PAGE_SIZE: std::sync::LazyLock<u32> = std::sync::LazyLock::new(|| {
51  #[cfg(not(windows))]
52  {
53    rustix::param::page_size() as u32
54  }
55
56  #[cfg(windows)]
57  {
58    use winapi::um::sysinfoapi::{GetSystemInfo, SYSTEM_INFO};
59
60    unsafe {
61      let mut system_info: SYSTEM_INFO = std::mem::zeroed();
62      GetSystemInfo(&mut system_info);
63      system_info.dwPageSize
64    }
65  }
66});
67
68#[cfg(not(all(feature = "std", not(target_family = "wasm"))))]
69static PAGE_SIZE: &u32 = &4096;
70
71/// Enumeration of possible methods to seek within an [`Allocator`].
72#[derive(Copy, PartialEq, Eq, Clone, Debug)]
73pub enum ArenaPosition {
74  /// Sets the offset to the provided number of bytes.
75  Start(u32),
76
77  /// Sets the offset to the capacity of the allocator minus the provided number of bytes.
78  End(u32),
79
80  /// Sets the offset to the current position plus the specified number of
81  /// bytes.
82  ///
83  /// It is possible to seek beyond the end of an object, but it's an error to
84  /// seek before byte 0.
85  Current(i64),
86}
87
88bitflags::bitflags! {
89  #[derive(Clone, Copy)]
90  struct MemoryFlags: u8 {
91    const ON_DISK = 0b0000_0001;
92    #[cfg(all(feature = "memmap", not(target_family = "wasm")))]
93    const MMAP = 0b0000_0010;
94  }
95}
96
97/// The memory chunk allocated from the allocator.
98pub trait Buffer {
99  /// Returns how many bytes of accessible buffer occupies.
100  fn capacity(&self) -> usize;
101
102  /// Returns the accessible offset to the pointer of the allocator.
103  fn offset(&self) -> usize;
104
105  /// Returns how many bytes of the whole buffer occupies.
106  fn buffer_capacity(&self) -> usize;
107
108  /// Returns the offset to the pointer of the allocator.
109  fn buffer_offset(&self) -> usize;
110
111  /// Detach the value from the ARENA, which means when the value is dropped,
112  /// the underlying buffer will not be collected for futhur allocation.
113  ///
114  /// ## Safety
115  /// - The caller must ensure the value is dropped before the ARENA is dropped.
116  unsafe fn detach(&mut self);
117
118  /// Flush the buffer to the disk.
119  #[cfg(all(feature = "memmap", not(target_family = "wasm")))]
120  #[cfg_attr(docsrs, doc(cfg(all(feature = "memmap", not(target_family = "wasm")))))]
121  fn flush(&self) -> std::io::Result<()>;
122
123  /// Asynchronously flush the buffer to the disk.
124  #[cfg(all(feature = "memmap", not(target_family = "wasm")))]
125  #[cfg_attr(docsrs, doc(cfg(all(feature = "memmap", not(target_family = "wasm")))))]
126  fn flush_async(&self) -> std::io::Result<()>;
127}
128
129#[inline]
130fn write_sanity(freelist: u8, magic_version: u16, data: &mut [u8]) {
131  data[FREELIST_OFFSET] = freelist;
132  data[MAGIC_TEXT_OFFSET..MAGIC_TEXT_OFFSET + MAGIC_TEXT_SIZE].copy_from_slice(MAGIC_TEXT.as_ref());
133  data[MAGIC_VERISON_OFFSET..MAGIC_VERISON_OFFSET + MAGIC_VERISON_SIZE]
134    .copy_from_slice(magic_version.to_le_bytes().as_ref());
135  data[VERSION_OFFSET..VERSION_OFFSET + VERSION_SIZE]
136    .copy_from_slice(CURRENT_VERSION.to_le_bytes().as_ref());
137}
138
139#[cfg(all(feature = "memmap", not(target_family = "wasm")))]
140#[inline]
141fn sanity_check(
142  freelist: Option<Freelist>,
143  magic_version: u16,
144  data: &[u8],
145) -> std::io::Result<Freelist> {
146  let stored_freelist = data[FREELIST_OFFSET];
147  let stored_freelist = Freelist::try_from(stored_freelist).map_err(invalid_data)?;
148
149  if let Some(freelist) = freelist {
150    if stored_freelist != freelist {
151      return Err(bad_freelist());
152    }
153  }
154
155  let stored_magic_version: u16 = u16::from_le_bytes(
156    data[MAGIC_VERISON_OFFSET..MAGIC_VERISON_OFFSET + MAGIC_VERISON_SIZE]
157      .try_into()
158      .unwrap(),
159  );
160  let version: u16 = u16::from_le_bytes(
161    data[VERSION_OFFSET..VERSION_OFFSET + VERSION_SIZE]
162      .try_into()
163      .unwrap(),
164  );
165
166  if stored_magic_version != magic_version {
167    return Err(invalid_data(MagicVersionMismatch::new(
168      magic_version,
169      stored_magic_version,
170    )));
171  }
172
173  if version != CURRENT_VERSION {
174    return Err(invalid_data(VersionMismatch::new(CURRENT_VERSION, version)));
175  }
176
177  if data[MAGIC_TEXT_OFFSET..MAGIC_TEXT_OFFSET + MAGIC_TEXT_SIZE] != MAGIC_TEXT {
178    return Err(bad_magic());
179  }
180  Ok(stored_freelist)
181}
182
183#[cfg(all(feature = "memmap", not(target_family = "wasm")))]
184#[inline]
185fn invalid_data<E: std::error::Error + Send + Sync + 'static>(e: E) -> std::io::Error {
186  std::io::Error::new(std::io::ErrorKind::InvalidData, e)
187}
188
189#[cfg(all(feature = "memmap", not(target_family = "wasm")))]
190#[inline]
191fn invalid_input<E: Into<std::boxed::Box<dyn std::error::Error + Send + Sync>>>(
192  e: E,
193) -> std::io::Error {
194  std::io::Error::new(std::io::ErrorKind::InvalidInput, e)
195}
196
197#[cfg(all(feature = "memmap", not(target_family = "wasm")))]
198#[inline]
199fn range_out_of_bounds(offset: usize, len: usize, cap: usize) -> std::io::Error {
200  #[derive(Debug)]
201  struct Err {
202    offset: usize,
203    len: usize,
204    cap: usize,
205  }
206
207  impl std::fmt::Display for Err {
208    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
209      write!(
210        f,
211        "range [{}..{}] out of bounds (cap={})",
212        self.offset, self.len, self.cap
213      )
214    }
215  }
216
217  impl std::error::Error for Err {}
218
219  std::io::Error::new(std::io::ErrorKind::InvalidInput, Err { offset, len, cap })
220}
221
222#[cfg(all(feature = "memmap", not(target_family = "wasm")))]
223#[inline]
224fn bad_magic() -> std::io::Error {
225  std::io::Error::new(std::io::ErrorKind::InvalidData, "arena has bad magic")
226}
227
228#[cfg(all(feature = "memmap", not(target_family = "wasm")))]
229#[inline]
230fn bad_freelist() -> std::io::Error {
231  std::io::Error::new(std::io::ErrorKind::InvalidData, "freelist mismatch")
232}
233
234#[inline]
235const fn decode_segment_node(val: u64) -> (u32, u32) {
236  ((val >> 32) as u32, val as u32)
237}
238
239#[inline]
240const fn encode_segment_node(size: u32, next: u32) -> u64 {
241  ((size as u64) << 32) | next as u64
242}
243
244/// Calculates the aligned offset for a given type `T` starting from `current_offset`.
245///
246/// This function aligns the given `current_offset` to the next boundary that satisfies the alignment requirements of type `T`.
247///
248/// # Parameters
249///
250/// - `current_offset`: The initial offset that needs to be aligned.
251///
252/// # Returns
253///
254/// The aligned offset that is the next multiple of the alignment requirement of type `T`.
255///
256/// ## Examples
257///
258/// ```ignore
259/// use std::mem;
260///
261/// #[repr(C, align(8))]
262/// struct Meta {
263///     a: u64,
264///     b: u64,
265/// }
266///
267/// let initial_offset: u32 = 1;
268/// let aligned_offset = align_offset::<Meta>(initial_offset);
269/// assert_eq!(aligned_offset, 8);
270/// ```
271///
272/// # Explanation
273///
274/// - Given an `alignment` of type `T`, this function calculates the next aligned offset from `current_offset`.
275/// - It ensures that the result is a multiple of `alignment` by adding `alignment - 1` to `current_offset` and then clearing the lower bits using bitwise AND with the negation of `alignment - 1`.
276///
277/// ```ignore
278/// let alignment = mem::align_of::<T>() as u32;
279/// (current_offset + alignment - 1) & !(alignment - 1)
280/// ```
281#[inline]
282pub const fn align_offset<T>(current_offset: u32) -> u32 {
283  let alignment = core::mem::align_of::<T>() as u32;
284  (current_offset + alignment - 1) & !(alignment - 1)
285}
286
287#[cfg(feature = "std")]
288macro_rules! write_byte_order {
289  ($write_name:ident::$put_name:ident::$converter:ident($ty:ident, $endian:literal)) => {
290    paste::paste! {
291      #[doc = "Write a `" $ty "` value into the buffer in " $endian " byte order, return an error if the buffer does not have enough space."]
292      #[inline]
293      #[cfg(feature = "std")]
294      pub fn $write_name(&mut self, value: $ty) -> std::io::Result<()> {
295        self.$put_name(value).map_err(|e| std::io::Error::new(std::io::ErrorKind::WriteZero, e))
296      }
297    }
298  }
299}
300
301macro_rules! put_byte_order {
302  ($name:ident::$converter:ident($ty:ident, $endian:literal)) => {
303    paste::paste! {
304      #[doc = "Put a `" $ty "` value into the buffer in " $endian " byte order, return an error if the buffer does not have enough space."]
305      #[inline]
306      pub fn $name(&mut self, value: $ty) -> Result<(), InsufficientBuffer> {
307        const SIZE: usize = core::mem::size_of::<$ty>();
308
309        if self.len + SIZE > self.capacity() {
310          return Err(InsufficientBuffer::with_information(SIZE as u64, (self.capacity() - self.len) as u64));
311        }
312
313        // SAFETY: We have checked the buffer size.
314        unsafe { self. [< $name _unchecked >](value); }
315        Ok(())
316      }
317
318      #[doc = "Put a `" $ty "` value into the buffer in " $endian " byte order, without doing bounds checking."]
319      ///
320      #[doc = "For a safe alternative see [`" $name "`](Self::" $name ")."]
321      ///
322      /// ## Safety
323      ///
324      /// Calling this method if the buffer does not have enough space to hold the value is *[undefined behavior]*.
325      ///
326      /// [undefined behavior]: https://doc.rust-lang.org/reference/behavior-considered-undefined.html
327      #[inline]
328      pub unsafe fn [< $name _unchecked >] (&mut self, value: $ty) {
329        const SIZE: usize = core::mem::size_of::<$ty>();
330
331        let cur = self.len;
332        let buf = self.buffer_mut();
333        buf[cur..cur + SIZE].copy_from_slice(&value.$converter());
334        self.len += SIZE;
335      }
336    }
337  }
338}
339
340#[cfg(feature = "std")]
341macro_rules! write_varint {
342  ($write_name:ident::$put_name:ident($ty:ident)) => {
343    paste::paste! {
344      #[doc = "Write a `" $ty "` value into the buffer in LEB128 format, return number of bytes written on success, or an error if the buffer does not have enough space."]
345      #[inline]
346      #[cfg(feature = "std")]
347      pub fn $write_name(&mut self, value: $ty) -> std::io::Result<usize> {
348        self.$put_name(value).map_err(|e| std::io::Error::new(std::io::ErrorKind::WriteZero, e))
349      }
350    }
351  }
352}
353
354macro_rules! put_varint {
355  ($name:ident($ty:ident)) => {
356    paste::paste! {
357      #[doc = "Put a `" $ty "` value into the buffer in LEB128 format, return an error if the buffer does not have enough space."]
358      ///
359      /// Returns the number of bytes written.
360      #[inline]
361      pub fn $name(&mut self, value: $ty) -> Result<usize, dbutils::error::InsufficientBuffer> {
362        let buf = unsafe {
363          core::slice::from_raw_parts_mut(self.as_mut_ptr().add(self.len), self.capacity() - self.len)
364        };
365        dbutils::leb128::[< encode_ $ty _varint_to >](value, buf)
366          .inspect(|len| self.len += *len)
367          .map_err(Into::into)
368      }
369
370      #[doc = "Put a `" $ty "` value into the buffer in LEB128 format, without doing error checking."]
371      ///
372      /// Returns the number of bytes written.
373      ///
374      /// # Panics
375      ///
376      #[doc = "- If the buffer does not have enough space to hold the `" $ty "`."]
377      #[inline]
378      pub fn [< $name _unchecked >] (&mut self, value: $ty) -> usize {
379        let buf = unsafe {
380          core::slice::from_raw_parts_mut(self.as_mut_ptr().add(self.len), self.capacity() - self.len)
381        };
382        dbutils::leb128::[< encode_ $ty _varint_to >] (value, buf).inspect(|len| self.len += *len).unwrap()
383      }
384    }
385  }
386}
387
388macro_rules! impl_bytes_mut_utils {
389  (align) => {
390    /// Add paddings to the buffer according to the alignment of `T`.
391    ///
392    /// Returns a well-aligned pointer for `T`
393    #[inline]
394    pub fn align_to<T>(&mut self) -> Result<core::ptr::NonNull<T>, InsufficientBuffer> {
395      if mem::size_of::<T>() == 0 {
396        return Ok(core::ptr::NonNull::dangling());
397      }
398
399      let align_offset = crate::align_offset::<T>(self.allocated.memory_offset + self.len as u32);
400
401      if align_offset > self.allocated.memory_offset + self.allocated.memory_size {
402        return Err(InsufficientBuffer::with_information((align_offset as u64 - self.len as u64 - self.allocated.memory_offset as u64), (self.allocated.memory_size as u64 - self.len as u64)));
403      }
404
405      self.len = (align_offset - self.allocated.memory_offset) as usize;
406      // SAFETY: We have checked the buffer size, and apply the align
407      Ok(unsafe {
408        core::ptr::NonNull::new_unchecked(self.as_mut_ptr().add(self.len).cast::<T>())
409      })
410    }
411
412    /// Set the buffer length.
413    ///
414    /// # Panics
415    /// - If `len` is greater than the capacity of the buffer.
416    #[inline]
417    pub fn set_len(&mut self, len: usize) {
418      if len > self.capacity() {
419        panic!("length out of bounds");
420      }
421
422      if len == self.len {
423        return;
424      }
425
426      let olen = self.len;
427      self.len = len;
428      if len > olen {
429        unsafe { core::ptr::write_bytes(self.as_mut_ptr().add(olen), 0, len - olen) };
430      } else {
431        unsafe { core::ptr::write_bytes(self.as_mut_ptr().add(len), 0, olen - len) };
432      }
433    }
434
435    /// Put `T` into the buffer, return an error if the buffer does not have enough space.
436    ///
437    /// You may want to use `put_aligned` instead of this method.
438    ///
439    /// ## Safety
440    ///
441    /// - Must invoke `align_to` before invoking this method, if `T` is not ZST.
442    /// - If `T` needs to be dropped and callers invoke [`detach`](crate::Buffer::detach),
443    ///   then the caller must ensure that the `T` is dropped before the allocator is dropped.
444    ///   Otherwise, it will lead to memory leaks.
445    ///
446    /// - If this is file backed allocator, then `T` must be recoverable from bytes.
447    ///   1. Types require allocation are not recoverable.
448    ///   2. Pointers are not recoverable, like `*const T`, `*mut T`, `NonNull` and any structs contains pointers,
449    ///      although those types are on stack, but they cannot be recovered, when reopens the file.
450    pub unsafe fn put<T>(&mut self, val: T) -> Result<&mut T, InsufficientBuffer> { unsafe {
451      let size = core::mem::size_of::<T>();
452
453      if self.len + size > self.capacity() {
454        return Err(InsufficientBuffer::with_information(size as u64, (self.capacity() - self.len) as u64));
455      }
456
457      // SAFETY: We have checked the buffer size.
458      let ptr = self.as_mut_ptr().add(self.len).cast::<T>();
459      ptr.write(val);
460      self.len += size;
461      Ok(&mut *ptr)
462    }}
463
464    /// Put `T` into the buffer, return an error if the buffer does not have enough space.
465    ///
466    /// ## Safety
467    ///
468    /// - If `T` needs to be dropped and callers invoke [`RefMut::detach`](crate::RefMut::detach),
469    ///   then the caller must ensure that the `T` is dropped before the allocator is dropped.
470    ///   Otherwise, it will lead to memory leaks.
471    ///
472    /// - If this is file backed allocator, then `T` must be recoverable from bytes.
473    ///   1. Types require allocation are not recoverable.
474    ///   2. Pointers are not recoverable, like `*const T`, `*mut T`, `NonNull` and any structs contains pointers,
475    ///      although those types are on stack, but they cannot be recovered, when reopens the file.
476    pub unsafe fn put_aligned<T>(&mut self, val: T) -> Result<&mut T, InsufficientBuffer> { unsafe {
477      let mut ptr = self.align_to::<T>()?;
478
479      ptr.as_ptr().write(val);
480      self.len += ::core::mem::size_of::<T>();
481      Ok(ptr.as_mut())
482    }}
483  };
484  (slice) => {
485    /// Put a bytes slice into the buffer, return an error if the buffer does not have enough space.
486    #[inline]
487    pub fn put_slice(&mut self, slice: &[u8]) -> Result<(), InsufficientBuffer> {
488      let size = slice.len();
489
490      if self.len + size > self.capacity() {
491        return Err(InsufficientBuffer::with_information(size as u64, (self.capacity() - self.len) as u64));
492      }
493
494      // SAFETY: We have checked the buffer size.
495      unsafe { self.put_slice_unchecked(slice); }
496      Ok(())
497    }
498
499    /// Put a bytes slice into the buffer, without doing bounds checking.
500    ///
501    /// For a safe alternative see [`put_slice`](Self::put_slice).
502    ///
503    /// ## Safety
504    ///
505    /// Calling this method if the buffer does not have enough space to hold the slice is *[undefined behavior]*.
506    ///
507    /// [undefined behavior]: https://doc.rust-lang.org/reference/behavior-considered-undefined.html
508    #[inline]
509    pub unsafe fn put_slice_unchecked(&mut self, slice: &[u8]) {
510      let size = slice.len();
511      let cur = self.len;
512      let buf = self.buffer_mut();
513      buf[cur..cur + size].copy_from_slice(slice);
514      self.len += size;
515    }
516  };
517  ($($ty:ident), +$(,)?) => {
518    $(
519      paste::paste! {
520        put_byte_order!([< put_ $ty _be>]::to_be_bytes($ty, "big-endian"));
521        put_byte_order!([< put_ $ty _le >]::to_le_bytes($ty, "little-endian"));
522        put_byte_order!([< put_ $ty _ne >]::to_ne_bytes($ty, "native-endian"));
523        #[cfg(feature="std")]
524        write_byte_order!([< write_ $ty _be>]::[< put_ $ty _be>]::to_be_bytes($ty, "big-endian"));
525        #[cfg(feature="std")]
526        write_byte_order!([< write_ $ty _le >]::[< put_ $ty _le >]::to_le_bytes($ty, "little-endian"));
527        #[cfg(feature="std")]
528        write_byte_order!([< write_ $ty _ne >]::[< put_ $ty _ne >]::to_ne_bytes($ty, "native-endian"));
529      }
530    )*
531  };
532  (leb($($ty:ident), +$(,)?)) => {
533    $(
534      paste::paste! {
535        put_varint!([< put_ $ty _varint>]($ty));
536        #[cfg(feature="std")]
537        write_varint!([< write_ $ty _varint>]::[< put_ $ty _varint>]($ty));
538      }
539    )*
540  };
541  (8) => {
542    /// Put a `u8` value into the buffer, return an error if the buffer does not have enough space.
543    #[inline]
544    pub fn put_u8(&mut self, value: u8) -> Result<(), InsufficientBuffer> {
545      const SIZE: usize = core::mem::size_of::<u8>();
546
547      if self.len + SIZE > self.capacity() {
548        return Err(InsufficientBuffer::with_information(SIZE as u64, (self.capacity() - self.len) as u64));
549      }
550
551      // SAFETY: We have checked the buffer size.
552      unsafe { self.put_u8_unchecked(value); }
553      Ok(())
554    }
555
556    /// Put a `u8` value into the buffer, without doing bounds checking.
557    ///
558    /// For a safe alternative see [`put_u8`](Self::put_u8).
559    ///
560    /// ## Safety
561    ///
562    /// Calling this method if the buffer does not have enough space to hold the value is *[undefined behavior]*.
563    ///
564    /// [undefined behavior]: https://doc.rust-lang.org/reference/behavior-considered-undefined.html
565    #[inline]
566    pub unsafe fn put_u8_unchecked(&mut self, value: u8) {
567      const SIZE: usize = core::mem::size_of::<u8>();
568
569      let cur = self.len;
570      let buf = self.buffer_mut();
571      buf[cur..cur + SIZE].copy_from_slice(&[value]);
572      self.len += SIZE;
573    }
574
575    /// Put a `i8` value into the buffer, return an error if the buffer does not have enough space.
576    #[inline]
577    pub fn put_i8(&mut self, value: i8) -> Result<(), InsufficientBuffer> {
578      self.put_u8(value as u8)
579    }
580
581    /// Put a `i8` value into the buffer, without doing bounds checking.
582    ///
583    /// For a safe alternative see [`put_i8`](Self::put_i8).
584    ///
585    /// ## Safety
586    ///
587    /// Calling this method if the buffer does not have enough space to hold the value is *[undefined behavior]*.
588    ///
589    /// [undefined behavior]: https://doc.rust-lang.org/reference/behavior-considered-undefined.html
590    #[inline]
591    pub unsafe fn put_i8_unchecked(&mut self, value: i8) { unsafe {
592      self.put_u8_unchecked(value as u8)
593    }}
594  };
595}
596
597macro_rules! get_byte_order {
598  ($name:ident::$converter:ident($ty:ident, $endian:literal)) => {
599    paste::paste! {
600      #[doc = "Get a `" $ty "` value from the buffer in " $endian " byte order, return an error if the buffer does not have enough bytes."]
601      #[inline]
602      pub fn $name(&mut self) -> Result<$ty, IncompleteBuffer> {
603        const SIZE: usize = core::mem::size_of::<$ty>();
604
605        if self.len < SIZE {
606          return Err(IncompleteBuffer::with_information(SIZE as u64, self.len as u64));
607        }
608
609        // SAFETY: We have checked the buffer size.
610        unsafe { Ok(self.[< $name _unchecked >]()) }
611      }
612
613      #[doc = "Get a `" $ty "` value from the buffer in " $endian " byte order, without doing bounds checking."]
614      ///
615      #[doc = "For a safe alternative see [`" $name "`](Self::" $name ")."]
616      ///
617      /// ## Safety
618      ///
619      /// Calling this method if the buffer does not have enough bytes to read the value is *[undefined behavior]*.
620      ///
621      /// [undefined behavior]: https://doc.rust-lang.org/reference/behavior-considered-undefined.html
622      #[inline]
623      pub unsafe fn [< $name _unchecked >](&mut self) -> $ty {
624        const SIZE: usize = core::mem::size_of::<$ty>();
625
626        let cur = self.len - SIZE;
627        let buf = self.buffer();
628        let value = <$ty>::from_be_bytes(buf[cur..cur + SIZE].try_into().unwrap());
629        self.len -= SIZE;
630        value
631      }
632    }
633  }
634}
635
636macro_rules! get_varint {
637  ($name:ident($ty:ident)) => {
638    paste::paste! {
639      #[doc = "Get a `" $ty "` value from the buffer in LEB128 format, return an error if the buffer does not have a valid LEB128 format `" $ty "`."]
640      ///
641      /// # Returns
642      ///
643      /// - The first element of the tuple is the number of bytes read from the buffer.
644      /// - The second element of the tuple is the decoded value.
645      #[doc = "- The second element of the tuple is the decoded `" $ty "`."]
646      #[inline]
647      pub fn $name(&self) -> Result<(usize, $ty), dbutils::leb128::DecodeVarintError> {
648        dbutils::leb128::[< decode_ $ty _varint >](self)
649          .map(|(bytes, value)| (bytes, value as $ty))
650      }
651
652      #[doc = "Get a `" $ty "` value from the buffer in LEB128 format, without doing checking."]
653      ///
654      #[doc = "For a safe alternative see [`" $name "`](Self::" $name ")."]
655      ///
656      /// # Panics
657      ///
658      #[doc = "- If the buffer does not have a valid LEB128 format `" $ty "`."]
659      #[inline]
660      pub fn [< $name _unchecked >](&mut self) -> (usize, $ty) {
661        dbutils::leb128::[< decode_ $ty _varint >](self)
662          .map(|(bytes, value)| (bytes, value as $ty))
663          .unwrap()
664      }
665    }
666  }
667}
668
669macro_rules! impl_bytes_utils {
670  (slice) => {
671    /// Get a byte slice from the buffer, return an error if the buffer does not have enough bytes.
672    #[inline]
673    pub fn get_slice(&self, size: usize) -> Result<&[u8], IncompleteBuffer> {
674      if self.len < size {
675        return Err(IncompleteBuffer::with_information(size as u64, self.len as u64));
676      }
677
678      // SAFETY: We have checked the buffer size.
679      unsafe { Ok(self.get_slice_unchecked(size)) }
680    }
681
682    /// Get a byte slice from the buffer, without doing bounds checking.
683    ///
684    /// For a safe alternative see [`get_slice`](Self::get_slice).
685    ///
686    /// ## Safety
687    ///
688    /// Calling this method if the buffer does not have enough bytes to read the slice is *[undefined behavior]*.
689    ///
690    /// [undefined behavior]: https://doc.rust-lang.org/reference/behavior-considered-undefined.html
691    #[inline]
692    pub unsafe fn get_slice_unchecked(&self, size: usize) -> &[u8] {
693      let buf = self.buffer();
694      &buf[..size]
695    }
696
697    /// Get a mutable byte slice from the buffer, return an error if the buffer does not have enough bytes.
698    #[inline]
699    pub fn get_slice_mut(&mut self, size: usize) -> Result<&mut [u8], IncompleteBuffer> {
700      if self.len < size {
701        return Err(IncompleteBuffer::with_information(size as u64, self.len as u64));
702      }
703
704      // SAFETY: We have checked the buffer size.
705      unsafe { Ok(self.get_slice_mut_unchecked(size)) }
706    }
707
708    /// Get a mutable byte slice from the buffer, without doing bounds checking.
709    ///
710    /// For a safe alternative see [`get_slice_mut`](Self::get_slice_mut).
711    ///
712    /// ## Safety
713    ///
714    /// Calling this method if the buffer does not have enough bytes to read the slice is *[undefined behavior]*.
715    ///
716    /// [undefined behavior]: https://doc.rust-lang.org/reference/behavior-considered-undefined.html
717    #[inline]
718    pub unsafe fn get_slice_mut_unchecked(&mut self, size: usize) -> &mut [u8] {
719      let buf = self.buffer_mut();
720      &mut buf[..size]
721    }
722  };
723  ($($ty:ident), +$(,)?) => {
724    $(
725      paste::paste! {
726        get_byte_order!([< get_ $ty _be >]::from_be_bytes($ty, "big-endian"));
727        get_byte_order!([< get_ $ty _le >]::from_le_bytes($ty, "little-endian"));
728        get_byte_order!([< get_ $ty _ne >]::from_ne_bytes($ty, "native-endian"));
729      }
730    )*
731  };
732  (leb($($ty:ident), +$(,)?)) => {
733    $(
734      paste::paste! {
735        get_varint!([< get_ $ty _varint >]($ty));
736      }
737    )*
738  };
739  (8) => {
740    /// Get a `u8` value from the buffer, return an error if the buffer does not have enough bytes.
741    #[inline]
742    pub fn get_u8(&mut self) -> Result<u8, IncompleteBuffer> {
743      if self.len < 1 {
744        return Err(IncompleteBuffer::with_information(1, self.len as u64));
745      }
746
747      // SAFETY: We have checked the buffer size.
748      unsafe { Ok(self.get_u8_unchecked()) }
749    }
750
751    /// Get a `u8` value from the buffer, without doing bounds checking.
752    ///
753    /// For a safe alternative see [`get_u8`](Self::get_u8).
754    ///
755    /// ## Safety
756    ///
757    /// Calling this method if the buffer does not have enough bytes to read the value is *[undefined behavior]*.
758    ///
759    /// [undefined behavior]: https://doc.rust-lang.org/reference/behavior-considered-undefined.html
760    #[inline]
761    pub unsafe fn get_u8_unchecked(&mut self) -> u8 {
762      let cur = self.len - 1;
763      let buf = self.buffer();
764      let value = buf[cur];
765      self.len -= 1;
766      value
767    }
768
769    /// Get a `i8` value from the buffer, return an error if the buffer does not have enough bytes.
770    #[inline]
771    pub fn get_i8(&mut self) -> Result<i8, IncompleteBuffer> {
772      self.get_u8().map(|v| v as i8)
773    }
774
775    /// Get a `i8` value from the buffer, without doing bounds checking.
776    ///
777    /// For a safe alternative see [`get_i8`](Self::get_i8).
778    ///
779    /// ## Safety
780    ///
781    /// Calling this method if the buffer does not have enough bytes to read the value is *[undefined behavior]*.
782    ///
783    /// [undefined behavior]: https://doc.rust-lang.org/reference/behavior-considered-undefined.html
784    #[inline]
785    pub unsafe fn get_i8_unchecked(&mut self) -> i8 { unsafe {
786      self.get_u8_unchecked() as i8
787    }}
788  };
789}
790
791#[cfg(feature = "std")]
792macro_rules! impl_write_in {
793  () => {
794    #[inline]
795    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
796      self
797        .put_slice(buf)
798        .map_err(|e| std::io::Error::new(std::io::ErrorKind::WriteZero, e))
799        .map(|_| buf.len())
800    }
801
802    #[inline(always)]
803    fn flush(&mut self) -> std::io::Result<()> {
804      Ok(())
805    }
806  };
807}
808
809macro_rules! impl_write {
810  ($ident: ident) => {
811    #[cfg(feature = "std")]
812    impl std::io::Write for $ident {
813      impl_write_in!();
814    }
815  };
816  ($ident: ident<'a>) => {
817    #[cfg(feature = "std")]
818    impl<'a> std::io::Write for $ident<'a> {
819      impl_write_in!();
820    }
821  };
822  ($ident: ident<T>) => {
823    #[cfg(feature = "std")]
824    impl<T> std::io::Write for $ident<T> {
825      impl_write_in!();
826    }
827  };
828  ($ident: ident<A>) => {
829    #[cfg(feature = "std")]
830    impl<A: $crate::Allocator> std::io::Write for $ident<A> {
831      impl_write_in!();
832    }
833  };
834  ($ident: ident<'a, A>) => {
835    #[cfg(feature = "std")]
836    impl<A: $crate::Allocator> std::io::Write for $ident<'_, A> {
837      impl_write_in!();
838    }
839  };
840}
841
842/// The metadata of the structs allocated from ARENA.
843#[derive(Debug, Copy, Clone, PartialEq, Eq)]
844pub struct Meta {
845  parent_ptr: *const u8,
846  memory_offset: u32,
847  memory_size: u32,
848  ptr_offset: u32,
849  ptr_size: u32,
850}
851
852unsafe impl Send for Meta {}
853unsafe impl Sync for Meta {}
854
855impl Meta {
856  #[inline]
857  const fn null(parent_ptr: *const u8) -> Self {
858    Self {
859      parent_ptr,
860      memory_offset: 0,
861      memory_size: 0,
862      ptr_offset: 0,
863      ptr_size: 0,
864    }
865  }
866
867  #[inline]
868  const fn new(parent_ptr: *const u8, memory_offset: u32, memory_size: u32) -> Self {
869    Self {
870      parent_ptr,
871      memory_offset,
872      memory_size,
873      // just set the ptr_offset to the memory_offset, and ptr_size to the memory_size.
874      // we will align the ptr_offset and ptr_size when it should be aligned.
875      ptr_offset: memory_offset,
876      ptr_size: memory_size,
877    }
878  }
879
880  #[inline]
881  unsafe fn clear<A: Allocator>(&self, arena: &A) {
882    unsafe {
883      let ptr = arena.raw_mut_ptr().add(self.ptr_offset as usize);
884      core::ptr::write_bytes(ptr, 0, self.ptr_size as usize);
885    }
886  }
887
888  #[inline]
889  fn align_to<T>(&mut self) {
890    let align_offset = align_offset::<T>(self.memory_offset);
891    self.ptr_offset = align_offset;
892    self.ptr_size = mem::size_of::<T>() as u32;
893  }
894
895  #[inline]
896  fn align_bytes_to<T>(&mut self) {
897    let align_offset = align_offset::<T>(self.memory_offset);
898    self.ptr_offset = align_offset;
899    self.ptr_size = self.memory_offset + self.memory_size - self.ptr_offset;
900  }
901}
902
903mod bytes;
904pub use bytes::*;
905
906mod object;
907pub use object::*;
908
909/// Lock-free allocator allocator can be used in concurrent environments.
910pub mod sync;
911
912/// allocator allocator for single-threaded environments.
913pub mod unsync;
rarena_allocator/lib.rs

rarena_allocator/
lib.rs
