]> git.lizzy.rs Git - rust.git/blob - library/core/src/slice/memchr.rs
Rollup merge of #88041 - Aaron1011:deny-proc-macro-hack, r=wesleywiser
[rust.git] / 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 #[inline]
41 pub fn memchr(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 fn memchr_general_case(x: u8, text: &[u8]) -> Option<usize> {
51     // Scan for a single byte value by reading two `usize` words at a time.
52     //
53     // Split `text` in three parts
54     // - unaligned initial part, before the first word aligned address in text
55     // - body, scan by 2 words at a time
56     // - the last remaining part, < 2 word size
57
58     // search up to an aligned boundary
59     let len = text.len();
60     let ptr = text.as_ptr();
61     let mut offset = ptr.align_offset(USIZE_BYTES);
62
63     if offset > 0 {
64         offset = cmp::min(offset, len);
65         if let Some(index) = text[..offset].iter().position(|elt| *elt == x) {
66             return Some(index);
67         }
68     }
69
70     // search the body of the text
71     let repeated_x = repeat_byte(x);
72     while offset <= len - 2 * USIZE_BYTES {
73         // SAFETY: the while's predicate guarantees a distance of at least 2 * usize_bytes
74         // between the offset and the end of the slice.
75         unsafe {
76             let u = *(ptr.add(offset) as *const usize);
77             let v = *(ptr.add(offset + USIZE_BYTES) as *const usize);
78
79             // break if there is a matching byte
80             let zu = contains_zero_byte(u ^ repeated_x);
81             let zv = contains_zero_byte(v ^ repeated_x);
82             if zu || zv {
83                 break;
84             }
85         }
86         offset += USIZE_BYTES * 2;
87     }
88
89     // Find the byte after the point the body loop stopped.
90     text[offset..].iter().position(|elt| *elt == x).map(|i| offset + i)
91 }
92
93 /// Returns the last index matching the byte `x` in `text`.
94 pub fn memrchr(x: u8, text: &[u8]) -> Option<usize> {
95     // Scan for a single byte value by reading two `usize` words at a time.
96     //
97     // Split `text` in three parts:
98     // - unaligned tail, after the last word aligned address in text,
99     // - body, scanned by 2 words at a time,
100     // - the first remaining bytes, < 2 word size.
101     let len = text.len();
102     let ptr = text.as_ptr();
103     type Chunk = usize;
104
105     let (min_aligned_offset, max_aligned_offset) = {
106         // We call this just to obtain the length of the prefix and suffix.
107         // In the middle we always process two chunks at once.
108         // SAFETY: transmuting `[u8]` to `[usize]` is safe except for size differences
109         // which are handled by `align_to`.
110         let (prefix, _, suffix) = unsafe { text.align_to::<(Chunk, Chunk)>() };
111         (prefix.len(), len - suffix.len())
112     };
113
114     let mut offset = max_aligned_offset;
115     if let Some(index) = text[offset..].iter().rposition(|elt| *elt == x) {
116         return Some(offset + index);
117     }
118
119     // Search the body of the text, make sure we don't cross min_aligned_offset.
120     // offset is always aligned, so just testing `>` is sufficient and avoids possible
121     // overflow.
122     let repeated_x = repeat_byte(x);
123     let chunk_bytes = mem::size_of::<Chunk>();
124
125     while offset > min_aligned_offset {
126         // SAFETY: offset starts at len - suffix.len(), as long as it is greater than
127         // min_aligned_offset (prefix.len()) the remaining distance is at least 2 * chunk_bytes.
128         unsafe {
129             let u = *(ptr.offset(offset as isize - 2 * chunk_bytes as isize) as *const Chunk);
130             let v = *(ptr.offset(offset as isize - chunk_bytes as isize) as *const Chunk);
131
132             // Break if there is a matching byte.
133             let zu = contains_zero_byte(u ^ repeated_x);
134             let zv = contains_zero_byte(v ^ repeated_x);
135             if zu || zv {
136                 break;
137             }
138         }
139         offset -= 2 * chunk_bytes;
140     }
141
142     // Find the byte before the point the body loop stopped.
143     text[..offset].iter().rposition(|elt| *elt == x)
144 }