Repository: tokio-rs/tokio
File: tokio/src/util/memchr.rs (Rust)

Code before:
fn memchr_inner(needle: u8, haystack: &[u8]) -> Option<usize> {
    haystack.iter().position(|val| needle == *val)
}

#[cfg(all(unix, feature = "libc"))]
fn memchr_inner(needle: u8, haystack: &[u8]) -> Option<usize> {
    let start = haystack.as_ptr();

    // SAFETY: `start` is valid for `haystack.len()` bytes.
    let ptr = (unsafe { libc::memchr(start.cast(), needle as _, haystack.len()) })
        .cast::<u8>()
        .cast_const();

    if ptr.is_null() {
        None
    } else {
        unsafe {
            // TODO(MSRV 1.87): When bumping MSRV, switch to `ptr.byte_offset_from_unsigned(start)`.
            Some(usize::try_from(ptr.offset_from(start)).unwrap_unchecked())
        }
    }
}

pub(crate) fn memchr(needle: u8, haystack: &[u8]) -> Option<usize> {
    let index = memchr_inner(needle, haystack)?;

    // TODO(MSRV 1.81): When bumping MSRV, switch to `std::hint::assert_unchecked(index < haystack.len());`
    unsafe {
        if index >= haystack.len() {
            std::hint::unreachable_unchecked()
        }
    }

    Some(index)
}

#[cfg(test)]
mod tests {
    use super::memchr;

    #[test]
    fn memchr_test() {
        let haystack = b"123abc456\0\xffabc\n";

        assert_eq!(memchr(b'1', haystack), Some(0));
        assert_eq!(memchr(b'2', haystack), Some(1));
        assert_eq!(memchr(b'3', haystack), Some(2));
        assert_eq!(memchr(b'4', haystack), Some(6));
        assert_eq!(memchr(b'5', haystack), Some(7));
        assert_eq!(memchr(b'6', haystack), Some(8));
        assert_eq!(memchr(b'7', haystack), None);

Review comment (suggestion):
New unsafe blocks should wrap only a single unsafe operation and use SAFETY comments.
```suggestion
        // SAFETY: `memchr` returns null or a pointer from same allocation as
        // `start`, and we know it's not null.
        let offset = unsafe { ptr.offset_from(start) } ;
        // TODO(MSRV 1.87): When bumping MSRV, switch to `ptr.byte_offset_from_unsigned(start)`.
        // SAFETY: `memchr` returns a pointer between `start` and `start+len`,
        // so the offset is positive.
        Some(unsafe { usize::try_from(offset).unwrap_unchecked() })
```


Code after:
fn memchr_inner(needle: u8, haystack: &[u8]) -> Option<usize> {
    haystack.iter().position(|val| needle == *val)
}

#[cfg(all(unix, feature = "libc"))]
fn memchr_inner(needle: u8, haystack: &[u8]) -> Option<usize> {
    let start = haystack.as_ptr();

    // SAFETY: `start` is valid for `haystack.len()` bytes.
    let ptr = (unsafe { libc::memchr(start.cast(), needle as _, haystack.len()) })
        .cast::<u8>()
        .cast_const();

    if ptr.is_null() {
        None
    } else {
        // SAFETY: `ptr` will always be in bounds, since libc guarentees that the ptr will either
        //          be to an element inside the array or the ptr will be null
        //          since the ptr is in bounds the offset must also always be non null
        //          and there can't be more than isize::MAX elements inside an array
        //          as rust guarentees that the maximum number of bytes a allocation
        //          may occupy is isize::MAX
        unsafe {
            // TODO(MSRV 1.87): When bumping MSRV, switch to `ptr.byte_offset_from_unsigned(start)`.
            Some(usize::try_from(ptr.offset_from(start)).unwrap_unchecked())
        }
    }
}

pub(crate) fn memchr(needle: u8, haystack: &[u8]) -> Option<usize> {
    let index = memchr_inner(needle, haystack)?;

    // SAFETY: `memchr_inner` returns Some(index) and in that case index must point to an element in haystack
    //         or `memchr_inner` None which is guarded by the `?` operator above
    //         therfore the index must **always** point to an element in the array
    //         and so this indexing operation is safe
    // TODO(MSRV 1.81): When bumping MSRV, switch to `std::hint::assert_unchecked(haystack.get(..=index).is_some());`
    unsafe {
        if haystack.get(..=index).is_none() {
            std::hint::unreachable_unchecked()
        }
    }

    Some(index)
}

#[cfg(test)]
mod tests {
    use super::memchr;

    #[test]