base64-ng 2.0.1

#![allow(unsafe_code)]

use crate::{Alphabet, DecodeError, scalar};

const SSSE3_DECODE_INPUT_BLOCK: usize = 16;
const SSSE3_DECODE_OUTPUT_BLOCK: usize = 12;
const AVX2_DECODE_INPUT_BLOCK: usize = 32;
const AVX2_DECODE_OUTPUT_BLOCK: usize = 24;
const AVX512_DECODE_INPUT_BLOCK: usize = 64;
const AVX512_DECODE_OUTPUT_BLOCK: usize = 48;

pub(crate) fn decode_slice_ssse3_sse41<A, const PAD: bool>(
    input: &[u8],
    output: &mut [u8],
) -> Result<usize, DecodeError>
where
    A: Alphabet,
{
    if input.len() < SSSE3_DECODE_INPUT_BLOCK || !super::ssse3_sse41_supports_decode_alphabet::<A>()
    {
        return scalar::decode_slice::<A, PAD>(input, output);
    }

    let required = scalar::validate_decode::<A, PAD>(input)?;
    if output.len() < required {
        return Err(DecodeError::OutputTooSmall {
            required,
            available: output.len(),
        });
    }

    let simd_input_len = unpadded_simd_prefix_len(input);
    // SAFETY: Health-gated dispatch or the static token contract proves the
    // target features. Validation and the output preflight above prove every
    // full block is canonical and every direct store remains in bounds.
    let (read, write, classified) =
        unsafe { decode_full_blocks_ssse3_sse41::<A>(input, output, simd_input_len) };
    if !classified {
        return scalar::decode_slice::<A, PAD>(input, output);
    }

    let tail_written = scalar::decode_slice::<A, PAD>(&input[read..], &mut output[write..])
        .map_err(|error| error.with_index_offset(read))?;
    Ok(write + tail_written)
}

pub(crate) fn decode_slice_avx2<A, const PAD: bool>(
    input: &[u8],
    output: &mut [u8],
) -> Result<usize, DecodeError>
where
    A: Alphabet,
{
    if input.len() < AVX2_DECODE_INPUT_BLOCK || !super::avx2_supports_decode_alphabet::<A>() {
        return decode_slice_ssse3_sse41::<A, PAD>(input, output);
    }

    let required = scalar::validate_decode::<A, PAD>(input)?;
    if output.len() < required {
        return Err(DecodeError::OutputTooSmall {
            required,
            available: output.len(),
        });
    }

    let simd_input_len = unpadded_simd_prefix_len(input);
    // SAFETY: Health-gated dispatch or the static token contract proves AVX2.
    // Validation and the output preflight above prove exact block bounds.
    let (read, write, classified) =
        unsafe { decode_full_blocks_avx2::<A>(input, output, simd_input_len) };
    if !classified {
        return scalar::decode_slice::<A, PAD>(input, output);
    }

    let tail_written = decode_slice_ssse3_sse41::<A, PAD>(&input[read..], &mut output[write..])
        .map_err(|error| error.with_index_offset(read))?;
    Ok(write + tail_written)
}

fn unpadded_simd_prefix_len(input: &[u8]) -> usize {
    if input.last() == Some(&b'=') {
        input.len().saturating_sub(4)
    } else {
        input.len()
    }
}

pub(crate) fn decode_slice_avx512<A, const PAD: bool>(
    input: &[u8],
    output: &mut [u8],
) -> Result<usize, DecodeError>
where
    A: Alphabet,
{
    if input.len() < AVX512_DECODE_INPUT_BLOCK || !super::avx512_supports_decode_alphabet::<A>() {
        return decode_slice_avx2::<A, PAD>(input, output);
    }

    let required = scalar::validate_decode::<A, PAD>(input)?;
    if output.len() < required {
        return Err(DecodeError::OutputTooSmall {
            required,
            available: output.len(),
        });
    }

    let simd_input_len = unpadded_simd_prefix_len(input);
    // SAFETY: Health-gated dispatch or the static token contract proves the
    // complete AVX-512 VBMI feature bundle. Validation and output preflight
    // prove every exact block access remains in bounds.
    let (read, write, classified) =
        unsafe { decode_full_blocks_avx512::<A>(input, output, simd_input_len) };
    if !classified {
        return scalar::decode_slice::<A, PAD>(input, output);
    }

    let tail_written = decode_slice_avx2::<A, PAD>(&input[read..], &mut output[write..])
        .map_err(|error| error.with_index_offset(read))?;
    Ok(write + tail_written)
}

