1use std::sync::atomic::{AtomicU8, Ordering};
15
16#[derive(Clone, Copy, Debug, PartialEq, Eq)]
18pub enum Level {
19 Scalar = 1,
21 Sse42 = 2,
23 Avx2 = 3,
25 Neon = 4,
27}
28
29#[derive(Debug)]
31pub struct ByteClass {
32 table: [bool; 256],
33 lo: [u8; 16],
34 high: bool,
35}
36
37impl ByteClass {
38 pub const fn new(singles: &[u8], ranges: &[(u8, u8)], high: bool) -> ByteClass {
41 let mut table = [false; 256];
42 let mut i = 0;
43 while i < singles.len() {
44 assert!(singles[i] < 0x80);
45 table[singles[i] as usize] = true;
46 i += 1;
47 }
48 let mut r = 0;
49 while r < ranges.len() {
50 let (a, b) = ranges[r];
51 assert!(a <= b && b < 0x80);
52 let mut c = a as usize;
53 while c <= b as usize {
54 table[c] = true;
55 c += 1;
56 }
57 r += 1;
58 }
59 if high {
60 let mut c = 0x80;
61 while c < 256 {
62 table[c] = true;
63 c += 1;
64 }
65 }
66 let mut lo = [0u8; 16];
67 let mut c = 0;
68 while c < 128 {
69 if table[c] {
70 lo[c & 15] |= 1 << (c >> 4);
71 }
72 c += 1;
73 }
74 ByteClass { table, lo, high }
75 }
76
77 #[inline]
79 pub fn contains(&self, b: u8) -> bool {
80 self.table[b as usize]
81 }
82}
83
84pub static TCHAR: ByteClass =
86 ByteClass::new(b"!#$%&'*+-.^_`|~", &[(b'0', b'9'), (b'A', b'Z'), (b'a', b'z')], false);
87
88pub static FIELD_VALUE: ByteClass = ByteClass::new(b"\t", &[(0x20, 0x7e)], true);
91
92pub static URI: ByteClass = ByteClass::new(
95 b"-._~!$&'()*+,;=:@/?",
96 &[(b'0', b'9'), (b'A', b'Z'), (b'a', b'z')],
97 false,
98);
99
100static LEVEL: AtomicU8 = AtomicU8::new(0);
101
102fn from_u8(n: u8) -> Level {
103 match n {
104 2 => Level::Sse42,
105 3 => Level::Avx2,
106 4 => Level::Neon,
107 _ => Level::Scalar,
108 }
109}
110
111pub fn available(level: Level) -> bool {
113 match level {
114 Level::Scalar => true,
115 #[cfg(target_arch = "x86_64")]
116 Level::Sse42 => {
117 use crate::cpu::x86;
118 crate::cpu::x86_features() & (x86::SSE42 | x86::SSSE3) == x86::SSE42 | x86::SSSE3
119 }
120 #[cfg(target_arch = "x86_64")]
121 Level::Avx2 => available(Level::Sse42) && crate::cpu::x86_features() & crate::cpu::x86::AVX2 != 0,
122 #[cfg(target_arch = "aarch64")]
123 Level::Neon => std::arch::is_aarch64_feature_detected!("neon"),
124 #[allow(unreachable_patterns)]
125 _ => false,
126 }
127}
128
129fn detect() -> Level {
130 for l in [Level::Avx2, Level::Sse42, Level::Neon] {
131 if available(l) {
132 return l;
133 }
134 }
135 Level::Scalar
136}
137
138#[inline]
140pub fn level() -> Level {
141 match LEVEL.load(Ordering::Relaxed) {
142 0 => {
143 let l = detect();
144 LEVEL.store(l as u8, Ordering::Relaxed);
145 l
146 }
147 n => from_u8(n),
148 }
149}
150
151pub fn force(level: Level) -> bool {
154 if !available(level) {
155 return false;
156 }
157 LEVEL.store(level as u8, Ordering::Relaxed);
158 true
159}
160
161#[inline]
164pub fn scan(buf: &[u8], class: &ByteClass) -> usize {
165 dispatch_scan(level(), buf, class)
166}
167
168pub fn scan_at(level: Level, buf: &[u8], class: &ByteClass) -> usize {
170 let level = if available(level) { level } else { Level::Scalar };
171 dispatch_scan(level, buf, class)
172}
173
174#[inline]
176pub fn find_byte(buf: &[u8], byte: u8) -> Option<usize> {
177 dispatch_find(level(), buf, byte)
178}
179
180pub fn find_byte_at(level: Level, buf: &[u8], byte: u8) -> Option<usize> {
182 let level = if available(level) { level } else { Level::Scalar };
183 dispatch_find(level, buf, byte)
184}
185
186#[inline]
192pub fn field_line(buf: &[u8]) -> (usize, usize, usize) {
193 match level() {
194 #[cfg(target_arch = "x86_64")]
195 Level::Avx2 => unsafe { x86::field_line_avx2(buf) },
197 #[cfg(target_arch = "x86_64")]
198 Level::Sse42 => unsafe { x86::field_line_sse(buf) },
200 #[cfg(target_arch = "aarch64")]
201 Level::Neon => unsafe { arm::field_line_neon(buf) },
203 _ => field_line_with(buf, scan_scalar),
204 }
205}
206
207#[inline(always)]
208fn field_line_with(buf: &[u8], scan: impl Fn(&[u8], &ByteClass) -> usize) -> (usize, usize, usize) {
209 let n = scan(buf, &TCHAR);
210 if buf.get(n) != Some(&b':') {
211 return (n, n, n);
212 }
213 let mut v = n + 1;
214 while v < buf.len() && (buf[v] == b' ' || buf[v] == b'\t') {
215 v += 1;
216 }
217 (n, v, v + scan(&buf[v..], &FIELD_VALUE))
218}
219
220fn scan_scalar(buf: &[u8], class: &ByteClass) -> usize {
221 buf.iter().position(|&b| !class.table[b as usize]).unwrap_or(buf.len())
222}
223
224#[inline]
226fn dispatch_scan(level: Level, buf: &[u8], class: &ByteClass) -> usize {
227 if buf.len() < 16 {
228 return scan_scalar(buf, class);
229 }
230 match level {
231 #[cfg(target_arch = "x86_64")]
232 Level::Sse42 => unsafe { x86::scan_sse(buf, class) },
235 #[cfg(target_arch = "x86_64")]
236 Level::Avx2 => unsafe { x86::scan_avx2(buf, class) },
239 #[cfg(target_arch = "aarch64")]
240 Level::Neon => unsafe { arm::scan_neon(buf, class) },
243 _ => scan_scalar(buf, class),
244 }
245}
246
247#[inline]
248fn dispatch_find(level: Level, buf: &[u8], byte: u8) -> Option<usize> {
249 if buf.len() < 16 {
250 return buf.iter().position(|&b| b == byte);
251 }
252 match level {
253 #[cfg(target_arch = "x86_64")]
254 Level::Sse42 => unsafe { x86::find_sse2(buf, byte) },
256 #[cfg(target_arch = "x86_64")]
257 Level::Avx2 => unsafe { x86::find_avx2(buf, byte) },
260 #[cfg(target_arch = "aarch64")]
261 Level::Neon => unsafe { arm::find_neon(buf, byte) },
264 _ => buf.iter().position(|&b| b == byte),
265 }
266}
267
268#[cfg(target_arch = "x86_64")]
269mod x86 {
270 use super::ByteClass;
271 use core::arch::x86_64::*;
272
273 macro_rules! notin128 {
275 ($v:expr, $lo_t:expr, $hi_t:expr, $high:expr) => {{
276 let v = $v;
277 let nib = _mm_set1_epi8(0x0f);
278 let lo = _mm_and_si128(v, nib);
279 let hi = _mm_and_si128(_mm_srli_epi16(v, 4), nib);
280 let bits = _mm_and_si128(_mm_shuffle_epi8($lo_t, lo), _mm_shuffle_epi8($hi_t, hi));
281 let mut out = _mm_cmpeq_epi8(bits, _mm_setzero_si128());
282 if $high {
283 out = _mm_andnot_si128(_mm_cmplt_epi8(v, _mm_setzero_si128()), out);
284 }
285 _mm_movemask_epi8(out) as u32
286 }};
287 }
288
289 macro_rules! hi_table {
290 () => {
291 _mm_setr_epi8(1, 2, 4, 8, 16, 32, 64, -128, 0, 0, 0, 0, 0, 0, 0, 0)
292 };
293 }
294
295 macro_rules! eq_mask128 {
296 ($p:expr, $needle:expr) => {
297 _mm_movemask_epi8(_mm_cmpeq_epi8(_mm_loadu_si128($p as *const __m128i), $needle)) as u32
298 };
299 }
300
301 #[target_feature(enable = "avx2")]
304 pub(super) unsafe fn field_line_avx2(buf: &[u8]) -> (usize, usize, usize) {
305 super::field_line_with(buf, |b, c| {
306 if b.len() < 16 {
307 return super::scan_scalar(b, c);
308 }
309 unsafe { scan_avx2(b, c) }
311 })
312 }
313
314 #[target_feature(enable = "ssse3,sse4.2")]
317 pub(super) unsafe fn field_line_sse(buf: &[u8]) -> (usize, usize, usize) {
318 super::field_line_with(buf, |b, c| {
319 if b.len() < 16 {
320 return super::scan_scalar(b, c);
321 }
322 unsafe { scan_sse(b, c) }
324 })
325 }
326
327 #[target_feature(enable = "ssse3,sse4.2")]
330 pub(super) unsafe fn scan_sse(buf: &[u8], class: &ByteClass) -> usize {
331 let n = buf.len();
332 let p = buf.as_ptr();
333 unsafe {
336 let lo_t = _mm_loadu_si128(class.lo.as_ptr() as *const __m128i);
337 let hi_t = hi_table!();
338 let mut i = 0;
339 while i + 16 <= n {
340 let m = notin128!(_mm_loadu_si128(p.add(i) as *const __m128i), lo_t, hi_t, class.high);
341 if m != 0 {
342 return i + m.trailing_zeros() as usize;
343 }
344 i += 16;
345 }
346 if i < n {
347 let start = n - 16;
349 let m = notin128!(_mm_loadu_si128(p.add(start) as *const __m128i), lo_t, hi_t, class.high)
350 >> (i - start);
351 if m != 0 {
352 return i + m.trailing_zeros() as usize;
353 }
354 }
355 n
356 }
357 }
358
359 #[target_feature(enable = "avx2")]
362 pub(super) unsafe fn scan_avx2(buf: &[u8], class: &ByteClass) -> usize {
363 let n = buf.len();
364 let p = buf.as_ptr();
365 unsafe {
368 let lo_t = _mm_loadu_si128(class.lo.as_ptr() as *const __m128i);
369 let hi_t = hi_table!();
370 let lo_t2 = _mm256_broadcastsi128_si256(lo_t);
371 let hi_t2 = _mm256_broadcastsi128_si256(hi_t);
372 let nib = _mm256_set1_epi8(0x0f);
373 let zero = _mm256_setzero_si256();
374 let mut i = 0;
375 while i + 32 <= n {
376 let v = _mm256_loadu_si256(p.add(i) as *const __m256i);
377 let lo = _mm256_and_si256(v, nib);
378 let hi = _mm256_and_si256(_mm256_srli_epi16(v, 4), nib);
379 let bits = _mm256_and_si256(_mm256_shuffle_epi8(lo_t2, lo), _mm256_shuffle_epi8(hi_t2, hi));
380 let mut out = _mm256_cmpeq_epi8(bits, zero);
381 if class.high {
382 out = _mm256_andnot_si256(_mm256_cmpgt_epi8(zero, v), out);
383 }
384 let m = _mm256_movemask_epi8(out) as u32;
385 if m != 0 {
386 return i + m.trailing_zeros() as usize;
387 }
388 i += 32;
389 }
390 if i + 16 <= n {
391 let m = notin128!(_mm_loadu_si128(p.add(i) as *const __m128i), lo_t, hi_t, class.high);
392 if m != 0 {
393 return i + m.trailing_zeros() as usize;
394 }
395 i += 16;
396 }
397 if i < n {
398 let start = n - 16;
399 let m = notin128!(_mm_loadu_si128(p.add(start) as *const __m128i), lo_t, hi_t, class.high)
400 >> (i - start);
401 if m != 0 {
402 return i + m.trailing_zeros() as usize;
403 }
404 }
405 n
406 }
407 }
408
409 #[target_feature(enable = "sse2")]
412 pub(super) unsafe fn find_sse2(buf: &[u8], byte: u8) -> Option<usize> {
413 let n = buf.len();
414 let p = buf.as_ptr();
415 unsafe {
417 let needle = _mm_set1_epi8(byte as i8);
418 let mut i = 0;
419 while i + 16 <= n {
420 let m = eq_mask128!(p.add(i), needle);
421 if m != 0 {
422 return Some(i + m.trailing_zeros() as usize);
423 }
424 i += 16;
425 }
426 if i < n {
427 let start = n - 16;
428 let m = eq_mask128!(p.add(start), needle) >> (i - start);
429 if m != 0 {
430 return Some(i + m.trailing_zeros() as usize);
431 }
432 }
433 None
434 }
435 }
436
437 #[target_feature(enable = "avx2")]
440 pub(super) unsafe fn find_avx2(buf: &[u8], byte: u8) -> Option<usize> {
441 let n = buf.len();
442 let p = buf.as_ptr();
443 unsafe {
446 let needle = _mm256_set1_epi8(byte as i8);
447 let mut i = 0;
448 while i + 32 <= n {
449 let m = _mm256_movemask_epi8(_mm256_cmpeq_epi8(_mm256_loadu_si256(p.add(i) as *const __m256i), needle))
450 as u32;
451 if m != 0 {
452 return Some(i + m.trailing_zeros() as usize);
453 }
454 i += 32;
455 }
456 let needle = _mm_set1_epi8(byte as i8);
457 if i + 16 <= n {
458 let m = eq_mask128!(p.add(i), needle);
459 if m != 0 {
460 return Some(i + m.trailing_zeros() as usize);
461 }
462 i += 16;
463 }
464 if i < n {
465 let start = n - 16;
466 let m = eq_mask128!(p.add(start), needle) >> (i - start);
467 if m != 0 {
468 return Some(i + m.trailing_zeros() as usize);
469 }
470 }
471 None
472 }
473 }
474}
475
476#[cfg(target_arch = "aarch64")]
477mod arm {
478 use super::ByteClass;
479 use core::arch::aarch64::*;
480
481 macro_rules! nibble_mask {
483 ($m:expr) => {
484 vget_lane_u64(vreinterpret_u64_u8(vshrn_n_u16(vreinterpretq_u16_u8($m), 4)), 0)
485 };
486 }
487
488 macro_rules! notin_neon {
489 ($v:expr, $lo_t:expr, $hi_t:expr, $high:expr) => {{
490 let v = $v;
491 let nib = vdupq_n_u8(0x0f);
492 let bits = vandq_u8(vqtbl1q_u8($lo_t, vandq_u8(v, nib)), vqtbl1q_u8($hi_t, vshrq_n_u8(v, 4)));
493 let mut out = vceqq_u8(bits, vdupq_n_u8(0));
494 if $high {
495 out = vbicq_u8(out, vcgeq_u8(v, vdupq_n_u8(0x80)));
496 }
497 nibble_mask!(out)
498 }};
499 }
500
501 const HI: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 0, 0, 0, 0, 0, 0, 0, 0];
502
503 #[target_feature(enable = "neon")]
506 pub(super) unsafe fn field_line_neon(buf: &[u8]) -> (usize, usize, usize) {
507 super::field_line_with(buf, |b, c| {
508 if b.len() < 16 {
509 return super::scan_scalar(b, c);
510 }
511 unsafe { scan_neon(b, c) }
513 })
514 }
515
516 #[target_feature(enable = "neon")]
519 pub(super) unsafe fn scan_neon(buf: &[u8], class: &ByteClass) -> usize {
520 let n = buf.len();
521 let p = buf.as_ptr();
522 unsafe {
525 let lo_t = vld1q_u8(class.lo.as_ptr());
526 let hi_t = vld1q_u8(HI.as_ptr());
527 let mut i = 0;
528 while i + 16 <= n {
529 let m = notin_neon!(vld1q_u8(p.add(i)), lo_t, hi_t, class.high);
530 if m != 0 {
531 return i + (m.trailing_zeros() / 4) as usize;
532 }
533 i += 16;
534 }
535 if i < n {
536 let start = n - 16;
537 let m = notin_neon!(vld1q_u8(p.add(start)), lo_t, hi_t, class.high) >> (4 * (i - start));
538 if m != 0 {
539 return i + (m.trailing_zeros() / 4) as usize;
540 }
541 }
542 n
543 }
544 }
545
546 #[target_feature(enable = "neon")]
549 pub(super) unsafe fn find_neon(buf: &[u8], byte: u8) -> Option<usize> {
550 let n = buf.len();
551 let p = buf.as_ptr();
552 unsafe {
554 let needle = vdupq_n_u8(byte);
555 let mut i = 0;
556 while i + 16 <= n {
557 let m = nibble_mask!(vceqq_u8(vld1q_u8(p.add(i)), needle));
558 if m != 0 {
559 return Some(i + (m.trailing_zeros() / 4) as usize);
560 }
561 i += 16;
562 }
563 if i < n {
564 let start = n - 16;
565 let m = nibble_mask!(vceqq_u8(vld1q_u8(p.add(start)), needle)) >> (4 * (i - start));
566 if m != 0 {
567 return Some(i + (m.trailing_zeros() / 4) as usize);
568 }
569 }
570 None
571 }
572 }
573}
574
575#[cfg(test)]
576mod tests {
577 use super::*;
578
579 struct Rng(u64);
580 impl Rng {
581 fn next(&mut self) -> u64 {
582 self.0 ^= self.0 << 13;
583 self.0 ^= self.0 >> 7;
584 self.0 ^= self.0 << 17;
585 self.0
586 }
587 }
588
589 fn levels() -> Vec<Level> {
590 [Level::Scalar, Level::Sse42, Level::Avx2, Level::Neon].into_iter().filter(|&l| available(l)).collect()
591 }
592
593 #[test]
594 fn classes_match_the_rfc_sets() {
595 for b in 0..=255u8 {
596 let tchar = b.is_ascii_alphanumeric() || b"!#$%&'*+-.^_`|~".contains(&b);
597 assert_eq!(TCHAR.contains(b), tchar, "{b}");
598 let fv = b == b'\t' || (0x20..0x7f).contains(&b) || b >= 0x80;
599 assert_eq!(FIELD_VALUE.contains(b), fv, "{b}");
600 let uri = b.is_ascii_alphanumeric() || b"-._~!$&'()*+,;=:@/?".contains(&b);
601 assert_eq!(URI.contains(b), uri, "{b}");
602 }
603 }
604
605 #[test]
606 fn simd_and_scalar_agree_on_random_inputs() {
607 let mut rng = Rng(0x9e37_79b9_7f4a_7c15);
608 let lv = levels();
609 let classes = [&TCHAR, &FIELD_VALUE, &URI];
610 let mut buf = Vec::new();
611 for round in 0..20_000 {
612 let len = (rng.next() % 100) as usize;
613 buf.clear();
614 let class = classes[round % 3];
616 while buf.len() < len {
617 let b = rng.next() as u8;
618 if class.contains(b) || rng.next() % 40 == 0 {
619 buf.push(b);
620 }
621 }
622 let needle = buf.get((rng.next() as usize) % (buf.len() + 1)).copied().unwrap_or(b'\r');
623 for &l in &lv {
624 for s in [0, 1, 3, 15, 17] {
626 if s <= buf.len() {
627 let b = &buf[s..];
628 assert_eq!(scan_at(l, b, class), scan_scalar(b, class), "{l:?} {b:?}");
629 assert_eq!(find_byte_at(l, b, needle), b.iter().position(|&x| x == needle), "{l:?}");
630 }
631 }
632 }
633 }
634 }
635
636 #[test]
637 fn every_byte_value_stops_correctly_at_every_position() {
638 for &l in &levels() {
639 for class in [&TCHAR, &FIELD_VALUE, &URI] {
640 for len in [16usize, 31, 32, 33, 64, 70] {
641 for pos in 0..len {
642 for b in 0..=255u8 {
643 let mut buf = vec![b'a'; len];
644 buf[pos] = b;
645 let want = if class.contains(b) { len } else { pos };
646 assert_eq!(scan_at(l, &buf, class), want, "{l:?} len {len} pos {pos} byte {b}");
647 }
648 let mut buf = vec![0u8; len];
649 buf[pos] = b'\n';
650 assert_eq!(find_byte_at(l, &buf, b'\n'), Some(pos));
651 }
652 }
653 }
654 }
655 }
656
657 #[test]
658 fn field_line_matches_two_scans() {
659 let mut rng = Rng(42);
660 let pool = b"abcXY-_:: \t\r\n\x00\x7f\xc3;=\"";
661 for _ in 0..20_000 {
662 let len = (rng.next() % 70) as usize;
663 let buf: Vec<u8> = (0..len)
664 .map(|_| {
665 if rng.next() % 3 == 0 {
666 pool[(rng.next() as usize) % pool.len()]
667 } else {
668 b'a' + (rng.next() % 26) as u8
669 }
670 })
671 .collect();
672 let want = field_line_with(&buf, scan_scalar);
673 assert_eq!(field_line(&buf), want, "{buf:?}");
674 let (n, vs, ve) = want;
675 assert_eq!(n, scan_scalar(&buf, &TCHAR));
676 if buf.get(n) == Some(&b':') {
677 assert!(vs > n && buf[n + 1..vs].iter().all(|&b| b == b' ' || b == b'\t'));
678 assert_eq!(ve, vs + scan_scalar(&buf[vs..], &FIELD_VALUE));
679 } else {
680 assert_eq!((vs, ve), (n, n));
681 }
682 }
683 assert_eq!(field_line(b"Host: example.com\r\n"), (4, 7, 18));
684 assert_eq!(field_line(b"Bad Name: x"), (3, 3, 3));
685 }
686
687 #[test]
688 fn force_and_level() {
689 let before = level();
690 assert!(available(before));
691 assert!(force(Level::Scalar));
692 assert_eq!(level(), Level::Scalar);
693 assert!(force(before));
694 assert_eq!(level(), before);
695 }
696}