tokio/util/
memchr.rs

1//! Search for a byte in a byte array using libc.
2//!
3//! When nothing pulls in libc, then just use a trivial implementation. Note
4//! that we only depend on libc on unix.
5
6#[cfg(not(all(unix, feature = "libc")))]
7fn memchr_inner(needle: u8, haystack: &[u8]) -> Option<usize> {
8    haystack.iter().position(|val| needle == *val)
9}
10
11#[cfg(all(unix, feature = "libc"))]
12fn memchr_inner(needle: u8, haystack: &[u8]) -> Option<usize> {
13    let start = haystack.as_ptr();
14
15    // SAFETY: `start` is valid for `haystack.len()` bytes.
16    let ptr = (unsafe { libc::memchr(start.cast(), needle as _, haystack.len()) })
17        .cast::<u8>()
18        .cast_const();
19
20    if ptr.is_null() {
21        None
22    } else {
23        // SAFETY: `ptr` will always be in bounds, since libc guarantees that the ptr will either
24        //          be to an element inside the array or the ptr will be null
25        //          since the ptr is in bounds the offset must also always be non null
26        //          and there can't be more than isize::MAX elements inside an array
27        //          as rust guarantees that the maximum number of bytes a allocation
28        //          may occupy is isize::MAX
29        unsafe {
30            // TODO(MSRV 1.87): When bumping MSRV, switch to `ptr.byte_offset_from_unsigned(start)`.
31            Some(usize::try_from(ptr.offset_from(start)).unwrap_unchecked())
32        }
33    }
34}
35
36pub(crate) fn memchr(needle: u8, haystack: &[u8]) -> Option<usize> {
37    let index = memchr_inner(needle, haystack)?;
38
39    // SAFETY: `memchr_inner` returns Some(index) and in that case index must point to an element in haystack
40    //         or `memchr_inner` None which is guarded by the `?` operator above
41    //         therefore the index must **always** point to an element in the array
42    //         and so this indexing operation is safe
43    // TODO(MSRV 1.81): When bumping MSRV, switch to `std::hint::assert_unchecked(haystack.get(..=index).is_some());`
44    unsafe {
45        if haystack.get(..=index).is_none() {
46            std::hint::unreachable_unchecked()
47        }
48    }
49
50    Some(index)
51}
52
53#[cfg(test)]
54mod tests {
55    use super::memchr;
56
57    #[test]
58    fn memchr_test() {
59        let haystack = b"123abc456\0\xffabc\n";
60
61        assert_eq!(memchr(b'1', haystack), Some(0));
62        assert_eq!(memchr(b'2', haystack), Some(1));
63        assert_eq!(memchr(b'3', haystack), Some(2));
64        assert_eq!(memchr(b'4', haystack), Some(6));
65        assert_eq!(memchr(b'5', haystack), Some(7));
66        assert_eq!(memchr(b'6', haystack), Some(8));
67        assert_eq!(memchr(b'7', haystack), None);
68        assert_eq!(memchr(b'a', haystack), Some(3));
69        assert_eq!(memchr(b'b', haystack), Some(4));
70        assert_eq!(memchr(b'c', haystack), Some(5));
71        assert_eq!(memchr(b'd', haystack), None);
72        assert_eq!(memchr(b'A', haystack), None);
73        assert_eq!(memchr(0, haystack), Some(9));
74        assert_eq!(memchr(0xff, haystack), Some(10));
75        assert_eq!(memchr(0xfe, haystack), None);
76        assert_eq!(memchr(1, haystack), None);
77        assert_eq!(memchr(b'\n', haystack), Some(14));
78        assert_eq!(memchr(b'\r', haystack), None);
79    }
80
81    #[test]
82    fn memchr_all() {
83        let mut arr = Vec::new();
84        for b in 0..=255 {
85            arr.push(b);
86        }
87        for b in 0..=255 {
88            assert_eq!(memchr(b, &arr), Some(b as usize));
89        }
90        arr.reverse();
91        for b in 0..=255 {
92            assert_eq!(memchr(b, &arr), Some(255 - b as usize));
93        }
94    }
95
96    #[test]
97    fn memchr_empty() {
98        for b in 0..=255 {
99            assert_eq!(memchr(b, b""), None);
100        }
101    }
102}
tokio/util/memchr.rs

tokio/util/
memchr.rs
