#![allow(unsafe_code)]
use crate::alphabet::Symbol;
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
use crate::{
engine::{
general_purpose::{
decode::decode_helper, encode_helper, GeneralPurpose, GeneralPurposeConfig,
GeneralPurposeEstimate,
},
DecodeMetadata, Engine,
},
DecodeSliceError,
};
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
trait SimdAlphabet {
const ENCODE_LUT: [i8; 16];
const DECODE_SHIFT_LUT: [i8; 16];
const DECODE_MASK_LUT: [u8; 16];
const DECODE_FIXUP_CHAR: i8;
const DECODE_FIXUP_SHIFT: i8;
}
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
enum Standard {}
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
enum UrlSafe {}
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
impl SimdAlphabet for Standard {
const ENCODE_LUT: [i8; 16] = [
65, 71, -4, -4, -4, -4, -4, -4, -4, -4, -4, -4, -19, -16, 0, 0,
];
const DECODE_SHIFT_LUT: [i8; 16] = [0, 0, 19, 4, -65, -65, -71, -71, 0, 0, 0, 0, 0, 0, 0, 0];
const DECODE_MASK_LUT: [u8; 16] = [
0xA8, 0xF8, 0xF8, 0xF8, 0xF8, 0xF8, 0xF8, 0xF8, 0xF8, 0xF8, 0xF0, 0x54, 0x50, 0x50, 0x50,
0x54,
];
const DECODE_FIXUP_CHAR: i8 = 0x2F;
const DECODE_FIXUP_SHIFT: i8 = 16;
}
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
impl SimdAlphabet for UrlSafe {
const ENCODE_LUT: [i8; 16] = [
65, 71, -4, -4, -4, -4, -4, -4, -4, -4, -4, -4, -17, 32, 0, 0,
];
const DECODE_SHIFT_LUT: [i8; 16] = [0, 0, 17, 4, -65, -65, -71, -71, 0, 0, 0, 0, 0, 0, 0, 0];
const DECODE_MASK_LUT: [u8; 16] = [
0xA8, 0xF8, 0xF8, 0xF8, 0xF8, 0xF8, 0xF8, 0xF8, 0xF8, 0xF8, 0xF0, 0x50, 0x50, 0x54, 0x50,
0x70,
];
const DECODE_FIXUP_CHAR: i8 = 0x5F;
const DECODE_FIXUP_SHIFT: i8 = -32;
}
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum SimdKind {
Standard,
UrlSafe,
}
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
const SIMD_MIN_INPUT_ENCODE: usize = 128;
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
const SIMD_MIN_INPUT_DECODE: usize = 64;
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
const BITPOS_LUT: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 0, 0, 0, 0, 0, 0, 0, 0];
#[cfg(target_arch = "x86_64")]
mod avx2 {
use super::{SimdAlphabet, BITPOS_LUT, SIMD_MIN_INPUT_DECODE, SIMD_MIN_INPUT_ENCODE};
use core::arch::x86_64::*;
#[target_feature(enable = "avx2")]
pub(super) unsafe fn encode_bulk<A: SimdAlphabet>(
input: &[u8],
output: &mut [u8],
) -> (usize, usize) {
if input.len() < SIMD_MIN_INPUT_ENCODE {
return (0, 0);
}
let lut_data = A::ENCODE_LUT;
let lut = _mm256_broadcastsi128_si256(_mm_loadu_si128(lut_data.as_ptr().cast()));
#[rustfmt::skip]
let shuf = _mm256_setr_epi8(
1, 0, 2, 1, 4, 3, 5, 4, 7, 6, 8, 7, 10, 9, 11, 10,
1, 0, 2, 1, 4, 3, 5, 4, 7, 6, 8, 7, 10, 9, 11, 10,
);
let mask_hi = _mm256_set1_epi32(0x0fc0_fc00_u32 as i32);
let mul_hi = _mm256_set1_epi32(0x0400_0040);
let mask_lo = _mm256_set1_epi32(0x003f_03f0);
let mul_lo = _mm256_set1_epi32(0x0100_0010);
let const51 = _mm256_set1_epi8(51);
let const25 = _mm256_set1_epi8(25);
let mut i = 0usize;
let mut o = 0usize;
while i + 32 <= input.len() && o + 32 <= output.len() {
let data = _mm256_loadu_si256(input.as_ptr().add(i).cast());
let perm = _mm256_permutevar8x32_epi32(data, _mm256_setr_epi32(0, 1, 2, 3, 3, 4, 5, 6));
let inb = _mm256_shuffle_epi8(perm, shuf);
let t0 = _mm256_and_si256(inb, mask_hi);
let t1 = _mm256_mulhi_epu16(t0, mul_hi);
let t2 = _mm256_and_si256(inb, mask_lo);
let t3 = _mm256_mullo_epi16(t2, mul_lo);
let indices = _mm256_or_si256(t1, t3);
let reduced = _mm256_subs_epu8(indices, const51);
let gt25 = _mm256_cmpgt_epi8(indices, const25);
let reduced = _mm256_sub_epi8(reduced, gt25);
let ascii = _mm256_add_epi8(indices, _mm256_shuffle_epi8(lut, reduced));
_mm256_storeu_si256(output.as_mut_ptr().add(o).cast(), ascii);
i += 24;
o += 32;
}
(i, o)
}
#[target_feature(enable = "avx2")]
pub(super) unsafe fn decode_bulk<A: SimdAlphabet>(
input: &[u8],
quads_end: usize,
output: &mut [u8],
) -> (usize, usize) {
if quads_end < SIMD_MIN_INPUT_DECODE {
return (0, 0);
}
let shift_data = A::DECODE_SHIFT_LUT;
let mask_data = A::DECODE_MASK_LUT;
let shift_lut = _mm256_broadcastsi128_si256(_mm_loadu_si128(shift_data.as_ptr().cast()));
let mask_lut = _mm256_broadcastsi128_si256(_mm_loadu_si128(mask_data.as_ptr().cast()));
let bitpos_lut = _mm256_broadcastsi128_si256(_mm_loadu_si128(BITPOS_LUT.as_ptr().cast()));
let low_nibble_mask = _mm256_set1_epi8(0x0f);
let fixup_char = _mm256_set1_epi8(A::DECODE_FIXUP_CHAR);
let fixup_shift = _mm256_set1_epi8(A::DECODE_FIXUP_SHIFT);
let zero = _mm256_setzero_si256();
let merge_mul1 = _mm256_set1_epi32(0x0140_0140);
let merge_mul2 = _mm256_set1_epi32(0x0001_1000);
#[rustfmt::skip]
let pack_shuf = _mm256_setr_epi8(
2, 1, 0, 6, 5, 4, 10, 9, 8, 14, 13, 12, -1, -1, -1, -1,
2, 1, 0, 6, 5, 4, 10, 9, 8, 14, 13, 12, -1, -1, -1, -1,
);
let lane_compact = _mm256_setr_epi32(0, 1, 2, 4, 5, 6, 6, 6);
let mut i = 0usize;
let mut o = 0usize;
while i + 32 <= quads_end && o + 24 <= output.len() {
let data = _mm256_loadu_si256(input.as_ptr().add(i).cast());
let hi_nibbles = _mm256_and_si256(_mm256_srli_epi32(data, 4), low_nibble_mask);
let lo_nibbles = _mm256_and_si256(data, low_nibble_mask);
let m = _mm256_shuffle_epi8(mask_lut, lo_nibbles);
let bit = _mm256_shuffle_epi8(bitpos_lut, hi_nibbles);
let non_match = _mm256_cmpeq_epi8(_mm256_and_si256(m, bit), zero);
if _mm256_movemask_epi8(non_match) != 0 {
break;
}
let sh = _mm256_shuffle_epi8(shift_lut, hi_nibbles);
let eq_fixup = _mm256_cmpeq_epi8(data, fixup_char);
let shift = _mm256_blendv_epi8(sh, fixup_shift, eq_fixup);
let values = _mm256_add_epi8(data, shift);
let merged = _mm256_maddubs_epi16(values, merge_mul1);
let packed = _mm256_madd_epi16(merged, merge_mul2);
let shuffled = _mm256_shuffle_epi8(packed, pack_shuf);
let compact = _mm256_permutevar8x32_epi32(shuffled, lane_compact);
let lo = _mm256_castsi256_si128(compact);
let hi = _mm256_extracti128_si256(compact, 1);
_mm_storeu_si128(output.as_mut_ptr().add(o).cast(), lo);
_mm_storel_epi64(output.as_mut_ptr().add(o + 16).cast(), hi);
i += 32;
o += 24;
}
(i, o)
}
}
#[cfg(target_arch = "aarch64")]
mod neon {
use super::{SimdAlphabet, BITPOS_LUT, SIMD_MIN_INPUT_DECODE, SIMD_MIN_INPUT_ENCODE};
use core::arch::aarch64::*;
#[target_feature(enable = "neon")]
pub(super) unsafe fn encode_bulk<A: SimdAlphabet>(
input: &[u8],
output: &mut [u8],
) -> (usize, usize) {
if input.len() < SIMD_MIN_INPUT_ENCODE {
return (0, 0);
}
let lut_data = A::ENCODE_LUT;
let split_bytes: [u8; 16] = [1, 0, 2, 1, 4, 3, 5, 4, 7, 6, 8, 7, 10, 9, 11, 10];
let split_shuf = vld1q_u8(split_bytes.as_ptr());
let translate = vld1q_u8(lut_data.as_ptr().cast());
let m1 = vdupq_n_u32(0x0000_fc00);
let m2 = vdupq_n_u32(0x0000_03f0);
let m3 = vdupq_n_u32(0x0fc0_0000);
let m4 = vdupq_n_u32(0x003f_0000);
let c51 = vdupq_n_u8(51);
let c25 = vdupq_n_s8(25);
let mut i = 0usize;
let mut o = 0usize;
while i + 16 <= input.len() && o + 16 <= output.len() {
let data = vld1q_u8(input.as_ptr().add(i));
let x0 = vreinterpretq_u32_u8(vqtbl1q_u8(data, split_shuf));
let x1 = vshrq_n_u16::<10>(vreinterpretq_u16_u32(vandq_u32(x0, m1)));
let x2 = vshlq_n_u16::<4>(vreinterpretq_u16_u32(vandq_u32(x0, m2)));
let x3 = vshrq_n_u16::<6>(vreinterpretq_u16_u32(vandq_u32(x0, m3)));
let x4 = vshlq_n_u16::<8>(vreinterpretq_u16_u32(vandq_u32(x0, m4)));
let indices = vreinterpretq_u8_u16(vorrq_u16(vorrq_u16(x1, x2), vorrq_u16(x3, x4)));
let reduced = vqsubq_u8(indices, c51);
let gt25 = vcgtq_s8(vreinterpretq_s8_u8(indices), c25);
let reduced = vsubq_u8(reduced, gt25); let ascii = vaddq_u8(indices, vqtbl1q_u8(translate, reduced));
vst1q_u8(output.as_mut_ptr().add(o), ascii);
i += 12;
o += 16;
}
(i, o)
}
#[target_feature(enable = "neon")]
pub(super) unsafe fn decode_bulk<A: SimdAlphabet>(
input: &[u8],
quads_end: usize,
output: &mut [u8],
) -> (usize, usize) {
if quads_end < SIMD_MIN_INPUT_DECODE {
return (0, 0);
}
let shift_lut = vld1q_u8(A::DECODE_SHIFT_LUT.as_ptr().cast());
let mask_lut = vld1q_u8(A::DECODE_MASK_LUT.as_ptr());
let bitpos_lut = vld1q_u8(BITPOS_LUT.as_ptr());
let low_nibble_mask = vdupq_n_u8(0x0f);
let fixup_char_v = vdupq_n_u8(A::DECODE_FIXUP_CHAR as u8);
let fixup_shift_v = vdupq_n_u8(A::DECODE_FIXUP_SHIFT as u8);
let zero = vdupq_n_u8(0);
let mm1 = vdupq_n_u32(0x003f_003f);
let mm2 = vdupq_n_u32(0x3f00_3f00);
let out_mask = vdupq_n_u32(0x00ff_ffff);
let pack_bytes: [u8; 16] = [
2, 1, 0, 6, 5, 4, 10, 9, 8, 14, 13, 12, 0x80, 0x80, 0x80, 0x80,
];
let pack_shuf = vld1q_u8(pack_bytes.as_ptr());
let mut i = 0usize;
let mut o = 0usize;
while i + 16 <= quads_end && o + 12 <= output.len() {
let data = vld1q_u8(input.as_ptr().add(i));
let hi = vshrq_n_u8::<4>(data);
let lo = vandq_u8(data, low_nibble_mask);
let m = vqtbl1q_u8(mask_lut, lo);
let bit = vqtbl1q_u8(bitpos_lut, hi);
let non_match = vceqq_u8(vandq_u8(m, bit), zero);
if vmaxvq_u8(non_match) != 0 {
break;
}
let sh = vqtbl1q_u8(shift_lut, hi);
let eq_fixup = vceqq_u8(data, fixup_char_v);
let shift = vbslq_u8(eq_fixup, fixup_shift_v, sh);
let values = vaddq_u8(data, shift);
let v = vreinterpretq_u32_u8(values);
let x1 = vandq_u32(v, mm1); let x2 = vandq_u32(v, mm2); let x3 = vorrq_u32(vshlq_n_u32::<18>(x1), vshrq_n_u32::<10>(x1));
let x4 = vorrq_u32(vshlq_n_u32::<4>(x2), vshrq_n_u32::<24>(x2));
let merged = vandq_u32(vorrq_u32(x3, x4), out_mask);
let packed = vqtbl1q_u8(vreinterpretq_u8_u32(merged), pack_shuf);
vst1_u8(output.as_mut_ptr().add(o), vget_low_u8(packed));
let tail = vgetq_lane_u32::<2>(vreinterpretq_u32_u8(packed));
core::ptr::write_unaligned(output.as_mut_ptr().add(o + 8).cast::<u32>(), tail);
i += 16;
o += 12;
}
(i, o)
}
}
#[cfg(target_arch = "x86_64")]
#[inline]
unsafe fn avx2_encode(kind: SimdKind, input: &[u8], output: &mut [u8]) -> (usize, usize) {
match kind {
SimdKind::Standard => avx2::encode_bulk::<Standard>(input, output),
SimdKind::UrlSafe => avx2::encode_bulk::<UrlSafe>(input, output),
}
}
#[cfg(target_arch = "x86_64")]
#[inline]
unsafe fn avx2_decode(
kind: SimdKind,
input: &[u8],
quads_end: usize,
output: &mut [u8],
) -> (usize, usize) {
match kind {
SimdKind::Standard => avx2::decode_bulk::<Standard>(input, quads_end, output),
SimdKind::UrlSafe => avx2::decode_bulk::<UrlSafe>(input, quads_end, output),
}
}
#[cfg(target_arch = "aarch64")]
#[inline]
unsafe fn neon_encode(kind: SimdKind, input: &[u8], output: &mut [u8]) -> (usize, usize) {
match kind {
SimdKind::Standard => neon::encode_bulk::<Standard>(input, output),
SimdKind::UrlSafe => neon::encode_bulk::<UrlSafe>(input, output),
}
}
#[cfg(target_arch = "aarch64")]
#[inline]
unsafe fn neon_decode(
kind: SimdKind,
input: &[u8],
quads_end: usize,
output: &mut [u8],
) -> (usize, usize) {
match kind {
SimdKind::Standard => neon::decode_bulk::<Standard>(input, quads_end, output),
SimdKind::UrlSafe => neon::decode_bulk::<UrlSafe>(input, quads_end, output),
}
}
#[cfg(all(feature = "std", any(target_arch = "x86_64", target_arch = "aarch64")))]
#[derive(Clone, Copy, Debug)]
enum Backend {
Scalar,
#[cfg(target_arch = "x86_64")]
Avx2,
#[cfg(target_arch = "aarch64")]
Neon,
}
#[cfg(all(feature = "std", any(target_arch = "x86_64", target_arch = "aarch64")))]
#[derive(Debug, Clone)]
pub struct Simd {
inner: GeneralPurpose,
kind: SimdKind,
backend: Backend,
}
#[cfg(all(feature = "std", any(target_arch = "x86_64", target_arch = "aarch64")))]
impl Simd {
#[must_use]
pub fn standard(config: GeneralPurposeConfig) -> Self {
Self::new(SimdKind::Standard, &crate::alphabet::STANDARD, config)
}
#[must_use]
pub fn url_safe(config: GeneralPurposeConfig) -> Self {
Self::new(SimdKind::UrlSafe, &crate::alphabet::URL_SAFE, config)
}
fn new(
kind: SimdKind,
alphabet: &crate::alphabet::Alphabet,
config: GeneralPurposeConfig,
) -> Self {
#[cfg(target_arch = "x86_64")]
let backend = if std::is_x86_feature_detected!("avx2") {
Backend::Avx2
} else {
Backend::Scalar
};
#[cfg(target_arch = "aarch64")]
let backend = if std::arch::is_aarch64_feature_detected!("neon") {
Backend::Neon
} else {
Backend::Scalar
};
Self {
inner: GeneralPurpose::new(alphabet, config),
kind,
backend,
}
}
}
#[cfg(all(feature = "std", any(target_arch = "x86_64", target_arch = "aarch64")))]
impl Engine for Simd {
type Config = GeneralPurposeConfig;
type DecodeEstimate = GeneralPurposeEstimate;
fn internal_encode(&self, input: &[u8], output: &mut [u8]) -> usize {
let kind = self.kind;
match self.backend {
Backend::Scalar => self.inner.internal_encode(input, output),
#[cfg(target_arch = "x86_64")]
Backend::Avx2 => encode_helper(self.inner.encode_table(), input, output, |i, o| {
unsafe { avx2_encode(kind, i, o) }
}),
#[cfg(target_arch = "aarch64")]
Backend::Neon => encode_helper(self.inner.encode_table(), input, output, |i, o| {
unsafe { neon_encode(kind, i, o) }
}),
}
}
fn internal_decoded_len_estimate(&self, input_len: usize) -> Self::DecodeEstimate {
self.inner.internal_decoded_len_estimate(input_len)
}
fn internal_decode(
&self,
input: &[u8],
output: &mut [u8],
estimate: Self::DecodeEstimate,
) -> Result<DecodeMetadata, DecodeSliceError> {
let kind = self.kind;
match self.backend {
Backend::Scalar => self.inner.internal_decode(input, output, estimate),
#[cfg(target_arch = "x86_64")]
Backend::Avx2 => decode_helper(
input,
&estimate,
output,
self.inner.decode_table(),
self.inner.config().decode_allow_trailing_bits(),
self.inner.padding(),
self.inner.config().decode_padding_mode(),
|i, end, o| unsafe { avx2_decode(kind, i, end, o) },
),
#[cfg(target_arch = "aarch64")]
Backend::Neon => decode_helper(
input,
&estimate,
output,
self.inner.decode_table(),
self.inner.config().decode_allow_trailing_bits(),
self.inner.padding(),
self.inner.config().decode_padding_mode(),
|i, end, o| unsafe { neon_decode(kind, i, end, o) },
),
}
}
fn config(&self) -> &Self::Config {
self.inner.config()
}
fn padding(&self) -> Symbol {
self.inner.padding()
}
}
#[cfg(target_arch = "x86_64")]
#[derive(Debug, Clone)]
pub struct Avx2 {
inner: GeneralPurpose,
kind: SimdKind,
}
#[cfg(target_arch = "x86_64")]
impl Avx2 {
#[cfg(feature = "std")]
#[must_use]
pub fn standard(config: GeneralPurposeConfig) -> Option<Self> {
if std::is_x86_feature_detected!("avx2") {
Some(unsafe { Self::standard_unchecked(config) })
} else {
None
}
}
#[cfg(feature = "std")]
#[must_use]
pub fn url_safe(config: GeneralPurposeConfig) -> Option<Self> {
if std::is_x86_feature_detected!("avx2") {
Some(unsafe { Self::url_safe_unchecked(config) })
} else {
None
}
}
#[must_use]
pub const unsafe fn standard_unchecked(config: GeneralPurposeConfig) -> Self {
Self {
inner: GeneralPurpose::new(&crate::alphabet::STANDARD, config),
kind: SimdKind::Standard,
}
}
#[must_use]
pub const unsafe fn url_safe_unchecked(config: GeneralPurposeConfig) -> Self {
Self {
inner: GeneralPurpose::new(&crate::alphabet::URL_SAFE, config),
kind: SimdKind::UrlSafe,
}
}
}
#[cfg(target_arch = "x86_64")]
impl Engine for Avx2 {
type Config = GeneralPurposeConfig;
type DecodeEstimate = GeneralPurposeEstimate;
fn internal_encode(&self, input: &[u8], output: &mut [u8]) -> usize {
let kind = self.kind;
encode_helper(self.inner.encode_table(), input, output, |i, o| {
unsafe { avx2_encode(kind, i, o) }
})
}
fn internal_decoded_len_estimate(&self, input_len: usize) -> Self::DecodeEstimate {
self.inner.internal_decoded_len_estimate(input_len)
}
fn internal_decode(
&self,
input: &[u8],
output: &mut [u8],
estimate: Self::DecodeEstimate,
) -> Result<DecodeMetadata, DecodeSliceError> {
let kind = self.kind;
decode_helper(
input,
&estimate,
output,
self.inner.decode_table(),
self.inner.config().decode_allow_trailing_bits(),
self.inner.padding(),
self.inner.config().decode_padding_mode(),
|i, end, o| unsafe { avx2_decode(kind, i, end, o) },
)
}
fn config(&self) -> &Self::Config {
self.inner.config()
}
fn padding(&self) -> Symbol {
self.inner.padding()
}
}
#[cfg(target_arch = "aarch64")]
#[derive(Debug, Clone)]
pub struct Neon {
inner: GeneralPurpose,
kind: SimdKind,
}
#[cfg(target_arch = "aarch64")]
impl Neon {
#[must_use]
pub const fn standard(config: GeneralPurposeConfig) -> Self {
Self {
inner: GeneralPurpose::new(&crate::alphabet::STANDARD, config),
kind: SimdKind::Standard,
}
}
#[must_use]
pub const fn url_safe(config: GeneralPurposeConfig) -> Self {
Self {
inner: GeneralPurpose::new(&crate::alphabet::URL_SAFE, config),
kind: SimdKind::UrlSafe,
}
}
}
#[cfg(target_arch = "aarch64")]
impl Engine for Neon {
type Config = GeneralPurposeConfig;
type DecodeEstimate = GeneralPurposeEstimate;
fn internal_encode(&self, input: &[u8], output: &mut [u8]) -> usize {
let kind = self.kind;
encode_helper(self.inner.encode_table(), input, output, |i, o| {
unsafe { neon_encode(kind, i, o) }
})
}
fn internal_decoded_len_estimate(&self, input_len: usize) -> Self::DecodeEstimate {
self.inner.internal_decoded_len_estimate(input_len)
}
fn internal_decode(
&self,
input: &[u8],
output: &mut [u8],
estimate: Self::DecodeEstimate,
) -> Result<DecodeMetadata, DecodeSliceError> {
let kind = self.kind;
decode_helper(
input,
&estimate,
output,
self.inner.decode_table(),
self.inner.config().decode_allow_trailing_bits(),
self.inner.padding(),
self.inner.config().decode_padding_mode(),
|i, end, o| unsafe { neon_decode(kind, i, end, o) },
)
}
fn config(&self) -> &Self::Config {
self.inner.config()
}
fn padding(&self) -> Symbol {
self.inner.padding()
}
}