library/core/src/slice/memchr.rs

   1 // Original implementation taken from rust-memchr.
   2 // Copyright 2015 Andrew Gallant, bluss and Nicolas Koch
   3
   4 use crate::cmp;
   5 use crate::intrinsics;
   6 use crate::mem;
   7
   8 const LO_USIZE: usize = usize::repeat_u8(0x01);
   9 const HI_USIZE: usize = usize::repeat_u8(0x80);
  10 const USIZE_BYTES: usize = mem::size_of::<usize>();
  11
  12 /// Returns `true` if `x` contains any zero byte.
  13 ///
  14 /// From *Matters Computational*, J. Arndt:
  15 ///
  16 /// "The idea is to subtract one from each of the bytes and then look for
  17 /// bytes where the borrow propagated all the way to the most significant
  18 /// bit."
  19 #[inline]
  20 fn contains_zero_byte(x: usize) -> bool {
  21     x.wrapping_sub(LO_USIZE) & !x & HI_USIZE != 0
  22 }
  23
  24 #[cfg(target_pointer_width = "16")]
  25 #[inline]
  26 fn repeat_byte(b: u8) -> usize {
  27     (b as usize) << 8 | b as usize
  28 }
  29
  30 #[cfg(not(target_pointer_width = "16"))]
  31 #[inline]
  32 fn repeat_byte(b: u8) -> usize {
  33     (b as usize) * (usize::MAX / 255)
  34 }
  35
  36 /// Returns the first index matching the byte `x` in `text`.
  37 #[must_use]
  38 #[inline]
  39 pub const fn memchr(x: u8, text: &[u8]) -> Option<usize> {
  40     #[inline]
  41     fn rt_impl(x: u8, text: &[u8]) -> Option<usize> {
  42         // Fast path for small slices
  43         if text.len() < 2 * USIZE_BYTES {
  44             return text.iter().position(|elt| *elt == x);
  45         }
  46
  47         memchr_general_case(x, text)
  48     }
  49
  50     const fn const_impl(x: u8, bytes: &[u8]) -> Option<usize> {
  51         let mut i = 0;
  52         while i < bytes.len() {
  53             if bytes[i] == x {
  54                 return Some(i);
  55             }
  56             i += 1;
  57         }
  58
  59         None
  60     }
  61
  62     // SAFETY: The const and runtime versions have identical behavior
  63     unsafe { intrinsics::const_eval_select((x, text), const_impl, rt_impl) }
  64 }
  65
  66 fn memchr_general_case(x: u8, text: &[u8]) -> Option<usize> {
  67     // Scan for a single byte value by reading two `usize` words at a time.
  68     //
  69     // Split `text` in three parts
  70     // - unaligned initial part, before the first word aligned address in text
  71     // - body, scan by 2 words at a time
  72     // - the last remaining part, < 2 word size
  73
  74     // search up to an aligned boundary
  75     let len = text.len();
  76     let ptr = text.as_ptr();
  77     let mut offset = ptr.align_offset(USIZE_BYTES);
  78
  79     if offset > 0 {
  80         offset = cmp::min(offset, len);
  81         if let Some(index) = text[..offset].iter().position(|elt| *elt == x) {
  82             return Some(index);
  83         }
  84     }
  85
  86     // search the body of the text
  87     let repeated_x = repeat_byte(x);
  88     while offset <= len - 2 * USIZE_BYTES {
  89         // SAFETY: the while's predicate guarantees a distance of at least 2 * usize_bytes
  90         // between the offset and the end of the slice.
  91         unsafe {
  92             let u = *(ptr.add(offset) as *const usize);
  93             let v = *(ptr.add(offset + USIZE_BYTES) as *const usize);
  94
  95             // break if there is a matching byte
  96             let zu = contains_zero_byte(u ^ repeated_x);
  97             let zv = contains_zero_byte(v ^ repeated_x);
  98             if zu || zv {
  99                 break;
 100             }
 101         }
 102         offset += USIZE_BYTES * 2;
 103     }
 104
 105     // Find the byte after the point the body loop stopped.
 106     text[offset..].iter().position(|elt| *elt == x).map(|i| offset + i)
 107 }
 108
 109 /// Returns the last index matching the byte `x` in `text`.
 110 #[must_use]
 111 pub fn memrchr(x: u8, text: &[u8]) -> Option<usize> {
 112     // Scan for a single byte value by reading two `usize` words at a time.
 113     //
 114     // Split `text` in three parts:
 115     // - unaligned tail, after the last word aligned address in text,
 116     // - body, scanned by 2 words at a time,
 117     // - the first remaining bytes, < 2 word size.
 118     let len = text.len();
 119     let ptr = text.as_ptr();
 120     type Chunk = usize;
 121
 122     let (min_aligned_offset, max_aligned_offset) = {
 123         // We call this just to obtain the length of the prefix and suffix.
 124         // In the middle we always process two chunks at once.
 125         // SAFETY: transmuting `[u8]` to `[usize]` is safe except for size differences
 126         // which are handled by `align_to`.
 127         let (prefix, _, suffix) = unsafe { text.align_to::<(Chunk, Chunk)>() };
 128         (prefix.len(), len - suffix.len())
 129     };
 130
 131     let mut offset = max_aligned_offset;
 132     if let Some(index) = text[offset..].iter().rposition(|elt| *elt == x) {
 133         return Some(offset + index);
 134     }
 135
 136     // Search the body of the text, make sure we don't cross min_aligned_offset.
 137     // offset is always aligned, so just testing `>` is sufficient and avoids possible
 138     // overflow.
 139     let repeated_x = repeat_byte(x);
 140     let chunk_bytes = mem::size_of::<Chunk>();
 141
 142     while offset > min_aligned_offset {
 143         // SAFETY: offset starts at len - suffix.len(), as long as it is greater than
 144         // min_aligned_offset (prefix.len()) the remaining distance is at least 2 * chunk_bytes.
 145         unsafe {
 146             let u = *(ptr.offset(offset as isize - 2 * chunk_bytes as isize) as *const Chunk);
 147             let v = *(ptr.offset(offset as isize - chunk_bytes as isize) as *const Chunk);
 148
 149             // Break if there is a matching byte.
 150             let zu = contains_zero_byte(u ^ repeated_x);
 151             let zv = contains_zero_byte(v ^ repeated_x);
 152             if zu || zv {
 153                 break;
 154             }
 155         }
 156         offset -= 2 * chunk_bytes;
 157     }
 158
 159     // Find the byte before the point the body loop stopped.
 160     text[..offset].iter().rposition(|elt| *elt == x)
 161 }