#[cfg(all(feature = "std", test))]
#[target_feature(enable = "ssse3,sse4.1")]
pub(crate) unsafe fn decode_16_bytes_ssse3_sse41<A, const PAD: bool>(
    input: &[u8; 16],
    output: &mut [u8; 12],
) -> Result<usize, DecodeError>
where
    A: Alphabet,
{
    let written = scalar::validate_decode::<A, PAD>(input)?;
    if written != SSSE3_DECODE_OUTPUT_BLOCK
        || !super::is_standard_or_url_safe_family::<A>()
        || input.contains(&b'=')
    {
        return scalar::decode_slice::<A, PAD>(input, output);
    }

    // SAFETY: This function carries the target-feature contract. Scalar
    // validation proves the unpadded block is canonical, and fixed arrays
    // prove the exact 16-byte load and 12-byte store bounds.
    let classified =
        unsafe { super::decode_direct::decode_16_bytes_ssse3_sse41::<A>(input, output) };
    // SAFETY: The direct block no longer needs SIMD register contents.
    unsafe { super::cleanup::clear_xmm_registers_after_encode_block() };
    if !classified {
        return scalar::decode_slice::<A, PAD>(input, output);
    }
    Ok(SSSE3_DECODE_OUTPUT_BLOCK)
}

#[cfg(all(feature = "std", test))]
#[target_feature(enable = "avx2")]
pub(crate) unsafe fn decode_32_bytes_avx2<A, const PAD: bool>(
    input: &[u8; 32],
    output: &mut [u8; 24],
) -> Result<usize, DecodeError>
where
    A: Alphabet,
{
    let written = scalar::validate_decode::<A, PAD>(input)?;
    if written != AVX2_DECODE_OUTPUT_BLOCK
        || !super::is_standard_or_url_safe_family::<A>()
        || input.contains(&b'=')
    {
        return scalar::decode_slice::<A, PAD>(input, output);
    }

    // SAFETY: This function carries the target-feature contract. Scalar
    // validation proves the unpadded block is canonical, and fixed arrays
    // prove the exact 32-byte load and two 12-byte stores remain in bounds.
    let classified = unsafe { super::decode_direct::decode_32_bytes_avx2::<A>(input, output) };
    // SAFETY: The direct block no longer needs SIMD register contents.
    unsafe { super::cleanup::clear_ymm_registers_after_encode_block() };
    if !classified {
        return scalar::decode_slice::<A, PAD>(input, output);
    }
    Ok(AVX2_DECODE_OUTPUT_BLOCK)
}

#[target_feature(enable = "ssse3,sse4.1")]
unsafe fn decode_full_blocks_ssse3_sse41<A>(
    input: &[u8],
    output: &mut [u8],
    simd_input_len: usize,
) -> (usize, usize, bool)
where
    A: Alphabet,
{
    let mut read = 0;
    let mut write = 0;
    let mut classified = true;
    while read + SSSE3_DECODE_INPUT_BLOCK <= simd_input_len {
        // SAFETY: The loop guards prove both exact fixed blocks are within the
        // preflighted slices. This function carries the ISA contract.
        let block_classified = unsafe {
            let block = &*(input
                .as_ptr()
                .add(read)
                .cast::<[u8; SSSE3_DECODE_INPUT_BLOCK]>());
            let decoded = &mut *(output
                .as_mut_ptr()
                .add(write)
                .cast::<[u8; SSSE3_DECODE_OUTPUT_BLOCK]>());
            super::decode_direct::decode_16_bytes_ssse3_sse41::<A>(block, decoded)
        };
        if !block_classified {
            classified = false;
            break;
        }
        read += SSSE3_DECODE_INPUT_BLOCK;
        write += SSSE3_DECODE_OUTPUT_BLOCK;
    }
    if read != 0 {
        // SAFETY: All vector results have already been stored.
        unsafe { super::cleanup::clear_xmm_registers_after_encode_block() };
    }
    (read, write, classified)
}

