core: use the platform's memchr

library/compiler-builtins/compiler-builtins/src/lib.rs

@@ -9,6 +9,7 @@
 #![feature(linkage)]
 #![feature(repr_simd)]
 #![feature(macro_metavar_expr_concat)]
+#![feature(portable_simd)]
 #![feature(rustc_attrs)]
 #![cfg_attr(f16_enabled, feature(f16))]
 #![cfg_attr(f128_enabled, feature(f128))]

library/compiler-builtins/compiler-builtins/src/mem/memchr_impl.rs

@@ -0,0 +1,133 @@
+// Some parts taken from rust-memchr.
+// Copyright 2015 Andrew Gallant, bluss and Nicolas Koch
+
+#![forbid(unsafe_op_in_unsafe_fn)]
+
+use core::ffi::{c_int, c_void};
+use core::ptr;
+
+unsafe fn memchr_naive(
+    mut current: *const u8,
+    needle: u8,
+    mut n: usize,
+) -> *mut c_void {
+    while n > 0 {
+        let byte = unsafe { current.read() };
+        if byte == needle {
+            return current as *mut c_void;
+        }
+
+        current = unsafe { current.add(1) };
+        n -= 1;
+    }
+
+    ptr::null_mut()
+}
+
+type Chunk = cfg_select! {
+    any(
+        all(target_arch = "x86_64", target_feature = "sse2"),
+        all(target_arch = "aarch64", target_feature = "neon"),
+    ) => core::simd::u8x16,
+    _ => [usize; 2],
+};
+
+pub unsafe fn memchr(
+    s: *const c_void,
+    c: c_int,
+    mut n: usize,
+) -> *mut c_void {
+    let mut current = s.cast::<u8>();
+    let needle = c as u8;
+
+    if n < size_of::<Chunk>() {
+        return unsafe { memchr_naive(current, needle, n) };
+    }
+
+    // Advance until the pointer is aligned to a chunk.
+    // Since the size of a chunk must be larger than its alignment and we
+    // only reach this point if `n` is larger or equal to the size of a chunk,
+    // this doesn't need bound checks.
+    while !current.cast::<Chunk>().is_aligned() {
+        let byte = unsafe { current.read() };
+        if byte == needle {
+            return current as *mut c_void;
+        }
+
+        current = unsafe { current.add(1) };
+        n -= 1;
+    }
+
+    cfg_select! {
+        // These targets have efficient SIMD acceleration.
+        any(
+            all(target_arch = "x86_64", target_feature = "sse2"),
+            all(target_arch = "aarch64", target_feature = "neon"),
+        ) => {
+            use core::simd::cmp::SimdPartialEq;
+
+            let mut current = current.cast::<Chunk>();
+            let chunk_needle = Chunk::splat(needle);
+            while n >= size_of::<Chunk>() {
+                let chunk = unsafe { current.read() };
+                if let Some(index) = chunk.simd_eq(chunk_needle).first_set() {
+                    return unsafe { current.cast::<u8>().add(index) as *mut c_void }
+                }
+
+                current = unsafe { current.add(1) };
+                n -= size_of::<Chunk>();
+            }
+        }
+        // Unfortunately LLVM is not smart enough to use SWAR (SIMD-within-a-register)
+        // techniques if native SIMD is not available, so we need to do SWAR manually.
+        _ => {
+            const LOW_BITS: usize = usize::from_ne_bytes([0x01; _]);
+            const HIGH_BITS: usize = usize::from_ne_bytes([0x80; _]);
+
+            /// Returns `true` if `x` contains any zero byte.
+            ///
+            /// From *Matters Computational*, J. Arndt:
+            ///
+            /// "The idea is to subtract one from each of the bytes and then look for
+            /// bytes where the borrow propagated all the way to the most significant
+            /// bit."
+            #[inline]
+            const fn contains_zero_byte(x: usize) -> bool {
+                x.wrapping_sub(LOW_BITS) & !x & HIGH_BITS != 0
+            }
+
+            let mut current = current.cast::<Chunk>();
+            // Since `needle` is 8-bit wide, this multiplication will splat those
+            // 8 bits over all bytes of `usize`.
+            let chunk_needle = LOW_BITS * needle as usize;
+            while n >= size_of::<Chunk>() {
+                // Compare two words in one go.
+                let [lower, upper] = unsafe { current.read() };
+                // If the byte matches, then the XOR will result in that byte
+                // being zero.
+                let a = contains_zero_byte(lower ^ chunk_needle);
+                let b = contains_zero_byte(upper ^ chunk_needle);
+                if a | b {
+                    // Find the matching byte. The loop doesn't need to be bounded
+                    // since we know there is a match.
+                    let mut current = current.cast::<u8>();
+                    loop {
+                        let byte = unsafe { current.read() };
+                        if byte == needle {
+                            return current as *mut c_void;
+                        }
+
+                        current = unsafe { current.add(1) };
+                    }
+                }
+
+                current = unsafe { current.add(1) };
+                n -= size_of::<Chunk>();
+            }
+        }
+    }
+
+    // Search in the remaining bytes.
+    let current = current.cast::<u8>();
+    unsafe { memchr_naive(current, needle, n) }
+}

