Buffer refactoring changes from e9b745e

andersk · Aug 30, 2023 · 957203f · 957203f
1 parent ee8759c
commit 957203f
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 37 deletions.
diff --git a/src/deserialize/yyjson.rs b/src/deserialize/yyjson.rs
@@ -60,33 +60,59 @@ fn unsafe_yyjson_get_next(val: *mut yyjson_val) -> *mut yyjson_val {
 pub fn deserialize_yyjson(
     data: &'static str,
 ) -> Result<NonNull<pyo3_ffi::PyObject>, DeserializeError<'static>> {
+    let mut err = yyjson_read_err {
+        code: YYJSON_READ_SUCCESS,
+        msg: null(),
+        pos: 0,
+    };
+    let doc = if yyjson_read_max_memory_usage(data.len()) < YYJSON_BUFFER_SIZE {
+        read_doc_with_buffer(data, &mut err)
+    } else {
+        read_doc_default(data, &mut err)
+    };
+    if unlikely!(doc.is_null()) {
+        let msg: Cow<str> = unsafe { std::ffi::CStr::from_ptr(err.msg).to_string_lossy() };
+        Err(DeserializeError::from_yyjson(msg, err.pos as i64, data))
+    } else {
+        let root = yyjson_doc_get_root(doc);
+        let ret = parse_node(root);
+        unsafe { yyjson_doc_free(doc) };
+        Ok(ret)
+    }
+}
+
+fn read_doc_default(data: &'static str, err: &mut yyjson_read_err) -> *mut yyjson_doc {
     unsafe {
-        let allocator = if yyjson_read_max_memory_usage(data.len()) < YYJSON_BUFFER_SIZE {
-            YYJSON_ALLOC.get_or_init(yyjson_init) as *const yyjson_alc
-        } else {
-            null_mut()
-        };
-        let mut err = yyjson_read_err {
-            code: YYJSON_READ_SUCCESS,
-            msg: null(),
-            pos: 0,
-        };
-        let doc: *mut yyjson_doc = yyjson_read_opts(
+        yyjson_read_opts(
             data.as_ptr() as *mut c_char,
             data.len(),
             YYJSON_READ_NOFLAG,
-            allocator,
-            &mut err,
+            null_mut(),
+            err,
+        )
+    }
+}
+
+fn read_doc_with_buffer(data: &'static str, err: &mut yyjson_read_err) -> *mut yyjson_doc {
+    unsafe {
+        let mut allocator = crate::yyjson::yyjson_alc {
+            malloc: None,
+            realloc: None,
+            free: None,
+            ctx: null_mut(),
+        };
+        crate::yyjson::yyjson_alc_pool_init(
+            &mut allocator,
+            YYJSON_ALLOC.get_or_init(yyjson_init).as_ptr() as *mut std::os::raw::c_void,
+            YYJSON_BUFFER_SIZE,
         );
-        if unlikely!(doc.is_null()) {
-            let msg: Cow<str> = std::ffi::CStr::from_ptr(err.msg).to_string_lossy();
-            Err(DeserializeError::from_yyjson(msg, err.pos as i64, data))
-        } else {
-            let root = yyjson_doc_get_root(doc);
-            let ret = parse_node(root);
-            yyjson_doc_free(doc);
-            Ok(ret)
-        }
+        yyjson_read_opts(
+            data.as_ptr() as *mut c_char,
+            data.len(),
+            YYJSON_READ_NOFLAG,
+            std::ptr::addr_of!(allocator),
+            err,
+        )
     }
 }
 

diff --git a/src/typeref.rs b/src/typeref.rs
@@ -88,27 +88,16 @@ pub fn ahash_init() -> Box<ahash::RandomState> {
 pub const YYJSON_BUFFER_SIZE: usize = 1024 * 1024 * 8;
 
 #[cfg(feature = "yyjson")]
-pub static mut YYJSON_ALLOC: OnceBox<crate::yyjson::yyjson_alc> = OnceBox::new();
+pub static mut YYJSON_ALLOC: OnceBox<[u8; YYJSON_BUFFER_SIZE]> = OnceBox::new();
 
 #[cfg(feature = "yyjson")]
-pub fn yyjson_init() -> Box<crate::yyjson::yyjson_alc> {
+pub fn yyjson_init() -> Box<[u8; YYJSON_BUFFER_SIZE]> {
     unsafe {
         let buffer = std::alloc::alloc(std::alloc::Layout::from_size_align_unchecked(
             YYJSON_BUFFER_SIZE,
             64,
-        ));
-        let mut alloc = crate::yyjson::yyjson_alc {
-            malloc: None,
-            realloc: None,
-            free: None,
-            ctx: null_mut(),
-        };
-        crate::yyjson::yyjson_alc_pool_init(
-            &mut alloc,
-            buffer as *mut std::os::raw::c_void,
-            YYJSON_BUFFER_SIZE,
-        );
-        Box::new(alloc)
+        )) as *mut [u8; YYJSON_BUFFER_SIZE];
+        Box::new(*buffer)
     }
 }