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
8fn is_header_name_byte(b: u8) -> bool {
9    b.is_ascii_alphanumeric()
10        || matches!(
11            b,
12            b'!' | 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                | b'~'
26        )
27}
28
29fn scan_header_name_scalar(bytes: &[u8], start: usize) -> HeaderNameOutcome {
30    let mut idx = start;
31    while idx < bytes.len() {
32        let b = bytes[idx];
33        if b == b':' || b == b'\r' {
34            return HeaderNameOutcome::Found { pos: idx, byte: b };
35        }
36        if !is_header_name_byte(b) {
37            return HeaderNameOutcome::Invalid;
38        }
39        idx += 1;
40    }
41    HeaderNameOutcome::None
42}
43
44cfg_select! {
45    target_arch = "aarch64" => {
46        pub fn scan_header_name(bytes: &[u8], start: usize) -> HeaderNameOutcome {
47            if start >= bytes.len() {
48                return HeaderNameOutcome::None;
49            }
50            // SAFETY: aarch64 baseline always includes the neon target feature.
51            unsafe { scan_header_name_neon(bytes, start) }
52        }
53
54        use core::arch::aarch64::{
55            uint8x16_t, vandq_u8, vceqq_u8, vcgeq_u8, vcleq_u8, vdupq_n_u8,
56            vget_lane_u64, vld1q_u8, vminvq_u8, vmvnq_u8, vorrq_u8,
57            vreinterpret_u64_u8, vreinterpretq_u16_u8, vshrn_n_u16,
58        };
59
60        #[target_feature(enable = "neon")]
61        fn first_match(mask: uint8x16_t) -> Option<usize> {
62            let narrow = vshrn_n_u16(vreinterpretq_u16_u8(mask), 4);
63            let bits = vget_lane_u64(vreinterpret_u64_u8(narrow), 0);
64            if bits == 0 {
65                None
66            } else {
67                Some((bits.trailing_zeros() as usize) >> 2)
68            }
69        }
70
71        #[target_feature(enable = "neon")]
72        fn valid_mask(chunk: uint8x16_t) -> uint8x16_t {
73            let upper = vandq_u8(
74                    vcgeq_u8(chunk, vdupq_n_u8(b'A')),
75                    vcleq_u8(chunk, vdupq_n_u8(b'Z')),
76                );
77            let lower = vandq_u8(
78                    vcgeq_u8(chunk, vdupq_n_u8(b'a')),
79                    vcleq_u8(chunk, vdupq_n_u8(b'z')),
80                );
81            let digit = vandq_u8(
82                    vcgeq_u8(chunk, vdupq_n_u8(b'0')),
83                    vcleq_u8(chunk, vdupq_n_u8(b'9')),
84                );
85            let alpha_num = vorrq_u8(vorrq_u8(upper, lower), digit);
86            let punct = vorrq_u8(
87                    vorrq_u8(
88                        vorrq_u8(
89                            vceqq_u8(chunk, vdupq_n_u8(b'!')),
90                            vceqq_u8(chunk, vdupq_n_u8(b'#')),
91                        ),
92                        vorrq_u8(
93                            vceqq_u8(chunk, vdupq_n_u8(b'$')),
94                            vceqq_u8(chunk, vdupq_n_u8(b'%')),
95                        ),
96                    ),
97                    vorrq_u8(
98                        vorrq_u8(
99                            vorrq_u8(
100                                vceqq_u8(chunk, vdupq_n_u8(b'&')),
101                                vceqq_u8(chunk, vdupq_n_u8(b'\'')),
102                            ),
103                            vorrq_u8(
104                                vceqq_u8(chunk, vdupq_n_u8(b'*')),
105                                vceqq_u8(chunk, vdupq_n_u8(b'+')),
106                            ),
107                        ),
108                        vorrq_u8(
109                            vorrq_u8(
110                                vorrq_u8(
111                                    vceqq_u8(chunk, vdupq_n_u8(b'-')),
112                                    vceqq_u8(chunk, vdupq_n_u8(b'.')),
113                                ),
114                                vorrq_u8(
115                                    vceqq_u8(chunk, vdupq_n_u8(b'^')),
116                                    vceqq_u8(chunk, vdupq_n_u8(b'_')),
117                                ),
118                            ),
119                            vorrq_u8(
120                                vceqq_u8(chunk, vdupq_n_u8(b'`')),
121                                vorrq_u8(
122                                    vceqq_u8(chunk, vdupq_n_u8(b'|')),
123                                    vceqq_u8(chunk, vdupq_n_u8(b'~')),
124                                ),
125                            ),
126                        ),
127                    ),
128                );
129            vorrq_u8(alpha_num, punct)
130        }
131
132        #[target_feature(enable = "neon")]
133        fn scan_header_name_neon(bytes: &[u8], start: usize) -> HeaderNameOutcome {
134            let mut idx = start;
135            let len = bytes.len();
136            let colon_v = vdupq_n_u8(b':');
137            let cr_v = vdupq_n_u8(b'\r');
138            while idx + 16 <= len {
139                // SAFETY: idx + 16 <= len bounds the 16-byte load inside bytes.
140                let chunk = unsafe { vld1q_u8(bytes.as_ptr().add(idx)) };
141                let hit = vorrq_u8(vceqq_u8(chunk, colon_v), vceqq_u8(chunk, cr_v));
142                let valid_or_hit = vorrq_u8(valid_mask(chunk), hit);
143                if vminvq_u8(valid_or_hit) != 0xFF {
144                    let invalid = vmvnq_u8(valid_or_hit);
145                    let hit_pos = first_match(hit);
146                    let invalid_pos = first_match(invalid);
147                    return match (hit_pos, invalid_pos) {
148                        (Some(hit_off), Some(invalid_off)) if invalid_off < hit_off => {
149                            HeaderNameOutcome::Invalid
150                        }
151                        (Some(hit_off), _) => {
152                            let pos = idx + hit_off;
153                            HeaderNameOutcome::Found { pos, byte: bytes[pos] }
154                        }
155                        (None, Some(_)) => HeaderNameOutcome::Invalid,
156                        (None, None) => unreachable!(),
157                    };
158                }
159                if let Some(hit_off) = first_match(hit) {
160                    let pos = idx + hit_off;
161                    return HeaderNameOutcome::Found { pos, byte: bytes[pos] };
162                }
163                idx += 16;
164            }
165            scan_header_name_scalar(bytes, idx)
166        }
167    }
168    target_arch = "x86_64" => {
169        pub fn scan_header_name(bytes: &[u8], start: usize) -> HeaderNameOutcome {
170            if start >= bytes.len() {
171                return HeaderNameOutcome::None;
172            }
173            // SAFETY: sse2 is guaranteed on the x86_64 baseline.
174            unsafe { scan_header_name_sse2(bytes, start) }
175        }
176
177        use core::arch::x86_64::{
178            __m128i, _mm_cmpeq_epi8, _mm_loadu_si128, _mm_max_epu8, _mm_min_epu8,
179            _mm_movemask_epi8, _mm_or_si128, _mm_set1_epi8,
180        };
181
182        const LANE_MASK: u32 = 0xFFFF;
183
184        #[target_feature(enable = "sse2")]
185        fn lane_eq(chunk: __m128i, byte: u8) -> __m128i {
186            _mm_cmpeq_epi8(chunk, _mm_set1_epi8(byte as i8))
187        }
188
189        #[target_feature(enable = "sse2")]
190        fn in_range(chunk: __m128i, lo: u8, hi: u8) -> __m128i {
191            let clamped =
192                _mm_min_epu8(_mm_max_epu8(chunk, _mm_set1_epi8(lo as i8)), _mm_set1_epi8(hi as i8));
193            _mm_cmpeq_epi8(chunk, clamped)
194        }
195
196        #[target_feature(enable = "sse2")]
197        fn valid_mask(chunk: __m128i) -> __m128i {
198            let alpha_num = _mm_or_si128(
199                _mm_or_si128(in_range(chunk, b'A', b'Z'), in_range(chunk, b'a', b'z')),
200                in_range(chunk, b'0', b'9'),
201            );
202            let punct = _mm_or_si128(
203                _mm_or_si128(
204                    _mm_or_si128(
205                        _mm_or_si128(lane_eq(chunk, b'!'), lane_eq(chunk, b'#')),
206                        _mm_or_si128(lane_eq(chunk, b'$'), lane_eq(chunk, b'%')),
207                    ),
208                    _mm_or_si128(
209                        _mm_or_si128(lane_eq(chunk, b'&'), lane_eq(chunk, b'\'')),
210                        _mm_or_si128(lane_eq(chunk, b'*'), lane_eq(chunk, b'+')),
211                    ),
212                ),
213                _mm_or_si128(
214                    _mm_or_si128(
215                        _mm_or_si128(lane_eq(chunk, b'-'), lane_eq(chunk, b'.')),
216                        _mm_or_si128(lane_eq(chunk, b'^'), lane_eq(chunk, b'_')),
217                    ),
218                    _mm_or_si128(
219                        lane_eq(chunk, b'`'),
220                        _mm_or_si128(lane_eq(chunk, b'|'), lane_eq(chunk, b'~')),
221                    ),
222                ),
223            );
224            _mm_or_si128(alpha_num, punct)
225        }
226
227        #[target_feature(enable = "sse2")]
228        fn scan_header_name_sse2(bytes: &[u8], start: usize) -> HeaderNameOutcome {
229            let mut idx = start;
230            let len = bytes.len();
231            while idx + 16 <= len {
232                // SAFETY: idx + 16 <= len bounds the 16-byte load inside bytes.
233                let chunk = unsafe { _mm_loadu_si128(bytes.as_ptr().add(idx) as *const __m128i) };
234                let hit_v = _mm_or_si128(lane_eq(chunk, b':'), lane_eq(chunk, b'\r'));
235                let valid_or_hit = _mm_or_si128(valid_mask(chunk), hit_v);
236                let hit = _mm_movemask_epi8(hit_v) as u32;
237                let invalid = _mm_movemask_epi8(valid_or_hit) as u32 ^ LANE_MASK;
238                if invalid != 0 {
239                    let inv_off = invalid.trailing_zeros();
240                    if hit != 0 && hit.trailing_zeros() < inv_off {
241                        let pos = idx + hit.trailing_zeros() as usize;
242                        return HeaderNameOutcome::Found { pos, byte: bytes[pos] };
243                    }
244                    return HeaderNameOutcome::Invalid;
245                }
246                if hit != 0 {
247                    let pos = idx + hit.trailing_zeros() as usize;
248                    return HeaderNameOutcome::Found { pos, byte: bytes[pos] };
249                }
250                idx += 16;
251            }
252            scan_header_name_scalar(bytes, idx)
253        }
254    }
255    _ => {
256        pub fn scan_header_name(bytes: &[u8], start: usize) -> HeaderNameOutcome {
257            if start >= bytes.len() {
258                return HeaderNameOutcome::None;
259            }
260            scan_header_name_scalar(bytes, start)
261        }
262    }
263}