#[target_feature(enable = "avx2")]
unsafe fn decode_full_blocks_avx2<A>(
    input: &[u8],
    output: &mut [u8],
    simd_input_len: usize,
) -> (usize, usize, bool)
where
    A: Alphabet,
{
    let mut read = 0;
    let mut write = 0;
    let mut classified = true;
    while read + AVX2_DECODE_INPUT_BLOCK <= simd_input_len {
        // SAFETY: The loop guards prove both exact fixed blocks are within the
        // preflighted slices. This function carries the ISA contract.
        let block_classified = unsafe {
            let block = &*(input
                .as_ptr()
                .add(read)
                .cast::<[u8; AVX2_DECODE_INPUT_BLOCK]>());
            let decoded = &mut *(output
                .as_mut_ptr()
                .add(write)
                .cast::<[u8; AVX2_DECODE_OUTPUT_BLOCK]>());
            super::decode_direct::decode_32_bytes_avx2::<A>(block, decoded)
        };
        if !block_classified {
            classified = false;
            break;
        }
        read += AVX2_DECODE_INPUT_BLOCK;
        write += AVX2_DECODE_OUTPUT_BLOCK;
    }
    if read != 0 {
        // SAFETY: All vector results have already been stored.
        unsafe { super::cleanup::clear_ymm_registers_after_encode_block() };
    }
    (read, write, classified)
}

#[cfg(all(feature = "std", test))]
#[target_feature(enable = "avx512f,avx512bw,avx512vl,avx512vbmi")]
pub(crate) unsafe fn decode_64_bytes_avx512<A, const PAD: bool>(
    input: &[u8; 64],
    output: &mut [u8; 48],
) -> Result<usize, DecodeError>
where
    A: Alphabet,
{
    let written = scalar::validate_decode::<A, PAD>(input)?;
    if written != AVX512_DECODE_OUTPUT_BLOCK
        || !super::is_standard_or_url_safe_family::<A>()
        || input.contains(&b'=')
    {
        return scalar::decode_slice::<A, PAD>(input, output);
    }

    // SAFETY: This function carries the complete target-feature contract.
    // Scalar validation proves a canonical unpadded block, and fixed arrays
    // prove the exact 64-byte load and 48-byte masked store bounds.
    let classified = unsafe { super::decode_direct::decode_64_bytes_avx512::<A>(input, output) };
    // SAFETY: The direct block no longer needs SIMD register contents.
    unsafe { super::cleanup::clear_zmm_registers_after_encode_block() };
    if !classified {
        return scalar::decode_slice::<A, PAD>(input, output);
    }
    Ok(AVX512_DECODE_OUTPUT_BLOCK)
}

#[target_feature(enable = "avx512f,avx512bw,avx512vl,avx512vbmi")]
unsafe fn decode_full_blocks_avx512<A>(
    input: &[u8],
    output: &mut [u8],
    simd_input_len: usize,
) -> (usize, usize, bool)
where
    A: Alphabet,
{
    let mut read = 0;
    let mut write = 0;
    let mut classified = true;
    while read + AVX512_DECODE_INPUT_BLOCK <= simd_input_len {
        // SAFETY: The loop guards prove both exact fixed blocks are within the
        // preflighted slices. This function carries the complete ISA contract.
        let block_classified = unsafe {
            let block = &*(input
                .as_ptr()
                .add(read)
                .cast::<[u8; AVX512_DECODE_INPUT_BLOCK]>());
            let decoded = &mut *(output
                .as_mut_ptr()
                .add(write)
                .cast::<[u8; AVX512_DECODE_OUTPUT_BLOCK]>());
            super::decode_direct::decode_64_bytes_avx512::<A>(block, decoded)
        };
        if !block_classified {
            classified = false;
            break;
        }
        read += AVX512_DECODE_INPUT_BLOCK;
        write += AVX512_DECODE_OUTPUT_BLOCK;
    }
    if read != 0 {
        // SAFETY: All vector results have already been stored.
        unsafe { super::cleanup::clear_zmm_registers_after_encode_block() };
    }
    (read, write, classified)
}