Skip to main content

sark_core/
simd.rs

1#[derive(Debug, PartialEq, Eq)]
2pub enum HeaderNameOutcome {
3    Found { pos: usize, byte: u8 },
4    Invalid,
5    None,
6}
7
8#[derive(Debug, PartialEq, Eq)]
9pub enum HeaderValueOutcome {
10    Found { pos: usize },
11    Invalid,
12    None,
13}
14
15fn is_header_name_byte(b: u8) -> bool {
16    b.is_ascii_alphanumeric()
17        || matches!(
18            b,
19            b'!' | b'#'
20                | b'$'
21                | b'%'
22                | b'&'
23                | b'\''
24                | b'*'
25                | b'+'
26                | b'-'
27                | b'.'
28                | b'^'
29                | b'_'
30                | b'`'
31                | b'|'
32                | b'~'
33        )
34}
35
36fn scan_header_name_scalar(bytes: &[u8], start: usize) -> HeaderNameOutcome {
37    let mut idx = start;
38    while idx < bytes.len() {
39        let b = bytes[idx];
40        if b == b':' || b == b'\r' {
41            return HeaderNameOutcome::Found { pos: idx, byte: b };
42        }
43        if !is_header_name_byte(b) {
44            return HeaderNameOutcome::Invalid;
45        }
46        idx += 1;
47    }
48    HeaderNameOutcome::None
49}
50
51fn scan_header_value_scalar(bytes: &[u8], start: usize) -> HeaderValueOutcome {
52    let mut idx = start;
53    while idx < bytes.len() {
54        let byte = bytes[idx];
55        if byte == b'\r' {
56            if idx + 1 == bytes.len() {
57                return HeaderValueOutcome::None;
58            }
59            return if bytes[idx + 1] == b'\n' {
60                HeaderValueOutcome::Found { pos: idx }
61            } else {
62                HeaderValueOutcome::Invalid
63            };
64        }
65        if (byte < 0x20 && byte != b'\t') || byte == 0x7f {
66            return HeaderValueOutcome::Invalid;
67        }
68        idx += 1;
69    }
70    HeaderValueOutcome::None
71}
72
73fn request_target_is_valid_scalar(bytes: &[u8]) -> bool {
74    !bytes.iter().any(|&byte| byte <= 0x20 || byte == 0x7f)
75}
76
77cfg_select! {
78    target_arch = "aarch64" => {
79        pub fn scan_header_name(bytes: &[u8], start: usize) -> HeaderNameOutcome {
80            if start >= bytes.len() {
81                return HeaderNameOutcome::None;
82            }
83            unsafe { scan_header_name_neon(bytes, start) }
84        }
85
86        pub fn scan_header_value(bytes: &[u8], start: usize) -> HeaderValueOutcome {
87            if start >= bytes.len() {
88                return HeaderValueOutcome::None;
89            }
90            unsafe { scan_header_value_neon(bytes, start) }
91        }
92
93        #[inline]
94        pub fn request_target_is_valid(bytes: &[u8]) -> bool {
95            unsafe { request_target_is_valid_neon(bytes) }
96        }
97
98        use core::arch::aarch64::{
99            uint8x16_t, vandq_u8, vceqq_u8, vcgeq_u8, vcleq_u8, vdupq_n_u8,
100            vget_lane_u64, vld1q_u8, vminvq_u8, vmvnq_u8, vorrq_u8,
101            vreinterpret_u64_u8, vreinterpretq_u16_u8, vshrn_n_u16,
102        };
103
104        #[target_feature(enable = "neon")]
105        fn first_match(mask: uint8x16_t) -> Option<usize> {
106            let narrow = vshrn_n_u16(vreinterpretq_u16_u8(mask), 4);
107            let bits = vget_lane_u64(vreinterpret_u64_u8(narrow), 0);
108            if bits == 0 {
109                None
110            } else {
111                Some((bits.trailing_zeros() as usize) >> 2)
112            }
113        }
114
115        #[target_feature(enable = "neon")]
116        fn valid_mask(chunk: uint8x16_t) -> uint8x16_t {
117            let upper = vandq_u8(
118                    vcgeq_u8(chunk, vdupq_n_u8(b'A')),
119                    vcleq_u8(chunk, vdupq_n_u8(b'Z')),
120                );
121            let lower = vandq_u8(
122                    vcgeq_u8(chunk, vdupq_n_u8(b'a')),
123                    vcleq_u8(chunk, vdupq_n_u8(b'z')),
124                );
125            let digit = vandq_u8(
126                    vcgeq_u8(chunk, vdupq_n_u8(b'0')),
127                    vcleq_u8(chunk, vdupq_n_u8(b'9')),
128                );
129            let alpha_num = vorrq_u8(vorrq_u8(upper, lower), digit);
130            let punct = vorrq_u8(
131                    vorrq_u8(
132                        vorrq_u8(
133                            vceqq_u8(chunk, vdupq_n_u8(b'!')),
134                            vceqq_u8(chunk, vdupq_n_u8(b'#')),
135                        ),
136                        vorrq_u8(
137                            vceqq_u8(chunk, vdupq_n_u8(b'$')),
138                            vceqq_u8(chunk, vdupq_n_u8(b'%')),
139                        ),
140                    ),
141                    vorrq_u8(
142                        vorrq_u8(
143                            vorrq_u8(
144                                vceqq_u8(chunk, vdupq_n_u8(b'&')),
145                                vceqq_u8(chunk, vdupq_n_u8(b'\'')),
146                            ),
147                            vorrq_u8(
148                                vceqq_u8(chunk, vdupq_n_u8(b'*')),
149                                vceqq_u8(chunk, vdupq_n_u8(b'+')),
150                            ),
151                        ),
152                        vorrq_u8(
153                            vorrq_u8(
154                                vorrq_u8(
155                                    vceqq_u8(chunk, vdupq_n_u8(b'-')),
156                                    vceqq_u8(chunk, vdupq_n_u8(b'.')),
157                                ),
158                                vorrq_u8(
159                                    vceqq_u8(chunk, vdupq_n_u8(b'^')),
160                                    vceqq_u8(chunk, vdupq_n_u8(b'_')),
161                                ),
162                            ),
163                            vorrq_u8(
164                                vceqq_u8(chunk, vdupq_n_u8(b'`')),
165                                vorrq_u8(
166                                    vceqq_u8(chunk, vdupq_n_u8(b'|')),
167                                    vceqq_u8(chunk, vdupq_n_u8(b'~')),
168                                ),
169                            ),
170                        ),
171                    ),
172                );
173            vorrq_u8(alpha_num, punct)
174        }
175
176        #[target_feature(enable = "neon")]
177        fn scan_header_name_neon(bytes: &[u8], start: usize) -> HeaderNameOutcome {
178            let mut idx = start;
179            let len = bytes.len();
180            let colon_v = vdupq_n_u8(b':');
181            let cr_v = vdupq_n_u8(b'\r');
182            while idx + 16 <= len {
183                let chunk = unsafe { vld1q_u8(bytes.as_ptr().add(idx)) };
184                let hit = vorrq_u8(vceqq_u8(chunk, colon_v), vceqq_u8(chunk, cr_v));
185                let valid_or_hit = vorrq_u8(valid_mask(chunk), hit);
186                if vminvq_u8(valid_or_hit) != 0xFF {
187                    let invalid = vmvnq_u8(valid_or_hit);
188                    let hit_pos = first_match(hit);
189                    let invalid_pos = first_match(invalid);
190                    return match (hit_pos, invalid_pos) {
191                        (Some(hit_off), Some(invalid_off)) if invalid_off < hit_off => {
192                            HeaderNameOutcome::Invalid
193                        }
194                        (Some(hit_off), _) => {
195                            let pos = idx + hit_off;
196                            HeaderNameOutcome::Found { pos, byte: bytes[pos] }
197                        }
198                        (None, Some(_)) => HeaderNameOutcome::Invalid,
199                        (None, None) => unreachable!(),
200                    };
201                }
202                if let Some(hit_off) = first_match(hit) {
203                    let pos = idx + hit_off;
204                    return HeaderNameOutcome::Found { pos, byte: bytes[pos] };
205                }
206                idx += 16;
207            }
208            scan_header_name_scalar(bytes, idx)
209        }
210
211        #[target_feature(enable = "neon")]
212        fn scan_header_value_neon(bytes: &[u8], start: usize) -> HeaderValueOutcome {
213            let mut idx = start;
214            let len = bytes.len();
215            let ctl_max = vdupq_n_u8(0x1f);
216            let tab = vdupq_n_u8(b'\t');
217            let del = vdupq_n_u8(0x7f);
218            while idx + 16 <= len {
219                let chunk = unsafe { vld1q_u8(bytes.as_ptr().add(idx)) };
220                let low_ctl = vandq_u8(vcleq_u8(chunk, ctl_max), vmvnq_u8(vceqq_u8(chunk, tab)));
221                let special = vorrq_u8(low_ctl, vceqq_u8(chunk, del));
222                if let Some(offset) = first_match(special) {
223                    return scan_header_value_scalar(bytes, idx + offset);
224                }
225                idx += 16;
226            }
227            scan_header_value_scalar(bytes, idx)
228        }
229
230        #[target_feature(enable = "neon")]
231        fn request_target_is_valid_neon(bytes: &[u8]) -> bool {
232            let mut chunks = bytes.chunks_exact(16);
233            let space = vdupq_n_u8(0x20);
234            let del = vdupq_n_u8(0x7f);
235            for chunk in &mut chunks {
236                let vector = unsafe { vld1q_u8(chunk.as_ptr()) };
237                let invalid = vorrq_u8(vcleq_u8(vector, space), vceqq_u8(vector, del));
238                if first_match(invalid).is_some() {
239                    return false;
240                }
241            }
242            request_target_is_valid_scalar(chunks.remainder())
243        }
244    }
245    target_arch = "x86_64" => {
246        pub fn scan_header_name(bytes: &[u8], start: usize) -> HeaderNameOutcome {
247            if start >= bytes.len() {
248                return HeaderNameOutcome::None;
249            }
250            unsafe { scan_header_name_sse2(bytes, start) }
251        }
252
253        pub fn scan_header_value(bytes: &[u8], start: usize) -> HeaderValueOutcome {
254            if start >= bytes.len() {
255                return HeaderValueOutcome::None;
256            }
257            unsafe { scan_header_value_sse2(bytes, start) }
258        }
259
260        #[inline]
261        pub fn request_target_is_valid(bytes: &[u8]) -> bool {
262            unsafe { request_target_is_valid_sse2(bytes) }
263        }
264
265        use core::arch::x86_64::{
266            __m128i, _mm_andnot_si128, _mm_cmpeq_epi8, _mm_loadu_si128,
267            _mm_max_epu8, _mm_min_epu8, _mm_movemask_epi8, _mm_or_si128, _mm_set1_epi8,
268        };
269
270        const LANE_MASK: u32 = 0xFFFF;
271
272        #[target_feature(enable = "sse2")]
273        fn lane_eq(chunk: __m128i, byte: u8) -> __m128i {
274            _mm_cmpeq_epi8(chunk, _mm_set1_epi8(byte as i8))
275        }
276
277        #[target_feature(enable = "sse2")]
278        fn in_range(chunk: __m128i, lo: u8, hi: u8) -> __m128i {
279            let clamped =
280                _mm_min_epu8(_mm_max_epu8(chunk, _mm_set1_epi8(lo as i8)), _mm_set1_epi8(hi as i8));
281            _mm_cmpeq_epi8(chunk, clamped)
282        }
283
284        #[target_feature(enable = "sse2")]
285        fn valid_mask(chunk: __m128i) -> __m128i {
286            let alpha_num = _mm_or_si128(
287                _mm_or_si128(in_range(chunk, b'A', b'Z'), in_range(chunk, b'a', b'z')),
288                in_range(chunk, b'0', b'9'),
289            );
290            let punct = _mm_or_si128(
291                _mm_or_si128(
292                    _mm_or_si128(
293                        _mm_or_si128(lane_eq(chunk, b'!'), lane_eq(chunk, b'#')),
294                        _mm_or_si128(lane_eq(chunk, b'$'), lane_eq(chunk, b'%')),
295                    ),
296                    _mm_or_si128(
297                        _mm_or_si128(lane_eq(chunk, b'&'), lane_eq(chunk, b'\'')),
298                        _mm_or_si128(lane_eq(chunk, b'*'), lane_eq(chunk, b'+')),
299                    ),
300                ),
301                _mm_or_si128(
302                    _mm_or_si128(
303                        _mm_or_si128(lane_eq(chunk, b'-'), lane_eq(chunk, b'.')),
304                        _mm_or_si128(lane_eq(chunk, b'^'), lane_eq(chunk, b'_')),
305                    ),
306                    _mm_or_si128(
307                        lane_eq(chunk, b'`'),
308                        _mm_or_si128(lane_eq(chunk, b'|'), lane_eq(chunk, b'~')),
309                    ),
310                ),
311            );
312            _mm_or_si128(alpha_num, punct)
313        }
314
315        #[target_feature(enable = "sse2")]
316        fn scan_header_name_sse2(bytes: &[u8], start: usize) -> HeaderNameOutcome {
317            let mut idx = start;
318            let len = bytes.len();
319            while idx + 16 <= len {
320                let chunk = unsafe { _mm_loadu_si128(bytes.as_ptr().add(idx) as *const __m128i) };
321                let hit_v = _mm_or_si128(lane_eq(chunk, b':'), lane_eq(chunk, b'\r'));
322                let valid_or_hit = _mm_or_si128(valid_mask(chunk), hit_v);
323                let hit = _mm_movemask_epi8(hit_v) as u32;
324                let invalid = _mm_movemask_epi8(valid_or_hit) as u32 ^ LANE_MASK;
325                if invalid != 0 {
326                    let inv_off = invalid.trailing_zeros();
327                    if hit != 0 && hit.trailing_zeros() < inv_off {
328                        let pos = idx + hit.trailing_zeros() as usize;
329                        return HeaderNameOutcome::Found { pos, byte: bytes[pos] };
330                    }
331                    return HeaderNameOutcome::Invalid;
332                }
333                if hit != 0 {
334                    let pos = idx + hit.trailing_zeros() as usize;
335                    return HeaderNameOutcome::Found { pos, byte: bytes[pos] };
336                }
337                idx += 16;
338            }
339            scan_header_name_scalar(bytes, idx)
340        }
341
342        #[target_feature(enable = "sse2")]
343        fn scan_header_value_sse2(bytes: &[u8], start: usize) -> HeaderValueOutcome {
344            let mut idx = start;
345            let len = bytes.len();
346            while idx + 16 <= len {
347                let chunk = unsafe { _mm_loadu_si128(bytes.as_ptr().add(idx) as *const __m128i) };
348                let low_ctl = _mm_andnot_si128(
349                    lane_eq(chunk, b'\t'),
350                    _mm_cmpeq_epi8(
351                        _mm_min_epu8(chunk, _mm_set1_epi8(0x1f)),
352                        chunk,
353                    ),
354                );
355                let special = _mm_or_si128(low_ctl, lane_eq(chunk, 0x7f));
356                let mask = _mm_movemask_epi8(special) as u32;
357                if mask != 0 {
358                    return scan_header_value_scalar(bytes, idx + mask.trailing_zeros() as usize);
359                }
360                idx += 16;
361            }
362            scan_header_value_scalar(bytes, idx)
363        }
364
365        #[target_feature(enable = "sse2")]
366        fn request_target_is_valid_sse2(bytes: &[u8]) -> bool {
367            let mut chunks = bytes.chunks_exact(16);
368            for chunk in &mut chunks {
369                let vector = unsafe { _mm_loadu_si128(chunk.as_ptr() as *const __m128i) };
370                let low = _mm_cmpeq_epi8(
371                    _mm_min_epu8(vector, _mm_set1_epi8(0x20)),
372                    vector,
373                );
374                let invalid = _mm_or_si128(low, lane_eq(vector, 0x7f));
375                if _mm_movemask_epi8(invalid) != 0 {
376                    return false;
377                }
378            }
379            request_target_is_valid_scalar(chunks.remainder())
380        }
381    }
382    _ => {
383        pub fn scan_header_name(bytes: &[u8], start: usize) -> HeaderNameOutcome {
384            if start >= bytes.len() {
385                return HeaderNameOutcome::None;
386            }
387            scan_header_name_scalar(bytes, start)
388        }
389
390
391        pub fn scan_header_value(bytes: &[u8], start: usize) -> HeaderValueOutcome {
392            if start >= bytes.len() {
393                return HeaderValueOutcome::None;
394            }
395            scan_header_value_scalar(bytes, start)
396        }
397
398        #[inline]
399        pub fn request_target_is_valid(bytes: &[u8]) -> bool {
400            request_target_is_valid_scalar(bytes)
401        }
402    }
403}