Reduce CString allocations in std as much as possible

Signed-off-by: Alex Saveau <[email protected]>

Author: SUPERCILEX

library/std/src/sys/common/mod.rs

@@ -11,3 +11,7 @@
 #![allow(dead_code)]
 
 pub mod alloc;
+pub mod small_c_string;
+
+#[cfg(test)]
+mod tests;

library/std/src/sys/common/small_c_string.rs

@@ -0,0 +1,58 @@
+use crate::ffi::{CStr, CString};
+use crate::mem::MaybeUninit;
+use crate::path::Path;
+use crate::slice;
+use crate::{io, ptr};
+
+// Make sure to stay under 4096 so the compiler doesn't insert a probe frame:
+// https://docs.rs/compiler_builtins/latest/compiler_builtins/probestack/index.html
+#[cfg(not(target_os = "espidf"))]
+const MAX_STACK_ALLOCATION: usize = 384;
+#[cfg(target_os = "espidf")]
+const MAX_STACK_ALLOCATION: usize = 32;
+
+const NUL_ERR: io::Error =
+    io::const_io_error!(io::ErrorKind::InvalidInput, "file name contained an unexpected NUL byte");
+
+#[inline]
+pub fn run_path_with_cstr<T, F>(path: &Path, f: F) -> io::Result<T>
+where
+    F: FnOnce(&CStr) -> io::Result<T>,
+{
+    run_with_cstr(path.as_os_str().bytes(), f)
+}
+
+#[inline]
+pub fn run_with_cstr<T, F>(bytes: &[u8], f: F) -> io::Result<T>
+where
+    F: FnOnce(&CStr) -> io::Result<T>,
+{
+    if bytes.len() >= MAX_STACK_ALLOCATION {
+        return run_with_cstr_allocating(bytes, f);
+    }
+
+    let mut buf = MaybeUninit::<[u8; MAX_STACK_ALLOCATION]>::uninit();
+    let buf_ptr = buf.as_mut_ptr() as *mut u8;
+
+    unsafe {
+        ptr::copy_nonoverlapping(bytes.as_ptr(), buf_ptr, bytes.len());
+        buf_ptr.add(bytes.len()).write(0);
+    }
+
+    match CStr::from_bytes_with_nul(unsafe { slice::from_raw_parts(buf_ptr, bytes.len() + 1) }) {
+        Ok(s) => f(s),
+        Err(_) => Err(NUL_ERR),
+    }
+}
+
+#[cold]
+#[inline(never)]
+fn run_with_cstr_allocating<T, F>(bytes: &[u8], f: F) -> io::Result<T>
+where
+    F: FnOnce(&CStr) -> io::Result<T>,
+{
+    match CString::new(bytes) {
+        Ok(s) => f(&s),
+        Err(_) => Err(NUL_ERR),
+    }
+}

library/std/src/sys/common/tests.rs

@@ -0,0 +1,66 @@
+use crate::ffi::CString;
+use crate::hint::black_box;
+use crate::path::Path;
+use crate::sys::common::small_c_string::run_path_with_cstr;
+use core::iter::repeat;
+
+#[test]
+fn stack_allocation_works() {
+    let path = Path::new("abc");
+    let result = run_path_with_cstr(path, |p| {
+        assert_eq!(p, &*CString::new(path.as_os_str().bytes()).unwrap());
+        Ok(42)
+    });
+    assert_eq!(result.unwrap(), 42);
+}
+
+#[test]
+fn stack_allocation_fails() {
+    let path = Path::new("ab\0");
+    assert!(run_path_with_cstr::<(), _>(path, |_| unreachable!()).is_err());
+}
+
+#[test]
+fn heap_allocation_works() {
+    let path = repeat("a").take(384).collect::<String>();
+    let path = Path::new(&path);
+    let result = run_path_with_cstr(path, |p| {
+        assert_eq!(p, &*CString::new(path.as_os_str().bytes()).unwrap());
+        Ok(42)
+    });
+    assert_eq!(result.unwrap(), 42);
+}
+
+#[test]
+fn heap_allocation_fails() {
+    let mut path = repeat("a").take(384).collect::<String>();
+    path.push('\0');
+    let path = Path::new(&path);
+    assert!(run_path_with_cstr::<(), _>(path, |_| unreachable!()).is_err());
+}
+
+#[bench]
+fn bench_stack_path_alloc(b: &mut test::Bencher) {
+    let path = repeat("a").take(383).collect::<String>();
+    let p = Path::new(&path);
+    b.iter(|| {
+        run_path_with_cstr(p, |cstr| {
+            black_box(cstr);
+            Ok(())
+        })
+        .unwrap();
+    });
+}
+
+#[bench]
+fn bench_heap_path_alloc(b: &mut test::Bencher) {
+    let path = repeat("a").take(384).collect::<String>();
+    let p = Path::new(&path);
+    b.iter(|| {
+        run_path_with_cstr(p, |cstr| {
+            black_box(cstr);
+            Ok(())
+        })
+        .unwrap();
+    });
+}

library/std/src/sys/hermit/fs.rs

