fast-hex/src/enc.rs

//! SIMD-accelerated hex encoding.

use std::mem::MaybeUninit;

use crate::{simd, util};

use util::array_op;

const HEX_CHARS_LOWER: [u8; 16] = array_op!(map[16, ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'a', 'b', 'c', 'd', 'e', 'f']] |_, c| c as u8);
const HEX_CHARS_UPPER: [u8; 16] =
    array_op!(map[16, HEX_CHARS_LOWER] |_, c| (c as char).to_ascii_uppercase() as u8);

macro_rules! select {
    ($cond:ident ? $true:ident : $false:ident) => {
        if $cond {
            $true
        } else {
            $false
        }
    };
    (($cond:expr) ? ($true:expr) : ($false:expr)) => {
        if $cond {
            $true
        } else {
            $false
        }
    };
}

macro_rules! const_impl {
    ($UPPER:ident, $src:ident, $dst:ident) => {{
        let mut i = 0;
        let ub = $dst.len();
        while i < ub {
            let b = $src[i >> 1];
            $dst[i] = MaybeUninit::new(select!($UPPER ? HEX_CHARS_UPPER : HEX_CHARS_LOWER)[(b >> 4) as usize]);
            $dst[i + 1] = MaybeUninit::new(select!($UPPER ? HEX_CHARS_UPPER : HEX_CHARS_LOWER)[(b & 0x0f) as usize]);
            i += 2;
        }
    }};
}

macro_rules! common_impl {
    ($UPPER:ident, $src:ident, $dst:ident) => {
        const_impl!($UPPER, $src, $dst)
    };
}

pub trait Encode {
    /// Encodes the sized input on the stack.
    fn enc_sized<const N: usize>(src: &[u8; N]) -> [u8; N * 2]
    where
        [u8; N * 2]:;

    /// Encodes the sized input on the heap.
    fn enc_sized_heap<const N: usize>(src: &[u8; N]) -> Box<[u8; N * 2]>
    where
        [u8; N * 2]:;

    /// Encodes the unsized input on the heap.
    fn enc_slice(src: &[u8]) -> Box<[u8]>;
}

pub struct Encoder<const UPPER: bool = false>;

impl<const UPPER: bool> Encode for Encoder<UPPER> {
    #[inline]
    fn enc_sized<const N: usize>(src: &[u8; N]) -> [u8; N * 2]
    where
        [u8; N * 2]:,
    {
        let mut buf = MaybeUninit::uninit_array();
        common_impl!(UPPER, src, buf);
        unsafe { MaybeUninit::array_assume_init(buf) }
    }

    #[inline]
    fn enc_sized_heap<const N: usize>(src: &[u8; N]) -> Box<[u8; N * 2]>
    where
        [u8; N * 2]:,
    {
        let mut buf: Box<[MaybeUninit<u8>; N * 2]> = unsafe { Box::new_uninit().assume_init() };
        common_impl!(UPPER, src, buf);
        unsafe { Box::from_raw(Box::into_raw(buf).cast()) }
    }

    #[inline]
    fn enc_slice(src: &[u8]) -> Box<[u8]> {
        let mut buf = Box::new_uninit_slice(src.len() * 2);
        common_impl!(UPPER, src, buf);
        unsafe { Box::<[_]>::assume_init(buf) }
    }
}

impl<const UPPER: bool> Encoder<UPPER> {
    #[inline]
    pub fn enc_const<const N: usize>(src: &[u8; N]) -> [u8; N * 2]
    where
        [u8; N * 2]:,
    {
        let mut buf = MaybeUninit::uninit_array();
        const_impl!(UPPER, src, buf);
        unsafe { MaybeUninit::array_assume_init(buf) }
    }
}

#[cfg(test)]
mod test {
    use super::*;

    use crate::test::*;

    macro_rules! for_each_sample {
        ($name:ident, |$sb:ident, $shb:ident| $expr:expr) => {
            #[test]
            fn $name() {
                let $sb = BYTES;
                let $shb = HEX_BYTES;
                $expr;

                let $sb = LONG_BYTES;
                let $shb = LONG_HEX_BYTES;
                $expr;
            }
        };
    }

    type Enc = Encoder::<true>;

    for_each_sample!(enc_const, |b, hb| assert_eq!(Enc::enc_const(b), *hb));
    for_each_sample!(enc_sized, |b, hb| assert_eq!(Enc::enc_sized(b), *hb));
    for_each_sample!(enc_sized_heap, |b, hb| assert_eq!(Enc::enc_sized_heap(b), Box::new(*hb)));
    for_each_sample!(enc_slice, |b, hb| assert_eq!(Enc::enc_slice(b), (*hb).into_iter().collect::<Vec<_>>().into_boxed_slice()));
}