rsonpath/crates/rsonpath-lib/src/classification/depth/shared/vector_128.rs at ab4ff3ac9cdcae7c2b52e9ea075ff7d4938f4669 · rsonquery/rsonpath · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
use crate::classification::structural::BracketType;

#[cfg(target_arch = "wasm32")]
use core::arch::wasm32::*;
#[cfg(target_arch = "x86")]
use core::arch::x86::*;
#[cfg(target_arch = "x86_64")]
use core::arch::x86_64::*;

pub(crate) struct DelimiterClassifierImpl128 {
    opening: i8,
}

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
impl DelimiterClassifierImpl128 {
    pub(crate) fn new(opening: BracketType) -> Self {
        let opening = match opening {
            BracketType::Square => b'[',
            BracketType::Curly => b'{',
        };

        Self { opening: opening as i8 }
    }

    #[inline(always)]
    unsafe fn opening_mask(&self) -> __m128i {
        _mm_set1_epi8(self.opening)
    }

    #[inline(always)]
    unsafe fn closing_mask(&self) -> __m128i {
        _mm_set1_epi8(self.opening + 2)
    }

    #[target_feature(enable = "sse2")]
    #[inline]
    pub(crate) unsafe fn get_opening_and_closing_masks(&self, bytes: &[u8]) -> (u16, u16) {
        assert_eq!(16, bytes.len());
        // SAFETY: target_feature invariant
        unsafe {
            let byte_vector = _mm_loadu_si128(bytes.as_ptr().cast::<__m128i>());
            let opening_brace_cmp = _mm_cmpeq_epi8(byte_vector, self.opening_mask());
            let closing_brace_cmp = _mm_cmpeq_epi8(byte_vector, self.closing_mask());
            let opening_mask = _mm_movemask_epi8(opening_brace_cmp) as u16;
            let closing_mask = _mm_movemask_epi8(closing_brace_cmp) as u16;

            (opening_mask, closing_mask)
        }
    }
}

#[cfg(target_arch = "wasm32")]
impl DelimiterClassifierImpl128 {
    pub(crate) fn new(opening: BracketType) -> Self {
        let opening = match opening {
            BracketType::Square => b'[',
            BracketType::Curly => b'{',
        };
        Self { opening: opening as i8 }
    }

    #[inline(always)]
    unsafe fn opening_mask(&self) -> v128 {
        i8x16_splat(self.opening)
    }

    #[inline(always)]
    unsafe fn closing_mask(&self) -> v128 {
        i8x16_splat(self.opening + 2)
    }

    #[target_feature(enable = "simd128")]
    #[inline]
    pub(crate) unsafe fn get_opening_and_closing_masks(&self, bytes: &[u8]) -> (u16, u16) {
        assert_eq!(16, bytes.len());

        let byte_vector = v128_load(bytes.as_ptr() as *const v128);

        let opening_brace_cmp = i8x16_eq(byte_vector, self.opening_mask());
        let closing_brace_cmp = i8x16_eq(byte_vector, self.closing_mask());

        let opening_mask = i8x16_bitmask(opening_brace_cmp) as u16;
        let closing_mask = i8x16_bitmask(closing_brace_cmp) as u16;

        (opening_mask, closing_mask)
    }
}