Add Lua::create_buffer_with_capacity method

This allow creating a preallocated buffer with specified size initialized to zero.

Author: khvzak

src/state.rs

@@ -1150,7 +1150,7 @@ impl Lua {
         }
     }
 
-    /// Create and return an interned Lua string.
+    /// Creates and returns an interned Lua string.
     ///
     /// Lua strings can be arbitrary `[u8]` data including embedded nulls, so in addition to `&str`
     /// and `&String`, you can also pass plain `&[u8]` here.
@@ -1159,27 +1159,32 @@ impl Lua {
         unsafe { self.lock().create_string(s) }
     }
 
-    /// Create and return a Luau [buffer] object from a byte slice of data.
+    /// Creates and returns a Luau [buffer] object from a byte slice of data.
     ///
     /// [buffer]: https://luau.org/library#buffer-library
     #[cfg(any(feature = "luau", doc))]
     #[cfg_attr(docsrs, doc(cfg(feature = "luau")))]
-    pub fn create_buffer(&self, buf: impl AsRef<[u8]>) -> Result<Buffer> {
+    pub fn create_buffer(&self, data: impl AsRef<[u8]>) -> Result<Buffer> {
         let lua = self.lock();
-        let state = lua.state();
+        let data = data.as_ref();
         unsafe {
-            if lua.unlikely_memory_error() {
-                crate::util::push_buffer(state, buf.as_ref(), false)?;
-                return Ok(Buffer(lua.pop_ref()));
-            }
-
-            let _sg = StackGuard::new(state);
-            check_stack(state, 3)?;
-            crate::util::push_buffer(state, buf.as_ref(), true)?;
-            Ok(Buffer(lua.pop_ref()))
+            let (ptr, buffer) = lua.create_buffer_with_capacity(data.len())?;
+            ptr.copy_from_nonoverlapping(data.as_ptr(), data.len());
+            Ok(buffer)
         }
     }
 
+    /// Creates and returns a Luau [buffer] object with the specified size.
+    ///
+    /// Size limit is 1GB. All bytes will be initialized to zero.
+    ///
+    /// [buffer]: https://luau.org/library#buffer-library
+    #[cfg(any(feature = "luau", doc))]
+    #[cfg_attr(docsrs, doc(cfg(feature = "luau")))]
+    pub fn create_buffer_with_capacity(&self, size: usize) -> Result<Buffer> {
+        unsafe { Ok(self.lock().create_buffer_with_capacity(size)?.1) }
+    }
+
     /// Creates and returns a new empty table.
     #[inline]
     pub fn create_table(&self) -> Result<Table> {

src/state/raw.rs

@@ -523,6 +523,20 @@ impl RawLua {
         Ok(String(self.pop_ref()))
     }
 
+    #[cfg(feature = "luau")]
+    pub(crate) unsafe fn create_buffer_with_capacity(&self, size: usize) -> Result<(*mut u8, crate::Buffer)> {
+        let state = self.state();
+        if self.unlikely_memory_error() {
+            let ptr = crate::util::push_buffer(state, size, false)?;
+            return Ok((ptr, crate::Buffer(self.pop_ref())));
+        }
+
+        let _sg = StackGuard::new(state);
+        check_stack(state, 3)?;
+        let ptr = crate::util::push_buffer(state, size, true)?;
+        Ok((ptr, crate::Buffer(self.pop_ref())))
+    }
+
     /// See [`Lua::create_table_with_capacity`]
     pub(crate) unsafe fn create_table_with_capacity(&self, narr: usize, nrec: usize) -> Result<Table> {
         let state = self.state();

src/util/mod.rs

@@ -101,15 +101,13 @@ pub(crate) unsafe fn push_string(state: *mut ffi::lua_State, s: &[u8], protect:
 // Uses 3 stack spaces (when protect), does not call checkstack.
 #[cfg(feature = "luau")]
 #[inline(always)]
-pub(crate) unsafe fn push_buffer(state: *mut ffi::lua_State, b: &[u8], protect: bool) -> Result<()> {
-    let data = if protect {
-        protect_lua!(state, 0, 1, |state| ffi::lua_newbuffer(state, b.len()))?
+pub(crate) unsafe fn push_buffer(state: *mut ffi::lua_State, size: usize, protect: bool) -> Result<*mut u8> {
+    let data = if protect || size > const { 1024 * 1024 * 1024 } {
+        protect_lua!(state, 0, 1, |state| ffi::lua_newbuffer(state, size))?
     } else {
-        ffi::lua_newbuffer(state, b.len())
+        ffi::lua_newbuffer(state, size)
     };
-    let buf = slice::from_raw_parts_mut(data as *mut u8, b.len());
-    buf.copy_from_slice(b);
-    Ok(())
+    Ok(data as *mut u8)
 }
 
 // Uses 3 stack spaces, does not call checkstack.

tests/buffer.rs

@@ -55,3 +55,14 @@ fn test_buffer_out_of_bounds_write() {
     let buf = lua.create_buffer(b"hello, world!").unwrap();
     buf.write_bytes(14, b"!!");
 }
+
+#[test]
+fn create_large_buffer() {
+    let lua = Lua::new();
+    let err = lua.create_buffer_with_capacity(1_073_741_824 + 1).unwrap_err(); // 1GB
+    assert!(err.to_string().contains("memory allocation error"));
+
+    // Normal buffer is okay
+    let buf = lua.create_buffer_with_capacity(1024 * 1024).unwrap();
+    assert_eq!(buf.len(), 1024 * 1024);
+}