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