Skip to content

Commit

Permalink
io: use memchr from libc
Browse files Browse the repository at this point in the history
  • Loading branch information
Darksonn committed Aug 29, 2023
1 parent edd172c commit 9ddfcf2
Show file tree
Hide file tree
Showing 4 changed files with 79 additions and 2 deletions.
3 changes: 1 addition & 2 deletions tokio/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ full = [
]

fs = []
io-util = ["memchr", "bytes"]
io-util = ["bytes"]
# stdin, stdout, stderr
io-std = []
macros = ["tokio-macros"]
Expand Down Expand Up @@ -108,7 +108,6 @@ pin-project-lite = "0.2.0"

# Everything else is optional...
bytes = { version = "1.0.0", optional = true }
memchr = { version = "2.2", optional = true }
mio = { version = "0.8.1", optional = true, default-features = false }
socket2 = { version = "0.4.4", optional = true, features = [ "all" ] }
num_cpus = { version = "1.8.0", optional = true }
Expand Down
1 change: 1 addition & 0 deletions tokio/src/io/util/read_until.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use crate::io::AsyncBufRead;
use crate::util::memchr;

use pin_project_lite::pin_project;
use std::future::Future;
Expand Down
74 changes: 74 additions & 0 deletions tokio/src/util/memchr.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
//! Search for a byte in a byte array using libc.
//!
//! When nothing pulls in libc, then just use a trivial implementation. Note
//! that we only depend on libc on unix.

#[cfg(not(all(unix, feature = "libc")))]
pub(crate) fn memchr(needle: u8, haystack: &[u8]) -> Option<usize> {
haystack.iter().position(|val| needle == *val)
}

#[cfg(all(unix, feature = "libc"))]
pub(crate) fn memchr(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()) };

if ptr.is_null() {
None
} else {
Some(ptr as usize - start as usize)
}
}

#[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);
assert_eq!(memchr(b'a', haystack), Some(3));
assert_eq!(memchr(b'b', haystack), Some(4));
assert_eq!(memchr(b'c', haystack), Some(5));
assert_eq!(memchr(b'd', haystack), None);
assert_eq!(memchr(b'A', haystack), None);
assert_eq!(memchr(0, haystack), Some(9));
assert_eq!(memchr(0xff, haystack), Some(10));
assert_eq!(memchr(0xfe, haystack), None);
assert_eq!(memchr(1, haystack), None);
assert_eq!(memchr(b'\n', haystack), Some(14));
assert_eq!(memchr(b'\r', haystack), None);
}

#[test]
fn memchr_all() {
let mut arr = Vec::new();
for b in 0..=255 {
arr.push(b);
}
for b in 0..=255 {
assert_eq!(memchr(b, &arr), Some(b as usize));
}
arr.reverse();
for b in 0..=255 {
assert_eq!(memchr(b, &arr), Some(255 - b as usize));
}
}

#[test]
fn memchr_empty() {
for b in 0..=255 {
assert_eq!(memchr(b, b""), None);
}
}
}
3 changes: 3 additions & 0 deletions tokio/src/util/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -90,3 +90,6 @@ pub use self::rand::thread_rng_n;
all(unix, feature = "signal")
))]
pub(crate) mod error;

#[cfg(feature = "io-util")]
pub(crate) mod memchr;

0 comments on commit 9ddfcf2

Please sign in to comment.