core/slice/ascii.rs
1//! Operations on ASCII `[u8]`.
2
3use core::ascii::EscapeDefault;
4
5use crate::fmt::{self, Write};
6#[cfg(not(all(target_arch = "loongarch64", target_feature = "lsx")))]
7use crate::intrinsics::const_eval_select;
8use crate::{ascii, iter, ops};
9
10impl [u8] {
11 /// Checks if all bytes in this slice are within the ASCII range.
12 ///
13 /// An empty slice returns `true`.
14 #[stable(feature = "ascii_methods_on_intrinsics", since = "1.23.0")]
15 #[rustc_const_stable(feature = "const_slice_is_ascii", since = "1.74.0")]
16 #[must_use]
17 #[inline]
18 pub const fn is_ascii(&self) -> bool {
19 is_ascii(self)
20 }
21
22 /// If this slice [`is_ascii`](Self::is_ascii), returns it as a slice of
23 /// [ASCII characters](`ascii::Char`), otherwise returns `None`.
24 #[unstable(feature = "ascii_char", issue = "110998")]
25 #[must_use]
26 #[inline]
27 pub const fn as_ascii(&self) -> Option<&[ascii::Char]> {
28 if self.is_ascii() {
29 // SAFETY: Just checked that it's ASCII
30 Some(unsafe { self.as_ascii_unchecked() })
31 } else {
32 None
33 }
34 }
35
36 /// Converts this slice of bytes into a slice of ASCII characters,
37 /// without checking whether they're valid.
38 ///
39 /// # Safety
40 ///
41 /// Every byte in the slice must be in `0..=127`, or else this is UB.
42 #[unstable(feature = "ascii_char", issue = "110998")]
43 #[must_use]
44 #[inline]
45 pub const unsafe fn as_ascii_unchecked(&self) -> &[ascii::Char] {
46 let byte_ptr: *const [u8] = self;
47 let ascii_ptr = byte_ptr as *const [ascii::Char];
48 // SAFETY: The caller promised all the bytes are ASCII
49 unsafe { &*ascii_ptr }
50 }
51
52 /// Checks that two slices are an ASCII case-insensitive match.
53 ///
54 /// Same as `to_ascii_lowercase(a) == to_ascii_lowercase(b)`,
55 /// but without allocating and copying temporaries.
56 #[stable(feature = "ascii_methods_on_intrinsics", since = "1.23.0")]
57 #[rustc_const_stable(feature = "const_eq_ignore_ascii_case", since = "1.89.0")]
58 #[must_use]
59 #[inline]
60 pub const fn eq_ignore_ascii_case(&self, other: &[u8]) -> bool {
61 if self.len() != other.len() {
62 return false;
63 }
64
65 #[cfg(any(
66 all(target_arch = "x86_64", target_feature = "sse2"),
67 all(target_arch = "aarch64", target_feature = "neon")
68 ))]
69 {
70 const CHUNK_SIZE: usize = 16;
71 // The following function has two invariants:
72 // 1. The slice lengths must be equal, which we checked above.
73 // 2. The slice lengths must greater than or equal to N, which this
74 // if-statement is checking.
75 if self.len() >= CHUNK_SIZE {
76 return self.eq_ignore_ascii_case_chunks::<CHUNK_SIZE>(other);
77 }
78 }
79
80 self.eq_ignore_ascii_case_simple(other)
81 }
82
83 /// ASCII case-insensitive equality check without chunk-at-a-time
84 /// optimization.
85 #[inline]
86 const fn eq_ignore_ascii_case_simple(&self, other: &[u8]) -> bool {
87 // FIXME(const-hack): This implementation can be reverted when
88 // `core::iter::zip` is allowed in const. The original implementation:
89 // self.len() == other.len() && iter::zip(self, other).all(|(a, b)| a.eq_ignore_ascii_case(b))
90 let mut a = self;
91 let mut b = other;
92
93 while let ([first_a, rest_a @ ..], [first_b, rest_b @ ..]) = (a, b) {
94 if first_a.eq_ignore_ascii_case(&first_b) {
95 a = rest_a;
96 b = rest_b;
97 } else {
98 return false;
99 }
100 }
101
102 true
103 }
104
105 /// Optimized version of `eq_ignore_ascii_case` to process chunks at a time.
106 ///
107 /// Platforms that have SIMD instructions may benefit from this
108 /// implementation over `eq_ignore_ascii_case_simple`.
109 ///
110 /// # Invariants
111 ///
112 /// The caller must guarantee that the slices are equal in length, and the
113 /// slice lengths are greater than or equal to `N` bytes.
114 #[cfg(any(
115 all(target_arch = "x86_64", target_feature = "sse2"),
116 all(target_arch = "aarch64", target_feature = "neon")
117 ))]
118 #[inline]
119 const fn eq_ignore_ascii_case_chunks<const N: usize>(&self, other: &[u8]) -> bool {
120 // FIXME(const-hack): The while-loops that follow should be replaced by
121 // for-loops when available in const.
122
123 let (self_chunks, self_rem) = self.as_chunks::<N>();
124 let (other_chunks, _) = other.as_chunks::<N>();
125
126 // Branchless check to encourage auto-vectorization
127 #[inline(always)]
128 const fn eq_ignore_ascii_inner<const L: usize>(lhs: &[u8; L], rhs: &[u8; L]) -> bool {
129 let mut equal_ascii = true;
130 let mut j = 0;
131 while j < L {
132 equal_ascii &= lhs[j].eq_ignore_ascii_case(&rhs[j]);
133 j += 1;
134 }
135
136 equal_ascii
137 }
138
139 // Process the chunks, returning early if an inequality is found
140 let mut i = 0;
141 while i < self_chunks.len() && i < other_chunks.len() {
142 if !eq_ignore_ascii_inner(&self_chunks[i], &other_chunks[i]) {
143 return false;
144 }
145 i += 1;
146 }
147
148 // Check the length invariant which is necessary for the tail-handling
149 // logic to be correct. This should have been upheld by the caller,
150 // otherwise lengths less than N will compare as true without any
151 // checking.
152 debug_assert!(self.len() >= N);
153
154 // If there are remaining tails, load the last N bytes in the slices to
155 // avoid falling back to per-byte checking.
156 if !self_rem.is_empty() {
157 if let (Some(a_rem), Some(b_rem)) = (self.last_chunk::<N>(), other.last_chunk::<N>()) {
158 if !eq_ignore_ascii_inner(a_rem, b_rem) {
159 return false;
160 }
161 }
162 }
163
164 true
165 }
166
167 /// Converts this slice to its ASCII upper case equivalent in-place.
168 ///
169 /// ASCII letters 'a' to 'z' are mapped to 'A' to 'Z',
170 /// but non-ASCII letters are unchanged.
171 ///
172 /// To return a new uppercased value without modifying the existing one, use
173 /// [`to_ascii_uppercase`].
174 ///
175 /// [`to_ascii_uppercase`]: #method.to_ascii_uppercase
176 #[stable(feature = "ascii_methods_on_intrinsics", since = "1.23.0")]
177 #[rustc_const_stable(feature = "const_make_ascii", since = "1.84.0")]
178 #[inline]
179 pub const fn make_ascii_uppercase(&mut self) {
180 // FIXME(const-hack): We would like to simply iterate using `for` loops but this isn't currently allowed in constant expressions.
181 let mut i = 0;
182 while i < self.len() {
183 let byte = &mut self[i];
184 byte.make_ascii_uppercase();
185 i += 1;
186 }
187 }
188
189 /// Converts this slice to its ASCII lower case equivalent in-place.
190 ///
191 /// ASCII letters 'A' to 'Z' are mapped to 'a' to 'z',
192 /// but non-ASCII letters are unchanged.
193 ///
194 /// To return a new lowercased value without modifying the existing one, use
195 /// [`to_ascii_lowercase`].
196 ///
197 /// [`to_ascii_lowercase`]: #method.to_ascii_lowercase
198 #[stable(feature = "ascii_methods_on_intrinsics", since = "1.23.0")]
199 #[rustc_const_stable(feature = "const_make_ascii", since = "1.84.0")]
200 #[inline]
201 pub const fn make_ascii_lowercase(&mut self) {
202 // FIXME(const-hack): We would like to simply iterate using `for` loops but this isn't currently allowed in constant expressions.
203 let mut i = 0;
204 while i < self.len() {
205 let byte = &mut self[i];
206 byte.make_ascii_lowercase();
207 i += 1;
208 }
209 }
210
211 /// Returns an iterator that produces an escaped version of this slice,
212 /// treating it as an ASCII string.
213 ///
214 /// # Examples
215 ///
216 /// ```
217 /// let s = b"0\t\r\n'\"\\\x9d";
218 /// let escaped = s.escape_ascii().to_string();
219 /// assert_eq!(escaped, "0\\t\\r\\n\\'\\\"\\\\\\x9d");
220 /// ```
221 #[must_use = "this returns the escaped bytes as an iterator, \
222 without modifying the original"]
223 #[stable(feature = "inherent_ascii_escape", since = "1.60.0")]
224 pub fn escape_ascii(&self) -> EscapeAscii<'_> {
225 EscapeAscii { inner: self.iter().flat_map(EscapeByte) }
226 }
227
228 /// Returns a byte slice with leading ASCII whitespace bytes removed.
229 ///
230 /// 'Whitespace' refers to the definition used by
231 /// [`u8::is_ascii_whitespace`]. Importantly, this definition excludes
232 /// the `\0x0B` byte even though it has the Unicode [`White_Space`] property
233 /// and is removed by [`str::trim_start`].
234 ///
235 /// [`White_Space`]: https://www.unicode.org/reports/tr44/#White_Space
236 ///
237 /// # Examples
238 ///
239 /// ```
240 /// assert_eq!(b" \t hello world\n".trim_ascii_start(), b"hello world\n");
241 /// assert_eq!(b" ".trim_ascii_start(), b"");
242 /// assert_eq!(b"".trim_ascii_start(), b"");
243 /// ```
244 #[stable(feature = "byte_slice_trim_ascii", since = "1.80.0")]
245 #[rustc_const_stable(feature = "byte_slice_trim_ascii", since = "1.80.0")]
246 #[inline]
247 pub const fn trim_ascii_start(&self) -> &[u8] {
248 let mut bytes = self;
249 // Note: A pattern matching based approach (instead of indexing) allows
250 // making the function const.
251 while let [first, rest @ ..] = bytes {
252 if first.is_ascii_whitespace() {
253 bytes = rest;
254 } else {
255 break;
256 }
257 }
258 bytes
259 }
260
261 /// Returns a byte slice with trailing ASCII whitespace bytes removed.
262 ///
263 /// 'Whitespace' refers to the definition used by
264 /// [`u8::is_ascii_whitespace`]. Importantly, this definition excludes
265 /// the `\0x0B` byte even though it has the Unicode [`White_Space`] property
266 /// and is removed by [`str::trim_end`].
267 ///
268 /// [`White_Space`]: https://www.unicode.org/reports/tr44/#White_Space
269 ///
270 /// # Examples
271 ///
272 /// ```
273 /// assert_eq!(b"\r hello world\n ".trim_ascii_end(), b"\r hello world");
274 /// assert_eq!(b" ".trim_ascii_end(), b"");
275 /// assert_eq!(b"".trim_ascii_end(), b"");
276 /// ```
277 #[stable(feature = "byte_slice_trim_ascii", since = "1.80.0")]
278 #[rustc_const_stable(feature = "byte_slice_trim_ascii", since = "1.80.0")]
279 #[inline]
280 pub const fn trim_ascii_end(&self) -> &[u8] {
281 let mut bytes = self;
282 // Note: A pattern matching based approach (instead of indexing) allows
283 // making the function const.
284 while let [rest @ .., last] = bytes {
285 if last.is_ascii_whitespace() {
286 bytes = rest;
287 } else {
288 break;
289 }
290 }
291 bytes
292 }
293
294 /// Returns a byte slice with leading and trailing ASCII whitespace bytes
295 /// removed.
296 ///
297 /// 'Whitespace' refers to the definition used by
298 /// [`u8::is_ascii_whitespace`]. Importantly, this definition excludes
299 /// the `\0x0B` byte even though it has the Unicode [`White_Space`] property
300 /// and is removed by [`str::trim`].
301 ///
302 /// [`White_Space`]: https://www.unicode.org/reports/tr44/#White_Space
303 ///
304 /// # Examples
305 ///
306 /// ```
307 /// assert_eq!(b"\r hello world\n ".trim_ascii(), b"hello world");
308 /// assert_eq!(b" ".trim_ascii(), b"");
309 /// assert_eq!(b"".trim_ascii(), b"");
310 /// ```
311 #[stable(feature = "byte_slice_trim_ascii", since = "1.80.0")]
312 #[rustc_const_stable(feature = "byte_slice_trim_ascii", since = "1.80.0")]
313 #[inline]
314 pub const fn trim_ascii(&self) -> &[u8] {
315 self.trim_ascii_start().trim_ascii_end()
316 }
317}
318
319impl_fn_for_zst! {
320 #[derive(Clone)]
321 struct EscapeByte impl Fn = |byte: &u8| -> ascii::EscapeDefault {
322 ascii::escape_default(*byte)
323 };
324}
325
326/// An iterator over the escaped version of a byte slice.
327///
328/// This `struct` is created by the [`slice::escape_ascii`] method. See its
329/// documentation for more information.
330#[stable(feature = "inherent_ascii_escape", since = "1.60.0")]
331#[derive(Clone)]
332#[must_use = "iterators are lazy and do nothing unless consumed"]
333pub struct EscapeAscii<'a> {
334 inner: iter::FlatMap<super::Iter<'a, u8>, ascii::EscapeDefault, EscapeByte>,
335}
336
337#[stable(feature = "inherent_ascii_escape", since = "1.60.0")]
338impl<'a> iter::Iterator for EscapeAscii<'a> {
339 type Item = u8;
340 #[inline]
341 fn next(&mut self) -> Option<u8> {
342 self.inner.next()
343 }
344 #[inline]
345 fn size_hint(&self) -> (usize, Option<usize>) {
346 self.inner.size_hint()
347 }
348 #[inline]
349 fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
350 where
351 Fold: FnMut(Acc, Self::Item) -> R,
352 R: ops::Try<Output = Acc>,
353 {
354 self.inner.try_fold(init, fold)
355 }
356 #[inline]
357 fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
358 where
359 Fold: FnMut(Acc, Self::Item) -> Acc,
360 {
361 self.inner.fold(init, fold)
362 }
363 #[inline]
364 fn last(mut self) -> Option<u8> {
365 self.next_back()
366 }
367}
368
369#[stable(feature = "inherent_ascii_escape", since = "1.60.0")]
370impl<'a> iter::DoubleEndedIterator for EscapeAscii<'a> {
371 fn next_back(&mut self) -> Option<u8> {
372 self.inner.next_back()
373 }
374}
375#[stable(feature = "inherent_ascii_escape", since = "1.60.0")]
376impl<'a> iter::FusedIterator for EscapeAscii<'a> {}
377#[stable(feature = "inherent_ascii_escape", since = "1.60.0")]
378impl<'a> fmt::Display for EscapeAscii<'a> {
379 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
380 // disassemble iterator, including front/back parts of flatmap in case it has been partially consumed
381 let (front, slice, back) = self.clone().inner.into_parts();
382 let front = front.unwrap_or(EscapeDefault::empty());
383 let mut bytes = slice.unwrap_or_default().as_slice();
384 let back = back.unwrap_or(EscapeDefault::empty());
385
386 // usually empty, so the formatter won't have to do any work
387 for byte in front {
388 f.write_char(byte as char)?;
389 }
390
391 fn needs_escape(b: u8) -> bool {
392 b > 0x7E || b < 0x20 || b == b'\\' || b == b'\'' || b == b'"'
393 }
394
395 while bytes.len() > 0 {
396 // fast path for the printable, non-escaped subset of ascii
397 let prefix = bytes.iter().take_while(|&&b| !needs_escape(b)).count();
398 // SAFETY: prefix length was derived by counting bytes in the same splice, so it's in-bounds
399 let (prefix, remainder) = unsafe { bytes.split_at_unchecked(prefix) };
400 // SAFETY: prefix is a valid utf8 sequence, as it's a subset of ASCII
401 let prefix = unsafe { crate::str::from_utf8_unchecked(prefix) };
402
403 f.write_str(prefix)?; // the fast part
404
405 bytes = remainder;
406
407 if let Some(&b) = bytes.first() {
408 // guaranteed to be non-empty, better to write it as a str
409 fmt::Display::fmt(&ascii::escape_default(b), f)?;
410 bytes = &bytes[1..];
411 }
412 }
413
414 // also usually empty
415 for byte in back {
416 f.write_char(byte as char)?;
417 }
418 Ok(())
419 }
420}
421#[stable(feature = "inherent_ascii_escape", since = "1.60.0")]
422impl<'a> fmt::Debug for EscapeAscii<'a> {
423 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
424 f.debug_struct("EscapeAscii").finish_non_exhaustive()
425 }
426}
427
428/// ASCII test *without* the chunk-at-a-time optimizations.
429///
430/// This is carefully structured to produce nice small code -- it's smaller in
431/// `-O` than what the "obvious" ways produces under `-C opt-level=s`. If you
432/// touch it, be sure to run (and update if needed) the assembly test.
433#[unstable(feature = "str_internals", issue = "none")]
434#[doc(hidden)]
435#[inline]
436pub const fn is_ascii_simple(mut bytes: &[u8]) -> bool {
437 while let [rest @ .., last] = bytes {
438 if !last.is_ascii() {
439 break;
440 }
441 bytes = rest;
442 }
443 bytes.is_empty()
444}
445
446/// Optimized ASCII test that will use usize-at-a-time operations instead of
447/// byte-at-a-time operations (when possible).
448///
449/// The algorithm we use here is pretty simple. If `s` is too short, we just
450/// check each byte and be done with it. Otherwise:
451///
452/// - Read the first word with an unaligned load.
453/// - Align the pointer, read subsequent words until end with aligned loads.
454/// - Read the last `usize` from `s` with an unaligned load.
455///
456/// If any of these loads produces something for which `contains_nonascii`
457/// (above) returns true, then we know the answer is false.
458#[cfg(not(any(
459 all(target_arch = "x86_64", target_feature = "sse2"),
460 all(target_arch = "loongarch64", target_feature = "lsx"),
461 all(target_arch = "aarch64", target_feature = "neon")
462)))]
463#[inline]
464#[rustc_allow_const_fn_unstable(const_eval_select)] // fallback impl has same behavior
465const fn is_ascii(s: &[u8]) -> bool {
466 // The runtime version behaves the same as the compiletime version, it's
467 // just more optimized.
468 const_eval_select!(
469 @capture { s: &[u8] } -> bool:
470 if const {
471 is_ascii_simple(s)
472 } else {
473 /// Returns `true` if any byte in the word `v` is nonascii (>= 128). Snarfed
474 /// from `../str/mod.rs`, which does something similar for utf8 validation.
475 const fn contains_nonascii(v: usize) -> bool {
476 const NONASCII_MASK: usize = usize::repeat_u8(0x80);
477 (NONASCII_MASK & v) != 0
478 }
479
480 const USIZE_SIZE: usize = size_of::<usize>();
481
482 let len = s.len();
483 let align_offset = s.as_ptr().align_offset(USIZE_SIZE);
484
485 // If we wouldn't gain anything from the word-at-a-time implementation, fall
486 // back to a scalar loop.
487 //
488 // We also do this for architectures where `size_of::<usize>()` isn't
489 // sufficient alignment for `usize`, because it's a weird edge case.
490 if len < USIZE_SIZE || len < align_offset || USIZE_SIZE < align_of::<usize>() {
491 return is_ascii_simple(s);
492 }
493
494 // We always read the first word unaligned, which means `align_offset` is
495 // 0, we'd read the same value again for the aligned read.
496 let offset_to_aligned = if align_offset == 0 { USIZE_SIZE } else { align_offset };
497
498 let start = s.as_ptr();
499 // SAFETY: We verify `len < USIZE_SIZE` above.
500 let first_word = unsafe { (start as *const usize).read_unaligned() };
501
502 if contains_nonascii(first_word) {
503 return false;
504 }
505 // We checked this above, somewhat implicitly. Note that `offset_to_aligned`
506 // is either `align_offset` or `USIZE_SIZE`, both of are explicitly checked
507 // above.
508 debug_assert!(offset_to_aligned <= len);
509
510 // SAFETY: word_ptr is the (properly aligned) usize ptr we use to read the
511 // middle chunk of the slice.
512 let mut word_ptr = unsafe { start.add(offset_to_aligned) as *const usize };
513
514 // `byte_pos` is the byte index of `word_ptr`, used for loop end checks.
515 let mut byte_pos = offset_to_aligned;
516
517 // Paranoia check about alignment, since we're about to do a bunch of
518 // unaligned loads. In practice this should be impossible barring a bug in
519 // `align_offset` though.
520 // While this method is allowed to spuriously fail in CTFE, if it doesn't
521 // have alignment information it should have given a `usize::MAX` for
522 // `align_offset` earlier, sending things through the scalar path instead of
523 // this one, so this check should pass if it's reachable.
524 debug_assert!(word_ptr.is_aligned_to(align_of::<usize>()));
525
526 // Read subsequent words until the last aligned word, excluding the last
527 // aligned word by itself to be done in tail check later, to ensure that
528 // tail is always one `usize` at most to extra branch `byte_pos == len`.
529 while byte_pos < len - USIZE_SIZE {
530 // Sanity check that the read is in bounds
531 debug_assert!(byte_pos + USIZE_SIZE <= len);
532 // And that our assumptions about `byte_pos` hold.
533 debug_assert!(word_ptr.cast::<u8>() == start.wrapping_add(byte_pos));
534
535 // SAFETY: We know `word_ptr` is properly aligned (because of
536 // `align_offset`), and we know that we have enough bytes between `word_ptr` and the end
537 let word = unsafe { word_ptr.read() };
538 if contains_nonascii(word) {
539 return false;
540 }
541
542 byte_pos += USIZE_SIZE;
543 // SAFETY: We know that `byte_pos <= len - USIZE_SIZE`, which means that
544 // after this `add`, `word_ptr` will be at most one-past-the-end.
545 word_ptr = unsafe { word_ptr.add(1) };
546 }
547
548 // Sanity check to ensure there really is only one `usize` left. This should
549 // be guaranteed by our loop condition.
550 debug_assert!(byte_pos <= len && len - byte_pos <= USIZE_SIZE);
551
552 // SAFETY: This relies on `len >= USIZE_SIZE`, which we check at the start.
553 let last_word = unsafe { (start.add(len - USIZE_SIZE) as *const usize).read_unaligned() };
554
555 !contains_nonascii(last_word)
556 }
557 )
558}
559
560/// Chunk size for SSE2 vectorized ASCII checking (4x 16-byte loads).
561#[cfg(all(target_arch = "x86_64", target_feature = "sse2"))]
562const SSE2_CHUNK_SIZE: usize = 64;
563
564#[cfg(all(target_arch = "x86_64", target_feature = "sse2"))]
565#[inline]
566fn is_ascii_sse2(bytes: &[u8]) -> bool {
567 use crate::arch::x86_64::{__m128i, _mm_loadu_si128, _mm_movemask_epi8, _mm_or_si128};
568
569 let (chunks, rest) = bytes.as_chunks::<SSE2_CHUNK_SIZE>();
570
571 for chunk in chunks {
572 let ptr = chunk.as_ptr();
573 // SAFETY: chunk is 64 bytes. SSE2 is baseline on x86_64.
574 let mask = unsafe {
575 let a1 = _mm_loadu_si128(ptr as *const __m128i);
576 let a2 = _mm_loadu_si128(ptr.add(16) as *const __m128i);
577 let b1 = _mm_loadu_si128(ptr.add(32) as *const __m128i);
578 let b2 = _mm_loadu_si128(ptr.add(48) as *const __m128i);
579 // OR all chunks - if any byte has high bit set, combined will too.
580 let combined = _mm_or_si128(_mm_or_si128(a1, a2), _mm_or_si128(b1, b2));
581 // Create a mask from the MSBs of each byte.
582 // If any byte is >= 128, its MSB is 1, so the mask will be non-zero.
583 _mm_movemask_epi8(combined)
584 };
585 if mask != 0 {
586 return false;
587 }
588 }
589
590 // Handle remaining bytes
591 rest.iter().all(|b| b.is_ascii())
592}
593
594/// Chunk size for NEON vectorized ASCII checking (4x 16-byte loads).
595#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
596const NEON_CHUNK_SIZE: usize = 64;
597
598/// Width of a single NEON vector, used to vectorize the tail left over by the
599/// unrolled `NEON_CHUNK_SIZE` loop.
600#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
601const NEON_VECTOR_SIZE: usize = 16;
602
603#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
604#[inline]
605fn is_ascii_neon(bytes: &[u8]) -> bool {
606 use crate::arch::aarch64::{vld1q_u8, vmaxvq_u8, vorrq_u8};
607
608 let (chunks, rest) = bytes.as_chunks::<NEON_CHUNK_SIZE>();
609
610 for chunk in chunks {
611 let ptr = chunk.as_ptr();
612 // SAFETY: chunk is 64 bytes, and `vld1q_u8` has no alignment requirement.
613 let max = unsafe {
614 let a1 = vld1q_u8(ptr);
615 let a2 = vld1q_u8(ptr.add(16));
616 let b1 = vld1q_u8(ptr.add(32));
617 let b2 = vld1q_u8(ptr.add(48));
618 // OR all chunks - if any byte has high bit set, combined will too.
619 let combined = vorrq_u8(vorrq_u8(a1, a2), vorrq_u8(b1, b2));
620 // `vmaxvq_u8` is a horizontal reduction with a longer latency than
621 // `vorrq_u8`, so it runs once per 64 bytes rather than once per load.
622 vmaxvq_u8(combined)
623 };
624 if max >= 128 {
625 return false;
626 }
627 }
628
629 // The unrolled loop above leaves up to 63 bytes, so sweep those a vector at
630 // a time before falling back to a byte-at-a-time check.
631 let (vectors, rest) = rest.as_chunks::<NEON_VECTOR_SIZE>();
632
633 for vector in vectors {
634 // SAFETY: vector is 16 bytes, and `vld1q_u8` has no alignment requirement.
635 let max = unsafe { vmaxvq_u8(vld1q_u8(vector.as_ptr())) };
636 if max >= 128 {
637 return false;
638 }
639 }
640
641 // Handle remaining bytes
642 rest.iter().all(|b| b.is_ascii())
643}
644
645/// Uses explicit SIMD intrinsics to prevent LLVM from auto-vectorizing with
646/// broken code (e.g., AVX-512 on x86-64 that extracts mask bits one-by-one).
647#[cfg(any(
648 all(target_arch = "x86_64", target_feature = "sse2"),
649 all(target_arch = "aarch64", target_feature = "neon")
650))]
651#[inline]
652#[rustc_allow_const_fn_unstable(const_eval_select)]
653const fn is_ascii(bytes: &[u8]) -> bool {
654 const USIZE_SIZE: usize = size_of::<usize>();
655 const NONASCII_MASK: usize = usize::MAX / 255 * 0x80;
656
657 #[cfg(all(target_arch = "x86_64", target_feature = "sse2"))]
658 const SIMD_MIN_LEN: usize = SSE2_CHUNK_SIZE;
659 #[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
660 const SIMD_MIN_LEN: usize = NEON_CHUNK_SIZE;
661
662 const_eval_select!(
663 @capture { bytes: &[u8] } -> bool:
664 if const {
665 is_ascii_simple(bytes)
666 } else {
667 // For small inputs, use usize-at-a-time processing to avoid SSE2 call overhead.
668 if bytes.len() < SIMD_MIN_LEN {
669 let chunks = bytes.chunks_exact(USIZE_SIZE);
670 let remainder = chunks.remainder();
671 for chunk in chunks {
672 let word = usize::from_ne_bytes(chunk.try_into().unwrap());
673 if (word & NONASCII_MASK) != 0 {
674 return false;
675 }
676 }
677 return remainder.iter().all(|b| b.is_ascii());
678 }
679
680 #[cfg(all(target_arch = "x86_64", target_feature = "sse2"))]
681 { is_ascii_sse2(bytes) }
682 #[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
683 { is_ascii_neon(bytes) }
684 }
685 )
686}
687
688/// ASCII test optimized to use the `vmskltz.b` instruction on `loongarch64`.
689///
690/// Other platforms are not likely to benefit from this code structure, so they
691/// use SWAR techniques to test for ASCII in `usize`-sized chunks.
692#[cfg(all(target_arch = "loongarch64", target_feature = "lsx"))]
693#[inline]
694const fn is_ascii(bytes: &[u8]) -> bool {
695 // Process chunks of 32 bytes at a time in the fast path to enable
696 // auto-vectorization and use of `vmskltz.b`. Two 128-bit vector registers
697 // can be OR'd together and then the resulting vector can be tested for
698 // non-ASCII bytes.
699 const CHUNK_SIZE: usize = 32;
700
701 let mut i = 0;
702
703 while i + CHUNK_SIZE <= bytes.len() {
704 let chunk_end = i + CHUNK_SIZE;
705
706 // Get LLVM to produce a `vmskltz.b` instruction on loongarch64 which
707 // creates a mask from the most significant bit of each byte.
708 // ASCII bytes are less than 128 (0x80), so their most significant
709 // bit is unset.
710 let mut count = 0;
711 while i < chunk_end {
712 count += bytes[i].is_ascii() as u8;
713 i += 1;
714 }
715
716 // All bytes should be <= 127 so count is equal to chunk size.
717 if count != CHUNK_SIZE as u8 {
718 return false;
719 }
720 }
721
722 // Process the remaining `bytes.len() % N` bytes.
723 let mut is_ascii = true;
724 while i < bytes.len() {
725 is_ascii &= bytes[i].is_ascii();
726 i += 1;
727 }
728
729 is_ascii
730}