Skip to main content

sark_core/
simd.rs

1pub enum HeaderNameOutcome {
2    Found { pos: usize, byte: u8 },
3    Invalid,
4    None,
5}
6
7fn is_header_name_byte(b: u8) -> bool {
8    b.is_ascii_alphanumeric()
9        || matches!(
10            b,
11            b'!' | b'#'
12                | b'$'
13                | b'%'
14                | b'&'
15                | b'\''
16                | b'*'
17                | b'+'
18                | b'-'
19                | b'.'
20                | b'^'
21                | b'_'
22                | b'`'
23                | b'|'
24                | b'~'
25        )
26}
27
28fn scan_header_name_scalar(bytes: &[u8], start: usize) -> HeaderNameOutcome {
29    let mut idx = start;
30    while idx < bytes.len() {
31        let b = bytes[idx];
32        if b == b':' || b == b'\r' {
33            return HeaderNameOutcome::Found { pos: idx, byte: b };
34        }
35        if !is_header_name_byte(b) {
36            return HeaderNameOutcome::Invalid;
37        }
38        idx += 1;
39    }
40    HeaderNameOutcome::None
41}
42
43cfg_select! {
44    target_arch = "aarch64" => {
45        pub fn scan_header_name(bytes: &[u8], start: usize) -> HeaderNameOutcome {
46            if start >= bytes.len() {
47                return HeaderNameOutcome::None;
48            }
49            // SAFETY: aarch64 baseline always includes the neon target feature.
50            unsafe { scan_header_name_neon(bytes, start) }
51        }
52
53        use core::arch::aarch64::{
54            uint8x16_t, vandq_u8, vceqq_u8, vcgeq_u8, vcleq_u8, vdupq_n_u8,
55            vget_lane_u64, vld1q_u8, vminvq_u8, vmvnq_u8, vorrq_u8,
56            vreinterpret_u64_u8, vreinterpretq_u16_u8, vshrn_n_u16,
57        };
58
59        #[target_feature(enable = "neon")]
60        fn first_match(mask: uint8x16_t) -> Option<usize> {
61            let narrow = vshrn_n_u16(vreinterpretq_u16_u8(mask), 4);
62            let bits = vget_lane_u64(vreinterpret_u64_u8(narrow), 0);
63            if bits == 0 {
64                None
65            } else {
66                Some((bits.trailing_zeros() as usize) >> 2)
67            }
68        }
69
70        #[target_feature(enable = "neon")]
71        fn valid_mask(chunk: uint8x16_t) -> uint8x16_t {
72            let upper = vandq_u8(
73                    vcgeq_u8(chunk, vdupq_n_u8(b'A')),
74                    vcleq_u8(chunk, vdupq_n_u8(b'Z')),
75                );
76            let lower = vandq_u8(
77                    vcgeq_u8(chunk, vdupq_n_u8(b'a')),
78                    vcleq_u8(chunk, vdupq_n_u8(b'z')),
79                );
80            let digit = vandq_u8(
81                    vcgeq_u8(chunk, vdupq_n_u8(b'0')),
82                    vcleq_u8(chunk, vdupq_n_u8(b'9')),
83                );
84            let alpha_num = vorrq_u8(vorrq_u8(upper, lower), digit);
85            let punct = vorrq_u8(
86                    vorrq_u8(
87                        vorrq_u8(
88                            vceqq_u8(chunk, vdupq_n_u8(b'!')),
89                            vceqq_u8(chunk, vdupq_n_u8(b'#')),
90                        ),
91                        vorrq_u8(
92                            vceqq_u8(chunk, vdupq_n_u8(b'$')),
93                            vceqq_u8(chunk, vdupq_n_u8(b'%')),
94                        ),
95                    ),
96                    vorrq_u8(
97                        vorrq_u8(
98                            vorrq_u8(
99                                vceqq_u8(chunk, vdupq_n_u8(b'&')),
100                                vceqq_u8(chunk, vdupq_n_u8(b'\'')),
101                            ),
102                            vorrq_u8(
103                                vceqq_u8(chunk, vdupq_n_u8(b'*')),
104                                vceqq_u8(chunk, vdupq_n_u8(b'+')),
105                            ),
106                        ),
107                        vorrq_u8(
108                            vorrq_u8(
109                                vorrq_u8(
110                                    vceqq_u8(chunk, vdupq_n_u8(b'-')),
111                                    vceqq_u8(chunk, vdupq_n_u8(b'.')),
112                                ),
113                                vorrq_u8(
114                                    vceqq_u8(chunk, vdupq_n_u8(b'^')),
115                                    vceqq_u8(chunk, vdupq_n_u8(b'_')),
116                                ),
117                            ),
118                            vorrq_u8(
119                                vceqq_u8(chunk, vdupq_n_u8(b'`')),
120                                vorrq_u8(
121                                    vceqq_u8(chunk, vdupq_n_u8(b'|')),
122                                    vceqq_u8(chunk, vdupq_n_u8(b'~')),
123                                ),
124                            ),
125                        ),
126                    ),
127                );
128            vorrq_u8(alpha_num, punct)
129        }
130
131        #[target_feature(enable = "neon")]
132        fn scan_header_name_neon(bytes: &[u8], start: usize) -> HeaderNameOutcome {
133            let mut idx = start;
134            let len = bytes.len();
135            let colon_v = vdupq_n_u8(b':');
136            let cr_v = vdupq_n_u8(b'\r');
137            while idx + 16 <= len {
138                // SAFETY: idx + 16 <= len bounds the 16-byte load inside bytes.
139                let chunk = unsafe { vld1q_u8(bytes.as_ptr().add(idx)) };
140                let hit = vorrq_u8(vceqq_u8(chunk, colon_v), vceqq_u8(chunk, cr_v));
141                let valid_or_hit = vorrq_u8(valid_mask(chunk), hit);
142                if vminvq_u8(valid_or_hit) != 0xFF {
143                    let invalid = vmvnq_u8(valid_or_hit);
144                    let hit_pos = first_match(hit);
145                    let invalid_pos = first_match(invalid);
146                    return match (hit_pos, invalid_pos) {
147                        (Some(hit_off), Some(invalid_off)) if invalid_off < hit_off => {
148                            HeaderNameOutcome::Invalid
149                        }
150                        (Some(hit_off), _) => {
151                            let pos = idx + hit_off;
152                            HeaderNameOutcome::Found { pos, byte: bytes[pos] }
153                        }
154                        (None, Some(_)) => HeaderNameOutcome::Invalid,
155                        (None, None) => unreachable!(),
156                    };
157                }
158                if let Some(hit_off) = first_match(hit) {
159                    let pos = idx + hit_off;
160                    return HeaderNameOutcome::Found { pos, byte: bytes[pos] };
161                }
162                idx += 16;
163            }
164            scan_header_name_scalar(bytes, idx)
165        }
166    }
167    _ => {
168        pub fn scan_header_name(bytes: &[u8], start: usize) -> HeaderNameOutcome {
169            if start >= bytes.len() {
170                return HeaderNameOutcome::None;
171            }
172            scan_header_name_scalar(bytes, start)
173        }
174    }
175}