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