Implement OsStr::split_at

library/core/src/wtf8.rs

@@ -454,25 +454,50 @@ impl Wtf8 {
     #[track_caller]
     #[inline]
     pub fn check_utf8_boundary(&self, index: usize) {
+        let Err(err) = self.try_check_utf8_boundary(index) else { return };
+        match err {
+            Utf8BoundaryError::NotABoundary => {
+                panic!("byte index {index} is not a codepoint boundary")
+            }
+            Utf8BoundaryError::OutOfBounds => panic!("byte index {index} is out of bounds"),
+            Utf8BoundaryError::BetweenSurrogates => {
+                panic!("byte index {index} lies between surrogate codepoints")
+            }
+        }
+    }
+
+    #[track_caller]
+    #[inline]
+    pub fn try_check_utf8_boundary(&self, index: usize) -> Result<(), Utf8BoundaryError> {
         if index == 0 {
-            return;
+            return Ok(());
         }
         match self.bytes.get(index) {
             Some(0xED) => (), // Might be a surrogate
-            Some(&b) if (b as i8) >= -0x40 => return,
-            Some(_) => panic!("byte index {index} is not a codepoint boundary"),
-            None if index == self.len() => return,
-            None => panic!("byte index {index} is out of bounds"),
+            Some(&b) if (b as i8) >= -0x40 => return Ok(()),
+            Some(_) => return Err(Utf8BoundaryError::NotABoundary),
+            None if index == self.len() => return Ok(()),
+            None => return Err(Utf8BoundaryError::OutOfBounds),
         }
         if self.bytes[index + 1] >= 0xA0 {
             // There's a surrogate after index. Now check before index.
             if index >= 3 && self.bytes[index - 3] == 0xED && self.bytes[index - 2] >= 0xA0 {
-                panic!("byte index {index} lies between surrogate codepoints");
+                return Err(Utf8BoundaryError::BetweenSurrogates);
             }
         }
+        Ok(())
     }
 }
 
