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