library/compiler-builtins/compiler-builtins/src/mem/mod.rs

@@ -6,6 +6,7 @@
 // memcpy/memmove/memset have optimized implementations on some architectures
 #[cfg_attr(all(feature = "arch", target_arch = "x86_64"), path = "x86_64.rs")]
 mod impls;
+mod memchr_impl;
 
 intrinsics! {
     #[mem_builtin]
@@ -67,4 +68,13 @@ intrinsics! {
     pub unsafe extern "C" fn strlen(s: *const core::ffi::c_char) -> usize {
         impls::c_string_length(s)
     }
+
+    #[mem_builtin]
+    pub unsafe extern "C" fn memchr(
+        s: *const core::ffi::c_void,
+        c: core::ffi::c_int,
+        n: usize
+    ) -> *mut core::ffi::c_void {
+        memchr_impl::memchr(s, c, n)
+    }
 }

library/core/src/lib.rs

@@ -20,9 +20,9 @@
 // FIXME: Fill me in with more detail when the interface settles
 //! This library is built on the assumption of a few existing symbols:
 //!
-//! * `memcpy`, `memmove`, `memset`, `memcmp`, `bcmp`, `strlen` - These are core memory routines
+//! * `memcpy`, `memmove`, `memset`, `memcmp`, `bcmp`, `strlen`, `memchr` - These are core memory routines
 //!   which are generated by Rust codegen backends. Additionally, this library can make explicit
-//!   calls to `strlen`. Their signatures are the same as found in C, but there are extra
+//!   calls to `strlen` and `memchr`. Their signatures are the same as found in C, but there are extra
 //!   assumptions about their semantics: For `memcpy`, `memmove`, `memset`, `memcmp`, and `bcmp`, if
 //!   the `n` parameter is 0, the function is assumed to not be UB, even if the pointers are NULL or
 //!   dangling. (Note that making extra assumptions about these functions is common among compilers:

library/core/src/slice/memchr.rs

@@ -1,11 +1,11 @@
 // Original implementation taken from rust-memchr.
 // Copyright 2015 Andrew Gallant, bluss and Nicolas Koch
 
+use crate::ffi::{c_int, c_void};
 use crate::intrinsics::const_eval_select;
 
 const LO_USIZE: usize = usize::repeat_u8(0x01);
 const HI_USIZE: usize = usize::repeat_u8(0x80);
-const USIZE_BYTES: usize = size_of::<usize>();
 
 /// Returns `true` if `x` contains any zero byte.
 ///
@@ -22,86 +22,47 @@ const fn contains_zero_byte(x: usize) -> bool {
 /// Returns the first index matching the byte `x` in `text`.
 #[inline]
 #[must_use]
+#[rustc_allow_const_fn_unstable(const_eval_select)] // both impls have the exact same behaviour
 pub const fn memchr(x: u8, text: &[u8]) -> Option<usize> {
-    // Fast path for small slices.
-    if text.len() < 2 * USIZE_BYTES {
-        return memchr_naive(x, text);
-    }
-
-    memchr_aligned(x, text)
-}
-
-#[inline]
-const fn memchr_naive(x: u8, text: &[u8]) -> Option<usize> {
-    let mut i = 0;
-
-    // FIXME(const-hack): Replace with `text.iter().pos(|c| *c == x)`.
-    while i < text.len() {
-        if text[i] == x {
-            return Some(i);
-        }
-
-        i += 1;
-    }
-
-    None
-}
-
-#[rustc_allow_const_fn_unstable(const_eval_select)] // fallback impl has same behavior
-const fn memchr_aligned(x: u8, text: &[u8]) -> Option<usize> {
-    // The runtime version behaves the same as the compiletime version, it's
-    // just more optimized.
     const_eval_select!(
-        @capture { x: u8, text: &[u8] } -> Option<usize>:
+        @capture { x: u8 = x, text: &[u8] = text } -> Option<usize>:
         if const {
-            memchr_naive(x, text)
-        } else {
-            // Scan for a single byte value by reading two `usize` words at a time.
-            //
-            // Split `text` in three parts
-            // - unaligned initial part, before the first word aligned address in text
-            // - body, scan by 2 words at a time
-            // - the last remaining part, < 2 word size
+            let mut i = 0;
 
-            // search up to an aligned boundary
-            let len = text.len();
-            let ptr = text.as_ptr();
-            let mut offset = ptr.align_offset(USIZE_BYTES);
-
-            if offset > 0 {
-                offset = offset.min(len);
-                let slice = &text[..offset];
-                if let Some(index) = memchr_naive(x, slice) {
-                    return Some(index);
+            // FIXME(const-hack): Replace with `text.iter().pos(|c| *c == x)`.
+            while i < text.len() {
+                if text[i] == x {
+                    return Some(i);
                 }
-            }
 
-            // search the body of the text
-            let repeated_x = usize::repeat_u8(x);
-            while offset <= len - 2 * USIZE_BYTES {
-                // SAFETY: the while's predicate guarantees a distance of at least 2 * usize_bytes
-                // between the offset and the end of the slice.
-                unsafe {
-                    let u = *(ptr.add(offset) as *const usize);
-                    let v = *(ptr.add(offset + USIZE_BYTES) as *const usize);
+                i += 1;
+            }
 
-                    // break if there is a matching byte
-                    let zu = contains_zero_byte(u ^ repeated_x);
-                    let zv = contains_zero_byte(v ^ repeated_x);
-                    if zu || zv {
-                        break;
-                    }
-                }
-                offset += USIZE_BYTES * 2;
+            None
+        } else {
+            unsafe extern "C" {
+                // Provided in either libc or compiler-builtins.
+                fn memchr(
+                    s: *const c_void,
+                    c: c_int,
+                    n: usize,
+                ) -> *mut c_void;
             }
 
-            // Find the byte after the point the body loop stopped.
-            // FIXME(const-hack): Use `?` instead.
-            // FIXME(const-hack, fee1-dead): use range slicing
-            let slice =
-            // SAFETY: offset is within bounds
-                unsafe { super::from_raw_parts(text.as_ptr().add(offset), text.len() - offset) };
-            if let Some(i) = memchr_naive(x, slice) { Some(offset + i) } else { None }
+            // SAFETY:
+            // The pointer and length come from a slice reference and thus
+            // describe a valid memory region. Since the reference is a `&[u8]`,
+            // every byte contained therein is interpretable as an initialized
+            // byte.
+            let res = unsafe { memchr(text.as_ptr().cast(), x as c_int, text.len()) };
+            if res.is_null() {
+                None
+            } else {
+                // SAFETY: `res` is non-null, and thus is guaranteed to lie
+                // within the memory region passed to `memchr`.
+                let index = unsafe { res.offset_from_unsigned(ptr) };
+                Some(index)
+            }
         }
     )
 }