+// This error type is only used temporarily to provide better panic messages
+// It does not implement Error.
+#[derive(Debug)]
+pub enum Utf8BoundaryError {
+    NotABoundary,
+    OutOfBounds,
+    BetweenSurrogates,
+}
+
 /// Copied from core::str::raw::slice_unchecked
 #[inline]
 unsafe fn slice_unchecked(s: &Wtf8, begin: usize, end: usize) -> &Wtf8 {

library/std/src/ffi/os_str.rs

@@ -1055,6 +1055,61 @@ impl OsStr {
         OsString { inner: Buf::from_box(boxed) }
     }
 
+    /// Divides one string slice into two at an index.
+    ///
+    /// The two slices returned go from the start of the string slice to `mid`, and from `mid` to the end of the string slice.
+    ///
+    /// The argument, `mid`, should be a byte offset from the start of the string.
+    /// It must also be on a valid `OsStr` boundary.
+    /// See [`split_at_checked`][Self::split_at_checked] for the definition of a valid boundary.
+    ///
+    /// Panics
+    ///
+    /// Panics if `mid` is not on a valid boundary, or if it is past the end of the last code point of the string slice.
+    /// For a non-panicking alternative see [`split_at_checked`][Self::split_at_checked].
+    #[unstable(feature = "os_str_split_at", issue = "none")]
+    pub fn split_at(&self, mid: usize) -> (&OsStr, &OsStr) {
+        self.inner.check_public_boundary(mid);
+
+        // SAFETY: we've checked it's in bounds and a valid boundary
+        unsafe { self.split_at_unchecked(mid) }
+    }
+
+    /// Divides one string slice into two at an index.
+    ///
+    /// The two slices returned go from the start of the string slice to `mid`, and from `mid` to the end of the string slice.
+    ///
+    /// The argument, `mid`, should be a valid byte offset from the start of the string.
+    /// It must also be on a valid `OsStr` boundary.
+    /// The method returns `None` if that’s not the case.
+    /// A valid `OsStr` boundary is one of:
+    /// - The start of the string
+    /// - The end of the string
+    /// - The start of a valid non-empty UTF-8 substring
+    /// - Immediately follows a valid non-empty UTF-8 substring
+    #[unstable(feature = "os_str_split_at", issue = "none")]
+    pub fn split_at_checked(&self, mid: usize) -> Option<(&OsStr, &OsStr)> {
+        self.inner.try_check_public_boundary(mid)?;
+
+        // SAFETY: we've checked it's in bounds and a valid boundary
+        unsafe { Some(self.split_at_unchecked(mid)) }
+    }
+
+    /// Splits an `OsStr` without checking if `mid` is a valid boundary.
+    /// You should use `split_at` or `split_at_checked` instead.
+    ///
+    /// # Safety
+    ///
+    /// Any caller must ensure `mid` is within bounds and lies on
+    /// a valid `OsStr` boundary for the platform.
+    unsafe fn split_at_unchecked(&self, mid: usize) -> (&OsStr, &OsStr) {
+        // SAFETY: it's up to the caller to ensure this is safe.
+        unsafe {
+            let (first, second) = self.as_encoded_bytes().split_at_unchecked(mid);
+            (Self::from_encoded_bytes_unchecked(first), Self::from_encoded_bytes_unchecked(second))
+        }
+    }
+
     /// Converts an OS string slice to a byte slice.  To convert the byte slice back into an OS
     /// string slice, use the [`OsStr::from_encoded_bytes_unchecked`] function.
     ///

library/std/src/ffi/os_str/tests.rs

@@ -289,6 +289,108 @@ fn slice_surrogate_edge() {
     assert_eq!(post_crab.slice_encoded_bytes(4..), surrogate);
 }
 
+#[test]
+fn os_str_slice_at() {
+    #[track_caller]
+    fn slice_at_ok(input: &OsStr, index: usize, expected: (&str, &str)) {
+        let expected = (OsStr::new(expected.0), OsStr::new(expected.1));
+        assert_eq!(input.split_at(index), expected);
+        assert_eq!(input.split_at_checked(index), Some(expected));
+    }
+
+    let os_str = OsStr::new("123გ🦀4");
+    slice_at_ok(os_str, 0, ("", "123გ🦀4"));
+    slice_at_ok(os_str, 1, ("1", "23გ🦀4"));
+    slice_at_ok(os_str, 2, ("12", "3გ🦀4"));
+    slice_at_ok(os_str, 3, ("123", "გ🦀4"));
+    slice_at_ok(os_str, 6, ("123გ", "🦀4"));
+    slice_at_ok(os_str, 10, ("123გ🦀", "4"));
+    slice_at_ok(os_str, 11, ("123გ🦀4", ""));
+
+    // Invalid boundaries should fail.
+    assert!(os_str.split_at_checked(4).is_none());
+    assert!(os_str.split_at_checked(5).is_none());
+    assert!(os_str.split_at_checked(7).is_none());
+    assert!(os_str.split_at_checked(8).is_none());
+    assert!(os_str.split_at_checked(9).is_none());
+    // Out of bounds
+    assert!(os_str.split_at_checked(12).is_none());
+}
+
+#[test]
+#[should_panic]
+fn os_str_slice_at_out_of_bounds() {
+    let crab = OsStr::new("🦀");
+    let _ = crab.split_at(5);
+}
+
+#[test]
+#[should_panic]
+fn os_str_slice_at_mid_char() {
+    let crab = OsStr::new("🦀");
+    let _ = crab.split_at(2);
+}
+
+#[cfg(unix)]
+#[test]
+fn os_str_slice_at_unix() {
+    use crate::os::unix::ffi::OsStrExt;
+
+    let broken_utf8 = OsStr::from_bytes(&"🦀".as_bytes()[..3]);
+    let invalid = OsStr::from_bytes(b"\xFF");
+
+    // Check that broken UTF-8 isn't treated as if it's valid.
+    let mut os_string = invalid.to_os_string();
+    os_string.push(broken_utf8);
+    assert_eq!(os_string.split_at_checked(1), None);
+
+    // We should be able to split on ascii with invalid UTF-8 between
+    let os_string = OsStr::from_bytes(b"a\xFFa");
+    assert_eq!(os_string.split_at_checked(1), Some(("a".as_ref(), OsStr::from_bytes(b"\xFFa"))));
+    assert_eq!(os_string.split_at_checked(2), Some((OsStr::from_bytes(b"a\xFF"), "a".as_ref(),)));
+
+    let os_string = OsStr::from_bytes(&"abc🦀".as_bytes()[..6]);
+    assert_eq!(
+        os_string.split_at_checked(3),
+        Some(("abc".as_ref(), OsStr::from_bytes(b"\xF0\x9F\xA6")))
+    );
+
+    let mut os_string = invalid.to_os_string();
+    os_string.push("🦀");
+    assert_eq!(os_string.split_at_checked(1), Some((invalid, "🦀".as_ref())));
+}
+
+#[test]
+#[cfg(windows)]
+fn os_str_slice_at_windows() {
+    use crate::os::windows::ffi::OsStringExt;
+
+    // slicing between unpaired surrogates should not be possible
+    // checking is implemented as a loop so we're agnostic towards
+    // the internal encoding
+    let os_string = OsString::from_wide(&[0xD800, 0xD800]);
+    for i in 1..os_string.len() {
+        assert_eq!(os_string.split_at_checked(i), None);
+    }
+    // For completeness, check that splitting at the start and end still works.
+    assert!(os_string.split_at_checked(0).is_some());
+    assert!(os_string.split_at_checked(os_string.len()).is_some());
+
+    // check that slicing before and after unpaired surrogates work
+    let surrogate = OsString::from_wide(&[0xD800]);
+
+    let mut os_string = surrogate.clone();
+    os_string.push("🦀");
+    assert_eq!(
+        os_string.split_at_checked(surrogate.len()),
+        Some((surrogate.as_ref(), "🦀".as_ref()))
+    );
+
+    let mut os_string = OsString::from("🦀");
+    os_string.push(&surrogate);
+    assert_eq!(os_string.split_at_checked("🦀".len()), Some(("🦀".as_ref(), surrogate.as_ref())));
+}
+
 #[test]
 fn clone_to_uninit() {
     let a = OsStr::new("hello.txt");

library/std/src/sys/os_str/bytes.rs

@@ -238,43 +238,51 @@ impl Slice {
     #[track_caller]
     #[inline]
     pub fn check_public_boundary(&self, index: usize) {
+        if self.try_check_public_boundary(index).is_none() {
+            panic!("byte index {index} is not an OsStr boundary");
+        }
+    }
+
+    #[track_caller]
+    #[inline]
+    pub fn try_check_public_boundary(&self, index: usize) -> Option<()> {
         if index == 0 || index == self.inner.len() {
-            return;
+            return Some(());
         }
         if index < self.inner.len()
             && (self.inner[index - 1].is_ascii() || self.inner[index].is_ascii())
         {
-            return;
+            return Some(());
         }
 
-        slow_path(&self.inner, index);
+        return slow_path(&self.inner, index);
 
         /// We're betting that typical splits will involve an ASCII character.
         ///
         /// Putting the expensive checks in a separate function generates notably
         /// better assembly.
         #[track_caller]
         #[inline(never)]
-        fn slow_path(bytes: &[u8], index: usize) {
-            let (before, after) = bytes.split_at(index);
+        fn slow_path(bytes: &[u8], index: usize) -> Option<()> {
+            let (before, after) = bytes.split_at_checked(index)?;
 
             // UTF-8 takes at most 4 bytes per codepoint, so we don't
             // need to check more than that.
             let after = after.get(..4).unwrap_or(after);
             match str::from_utf8(after) {
-                Ok(_) => return,
-                Err(err) if err.valid_up_to() != 0 => return,
+                Ok(_) => return Some(()),
+                Err(err) if err.valid_up_to() != 0 => return Some(()),
                 Err(_) => (),
             }
 
             for len in 2..=4.min(index) {
                 let before = &before[index - len..];
                 if str::from_utf8(before).is_ok() {
-                    return;
+                    return Some(());
                 }
             }
 
-            panic!("byte index {index} is not an OsStr boundary");
+            None
         }
     }
 

library/std/src/sys/os_str/utf8.rs

@@ -224,10 +224,16 @@ impl Slice {
         Slice::from_str(unsafe { str::from_utf8_unchecked(s) })
     }
 
+    #[track_caller]
+    #[inline]
+    pub fn try_check_public_boundary(&self, index: usize) -> Option<()> {
+        if self.inner.is_char_boundary(index) { Some(()) } else { None }
+    }
+
     #[track_caller]
     #[inline]
     pub fn check_public_boundary(&self, index: usize) {
-        if !self.inner.is_char_boundary(index) {
+        if self.try_check_public_boundary(index).is_none() {
             panic!("byte index {index} is not an OsStr boundary");
         }
     }

library/std/src/sys/os_str/wtf8.rs

@@ -240,6 +240,11 @@ impl Slice {
         unsafe { mem::transmute(Wtf8::from_bytes_unchecked(s)) }
     }
 
+    #[inline]
+    pub fn try_check_public_boundary(&self, index: usize) -> Option<()> {
+        self.inner.try_check_utf8_boundary(index).ok()
+    }
+
     #[track_caller]
     #[inline]
     pub fn check_public_boundary(&self, index: usize) {