basekit 0.3.0

Universal library for encoding in various bases
Documentation 
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
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

use core::arch::x86_64::*;

use super::super::decode_full_group_into::decode_full_group_into;
use crate::base64::config::Base64DecodeConfig;
use crate::base64::error::Base64Error;

/// Decodes full 4-group blocks (src.len() is a multiple of 16, no padding).
/// Each 16 base64 characters → 12 output bytes.
/// Tail groups and padding are handled by the caller.
#[target_feature(enable = "ssse3")]
#[inline]
#[allow(unsafe_op_in_unsafe_fn)]
pub(crate) unsafe fn ssse3_decode_full_groups_into(
    config: &Base64DecodeConfig,
    dst: &mut [u8],
    src: &[u8],
) -> Result<usize, Base64Error> {
    debug_assert_eq!(src.len() % 16, 0);

    let table_ptr = config.decode_table.as_ptr() as *const __m128i;
    let tbl0 = _mm_loadu_si128(table_ptr);
    let tbl1 = _mm_loadu_si128(table_ptr.add(1));
    let tbl2 = _mm_loadu_si128(table_ptr.add(2));
    let tbl3 = _mm_loadu_si128(table_ptr.add(3));
    let tbl4 = _mm_loadu_si128(table_ptr.add(4));
    let tbl5 = _mm_loadu_si128(table_ptr.add(5));
    let tbl6 = _mm_loadu_si128(table_ptr.add(6));
    let tbl7 = _mm_loadu_si128(table_ptr.add(7));

    // pshufb shuffle: compacts [o2,o1,o0, 0 | o2,o1,o0, 0 | ...] → [o0,o1,o2, o0,o1,o2, ...]
    // indices: [2,1,0, 6,5,4, 10,9,8, 14,13,12, -1,-1,-1,-1]
    // _mm_set_epi8(e15..e0): e0 is byte 0 of result
    let pack_shuf = _mm_set_epi8(-1, -1, -1, -1, 12, 13, 14, 8, 9, 10, 4, 5, 6, 0, 1, 2);

    let mut src_offset = 0usize;
    let mut dst_offset = 0usize;

    while src_offset + 16 <= src.len() {
        let input = _mm_loadu_si128(src.as_ptr().add(src_offset) as *const __m128i);

        // Any byte >= 128 → invalid ASCII, fall back to scalar for this block.
        if _mm_movemask_epi8(input) != 0 {
            let mut written = 0usize;
            for group_offset in (0..16).step_by(4) {
                written += decode_full_group_into(
                    config,
                    &mut dst[dst_offset + written..],
                    &src[src_offset + group_offset..src_offset + group_offset + 4],
                    src_offset + group_offset,
                )?;
            }
            dst_offset += written;
            src_offset += 16;
            continue;
        }

        // Table lookup via 8 × pshufb.
        //
        // upper = bits 6:4 of each input byte (selects one of 8 decode-table blocks).
        // low   = bits 3:0 of each input byte (index within the 16-byte block).
        // sel_k = 0xFF where input's upper nibble == k, else 0x00.
        // r_k   = pshufb(tbl_k, low) & sel_k  (pshufb reads tbl_k[low[i]])
        // decoded = r_0 | r_1 | ... | r_7
        let upper = _mm_and_si128(input, _mm_set1_epi8(0x70u8 as i8));
        let low = _mm_and_si128(input, _mm_set1_epi8(0x0F));

        let sel0 = _mm_cmpeq_epi8(upper, _mm_setzero_si128());
        let sel1 = _mm_cmpeq_epi8(upper, _mm_set1_epi8(0x10u8 as i8));
        let sel2 = _mm_cmpeq_epi8(upper, _mm_set1_epi8(0x20u8 as i8));
        let sel3 = _mm_cmpeq_epi8(upper, _mm_set1_epi8(0x30u8 as i8));
        let sel4 = _mm_cmpeq_epi8(upper, _mm_set1_epi8(0x40u8 as i8));
        let sel5 = _mm_cmpeq_epi8(upper, _mm_set1_epi8(0x50u8 as i8));
        let sel6 = _mm_cmpeq_epi8(upper, _mm_set1_epi8(0x60u8 as i8));
        let sel7 = _mm_cmpeq_epi8(upper, _mm_set1_epi8(0x70u8 as i8));

        let r0 = _mm_and_si128(_mm_shuffle_epi8(tbl0, low), sel0);
        let r1 = _mm_and_si128(_mm_shuffle_epi8(tbl1, low), sel1);
        let r2 = _mm_and_si128(_mm_shuffle_epi8(tbl2, low), sel2);
        let r3 = _mm_and_si128(_mm_shuffle_epi8(tbl3, low), sel3);
        let r4 = _mm_and_si128(_mm_shuffle_epi8(tbl4, low), sel4);
        let r5 = _mm_and_si128(_mm_shuffle_epi8(tbl5, low), sel5);
        let r6 = _mm_and_si128(_mm_shuffle_epi8(tbl6, low), sel6);
        let r7 = _mm_and_si128(_mm_shuffle_epi8(tbl7, low), sel7);

        let decoded = _mm_or_si128(
            _mm_or_si128(_mm_or_si128(r0, r1), _mm_or_si128(r2, r3)),
            _mm_or_si128(_mm_or_si128(r4, r5), _mm_or_si128(r6, r7)),
        );

        // Any 0xFF in decoded means an invalid character in the table.
        if _mm_movemask_epi8(_mm_cmpeq_epi8(decoded, _mm_set1_epi8(-1))) != 0 {
            let mut written = 0usize;
            for group_offset in (0..16).step_by(4) {
                written += decode_full_group_into(
                    config,
                    &mut dst[dst_offset + written..],
                    &src[src_offset + group_offset..src_offset + group_offset + 4],
                    src_offset + group_offset,
                )?;
            }
            dst_offset += written;
            src_offset += 16;
            continue;
        }

        // Pack 16 × 6-bit decoded values into 12 output bytes (Muła/Lemire algorithm).
        //
        // Step A — maddubs: merge adjacent pairs (a, b) → a*64 + b (12-bit in 16-bit field).
        //   decoded[i] ∈ [0..63], so a*64+b ≤ 4095, no i16 saturation.
        let t0 = _mm_maddubs_epi16(decoded, _mm_set1_epi32(0x01400140u32 as i32));

        // Step B — madd: merge adjacent 12-bit pairs (c, d) → c<<12 | d (24-bit in 32-bit field).
        //   In LE memory the 32-bit word is: [o2, o1, o0, 0x00].
        let t1 = _mm_madd_epi16(t0, _mm_set1_epi32(0x00011000u32 as i32));

        // Step C — pshufb: reorder bytes to produce [o0_a,o1_a,o2_a, o0_b,o1_b,o2_b, ...].
        let packed = _mm_shuffle_epi8(t1, pack_shuf);

        // Write exactly 12 bytes: 8 via storel_epi64, then 4 via scalar.
        let out_ptr = dst.as_mut_ptr().add(dst_offset);
        _mm_storel_epi64(out_ptr as *mut __m128i, packed);
        let shifted = _mm_srli_si128(packed, 8);
        let word = _mm_cvtsi128_si32(shifted);
        core::ptr::write_unaligned(out_ptr.add(8) as *mut i32, word);

        src_offset += 16;
        dst_offset += 12;
    }

    Ok(dst_offset)
}