@@ -1,10 +1,12 @@
+use crate::convert::TryFrom;
 use crate::ffi::{CStr, CString, OsString};
 use crate::fmt;
 use crate::hash::{Hash, Hasher};
 use crate::io::{self, Error, ErrorKind};
 use crate::io::{BorrowedCursor, IoSlice, IoSliceMut, SeekFrom};
 use crate::os::unix::ffi::OsStrExt;
 use crate::path::{Path, PathBuf};
+use crate::sys::common::small_c_string::run_path_with_cstr;
 use crate::sys::cvt;
 use crate::sys::hermit::abi;
 use crate::sys::hermit::abi::{O_APPEND, O_CREAT, O_EXCL, O_RDONLY, O_RDWR, O_TRUNC, O_WRONLY};
@@ -15,10 +17,6 @@ use crate::sys::unsupported;
 pub use crate::sys_common::fs::{copy, try_exists};
 //pub use crate::sys_common::fs::remove_dir_all;
 
-fn cstr(path: &Path) -> io::Result<CString> {
-    Ok(CString::new(path.as_os_str().as_bytes())?)
-}
-
 #[derive(Debug)]
 pub struct File(FileDesc);
 
@@ -272,8 +270,7 @@ impl OpenOptions {
 
 impl File {
     pub fn open(path: &Path, opts: &OpenOptions) -> io::Result<File> {
-        let path = cstr(path)?;
-        File::open_c(&path, opts)
+        run_path_with_cstr(path, |path| File::open_c(&path, opts))
     }
 
     pub fn open_c(path: &CStr, opts: &OpenOptions) -> io::Result<File> {
@@ -373,9 +370,7 @@ pub fn readdir(_p: &Path) -> io::Result<ReadDir> {
 }
 
 pub fn unlink(path: &Path) -> io::Result<()> {
-    let name = cstr(path)?;
-    let _ = unsafe { cvt(abi::unlink(name.as_ptr()))? };
-    Ok(())
+    run_path_with_cstr(path, |path| cvt(unsafe { abi::unlink(path.as_ptr()) }).map(|_| ()))
 }
 
 pub fn rename(_old: &Path, _new: &Path) -> io::Result<()> {

library/std/src/sys/mod.rs

@@ -22,7 +22,7 @@
 
 #![allow(missing_debug_implementations)]
 
-mod common;
+pub mod common;
 
 cfg_if::cfg_if! {
     if #[cfg(unix)] {

library/std/src/sys/solid/os.rs

@@ -1,4 +1,5 @@
 use super::unsupported;
+use crate::convert::TryFrom;
 use crate::error::Error as StdError;
 use crate::ffi::{CStr, CString, OsStr, OsString};
 use crate::fmt;
@@ -9,6 +10,7 @@ use crate::os::{
 };
 use crate::path::{self, PathBuf};
 use crate::sync::RwLock;
+use crate::sys::common::small_c_string::run_with_cstr;
 use crate::vec;
 
 use super::{error, itron, memchr};
@@ -139,35 +141,33 @@ pub fn env() -> Env {
 pub fn getenv(k: &OsStr) -> Option<OsString> {
     // environment variables with a nul byte can't be set, so their value is
     // always None as well
-    let k = CString::new(k.as_bytes()).ok()?;
-    unsafe {
+    let s = run_with_cstr(k.as_bytes(), |k| {
         let _guard = ENV_LOCK.read();
-        let s = libc::getenv(k.as_ptr()) as *const libc::c_char;
-        if s.is_null() {
-            None
-        } else {
-            Some(OsStringExt::from_vec(CStr::from_ptr(s).to_bytes().to_vec()))
-        }
+        Ok(unsafe { libc::getenv(k.as_ptr()) } as *const libc::c_char)
+    })
+    .ok()?;
+
+    if s.is_null() {
+        None
+    } else {
+        Some(OsStringExt::from_vec(unsafe { CStr::from_ptr(s) }.to_bytes().to_vec()))
     }
 }
 
 pub fn setenv(k: &OsStr, v: &OsStr) -> io::Result<()> {
-    let k = CString::new(k.as_bytes())?;
-    let v = CString::new(v.as_bytes())?;
-
-    unsafe {
-        let _guard = ENV_LOCK.write();
-        cvt_env(libc::setenv(k.as_ptr(), v.as_ptr(), 1)).map(drop)
-    }
+    run_with_cstr(k.as_bytes(), |k| {
+        run_with_cstr(v.as_bytes(), |v| {
+            let _guard = ENV_LOCK.write();
+            cvt_env(unsafe { libc::setenv(k.as_ptr(), v.as_ptr(), 1) }).map(drop)
+        })
+    })
 }
 
 pub fn unsetenv(n: &OsStr) -> io::Result<()> {
-    let nbuf = CString::new(n.as_bytes())?;
-
-    unsafe {
+    run_with_cstr(n.as_bytes(), |nbuf| {
         let _guard = ENV_LOCK.write();
-        cvt_env(libc::unsetenv(nbuf.as_ptr())).map(drop)
-    }
+        cvt_env(unsafe { libc::unsetenv(nbuf.as_ptr()) }).map(drop)
+    })
 }
 
 /// In kmclib, `setenv` and `unsetenv` don't always set `errno`, so this