yyjson_rs/
allocator.rs

1use std::alloc::{self, Layout};
2use std::ffi::c_void;
3use std::ptr;
4
5use yyjson_sys as ffi;
6
7#[repr(transparent)]
8pub struct YyjsonAllocator<'a> {
9    pub(crate) p: *mut ffi::yyjson_alc,
10    // just in case the Allocator is an embedded field within a parent/holder it
11    // should not outlive the parent
12    pub(crate) _lifetime: std::marker::PhantomData<&'a ()>,
13}
14
15pub trait YyjsonAllocProvider {
16    /// Creates and returns a `YyjsonAllocator`.
17    ///
18    /// # Lifetimes
19    /// - `'a` ensures that the returned `YyjsonAllocator` cannot outlive the `self` reference,
20    ///   which prevents use-after-free issues if `self` owns or manages the underlying memory.
21    fn get_allocator<'a>(&'a self) -> YyjsonAllocator<'a>;
22}
23
24// =============================================================================
25//                  DEFAULT ALLOCATOR
26// =============================================================================
27
28pub struct BasicAllocProvider {}
29impl Default for BasicAllocProvider {
30    fn default() -> Self {
31        return BasicAllocProvider {};
32    }
33}
34impl YyjsonAllocProvider for BasicAllocProvider {
35    fn get_allocator<'a>(&'a self) -> YyjsonAllocator<'a> {
36        YyjsonAllocator {
37            p: ptr::null_mut(),
38            _lifetime: std::marker::PhantomData,
39        }
40    }
41}
42
43// =============================================================================
44//                  SINGLE BUF ALLOCATOR
45// =============================================================================
46
47// Uses a fixed length pre-allocated memory
48pub struct PoolAllocProvider {
49    alc: *mut ffi::yyjson_alc,
50    buf: *mut u8,
51    layout: Layout,
52}
53
54impl Drop for PoolAllocProvider {
55    fn drop(&mut self) {
56        // SAFETY: SingelBufAllocator can only be created if buf is non-null and
57        // points to valid memory and layout is valid
58        unsafe { alloc::dealloc(self.buf, self.layout) };
59
60        // SAFETY:: alc is initialized from a pointer from the boxed value.
61        let _ = unsafe { Box::from_raw(self.alc) };
62    }
63}
64
65impl YyjsonAllocProvider for PoolAllocProvider {
66    #[inline(always)]
67    fn get_allocator<'a>(&'a self) -> YyjsonAllocator<'a> {
68        YyjsonAllocator {
69            p: self.alc,
70            _lifetime: std::marker::PhantomData,
71        }
72    }
73}
74
75impl PoolAllocProvider {
76    pub fn new(
77        expected_json_size: usize,
78        options: Option<&crate::ReadOptions>,
79    ) -> Result<Self, &'static str> {
80        let flag = if let Some(v) = options {
81            v.to_read_flag()
82        } else {
83            ffi::YYJSON_READ_NOFLAG
84        };
85        // SAFETY: on error (e.g.overflow), this function returns 0. When
86        // creating layout below, if buf_size is 0, it will return an error to
87        // the caller
88        let buf_size = unsafe { ffi::yyjson_read_max_memory_usage(expected_json_size, flag) };
89        let layout = match Layout::from_size_align(buf_size, 64) {
90            Ok(l) => l,
91            Err(_) => {
92                return Err("align should be non-zero and a power of 2, size should not overflow")
93            }
94        };
95        let buf: *mut u8 = unsafe { alloc::alloc(layout) };
96        if buf.is_null() {
97            return Err("allocator returned null");
98        }
99        let alc: *mut ffi::yyjson_alc = {
100            let b = Box::new(ffi::yyjson_alc {
101                malloc: None,
102                realloc: None,
103                free: None,
104                ctx: ptr::null_mut(),
105            });
106            Box::into_raw(b)
107        };
108
109        // this has to be created here so that if `yyjson_alc_pool_init` fails
110        // then both alc and buf will be dropped
111        let sb = PoolAllocProvider { alc, buf, layout };
112
113        // init pool
114        // SAFETY: alc is non-null since it's retrieved from a Boxed value and
115        // initialized with valid values, buf is non-null and of given
116        // `buf_size`
117        let ok: bool = unsafe { ffi::yyjson_alc_pool_init(alc, buf.cast::<c_void>(), buf_size) };
118        if !ok {
119            return Err("Failed to initialized pool allocator, buf or size invalid");
120        }
121
122        Ok(sb)
123    }
124}
125
126// =============================================================================
127//                  DYN ALLOCATOR
128// =============================================================================
129
130#[repr(transparent)]
131pub struct DynamicAllocProvider {
132    alc: *mut ffi::yyjson_alc,
133}
134
135impl Default for DynamicAllocProvider {
136    fn default() -> Self {
137        let alc: *mut ffi::yyjson_alc = unsafe { ffi::yyjson_alc_dyn_new() };
138        Self { alc }
139    }
140}
141
142impl YyjsonAllocProvider for DynamicAllocProvider {
143    #[inline(always)]
144    fn get_allocator<'a>(&'a self) -> YyjsonAllocator<'a> {
145        YyjsonAllocator {
146            p: self.alc,
147            _lifetime: std::marker::PhantomData,
148        }
149    }
150}
151
152impl Drop for DynamicAllocProvider {
153    fn drop(&mut self) {
154        // SAFETY: by construction, either DynamicAllocator.alc is null in which
155        // case `yyjson_alc_dyn_free` will not error out or is non-null in which
156        // case it points to a valid yyjson_alc that will be freed appropriately
157        unsafe { ffi::yyjson_alc_dyn_free(self.alc) }
158    }
159}