diff --git a/library/core/src/slice/memchr.rs b/library/core/src/slice/memchr.rs index 1e1053583a617..68826ecac31f3 100644 --- a/library/core/src/slice/memchr.rs +++ b/library/core/src/slice/memchr.rs @@ -25,10 +25,20 @@ const fn contains_zero_byte(x: usize) -> bool { pub const fn memchr(x: u8, text: &[u8]) -> Option { // Fast path for small slices. if text.len() < 2 * USIZE_BYTES { - return memchr_naive(x, text); + let result = memchr_naive(x, text); + if let Some(index) = result { + // SAFETY: `memchr_naive` only returns an index from within `text`. + unsafe { crate::hint::assert_unchecked(index < text.len()) }; + } + return result; } - memchr_aligned(x, text) + let result = memchr_aligned(x, text); + if let Some(index) = result { + // SAFETY: `memchr_aligned` only returns an index from within `text`. + unsafe { crate::hint::assert_unchecked(index < text.len()) }; + } + result } #[inline] @@ -107,8 +117,18 @@ const fn memchr_aligned(x: u8, text: &[u8]) -> Option { } /// Returns the last index matching the byte `x` in `text`. +#[inline] #[must_use] pub fn memrchr(x: u8, text: &[u8]) -> Option { + let result = memrchr_aligned(x, text); + if let Some(index) = result { + // SAFETY: `memrchr_aligned` only returns the index of a matching byte in `text`. + unsafe { crate::hint::assert_unchecked(index < text.len()) }; + } + result +} + +fn memrchr_aligned(x: u8, text: &[u8]) -> Option { // Scan for a single byte value by reading two `usize` words at a time. // // Split `text` in three parts: diff --git a/library/coretests/tests/slice.rs b/library/coretests/tests/slice.rs index a4db7304fff90..b05f54d4df0a2 100644 --- a/library/coretests/tests/slice.rs +++ b/library/coretests/tests/slice.rs @@ -1781,6 +1781,17 @@ pub mod memchr { assert_eq!(None, memchr(b'a', b"xyz")); } + #[test] + fn each_alignment() { + let mut data = [1u8; 64]; + let needle = 2; + let pos = 40; + data[pos] = needle; + for start in 0..16 { + assert_eq!(Some(pos - start), memchr(needle, &data[start..])); + } + } + #[test] fn matches_one_reversed() { assert_eq!(Some(0), memrchr(b'a', b"a")); diff --git a/tests/codegen-llvm/lib-optimizations/memchr-result.rs b/tests/codegen-llvm/lib-optimizations/memchr-result.rs new file mode 100644 index 0000000000000..77abc33adde83 --- /dev/null +++ b/tests/codegen-llvm/lib-optimizations/memchr-result.rs @@ -0,0 +1,26 @@ +// Ensure `memchr` communicates that a returned index is in bounds. +//@ compile-flags: -Copt-level=3 -Zinline-mir=false +//@ only-x86_64 + +#![crate_type = "lib"] +#![feature(slice_internals)] + +extern crate core; + +use core::slice::memchr::memrchr; + +// CHECK-LABEL: @find_char +#[no_mangle] +pub fn find_char(haystack: &str, needle: char) -> Option { + // CHECK-NOT: phi { i64, i64 } + // CHECK: ret { i64, i64 } + haystack.find(needle) +} + +// CHECK-LABEL: @rfind_byte +#[no_mangle] +pub fn rfind_byte(haystack: &[u8], needle: u8) -> Option { + // CHECK-NOT: panic_bounds_check + // CHECK: ret { i1, i8 } + memrchr(needle, haystack).map(|index| haystack[index]) +}