1 | // Original implementation taken from rust-memchr. |
2 | // Copyright 2015 Andrew Gallant, bluss and Nicolas Koch |
3 | |
4 | use crate::mem; |
5 | |
6 | const LO_USIZE: usize = usize::repeat_u8(0x01); |
7 | const HI_USIZE: usize = usize::repeat_u8(0x80); |
8 | const 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" )] |
19 | const 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" )] |
27 | pub 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" )] |
38 | const 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" )] |
57 | const 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 ] |
110 | pub 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 | |