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