ratproto-did 0.0.3

#![doc = include_str!("plc_codec.md")]

use std::{arch::x86_64, mem};

use crate::DidInner;

/// A helper type representing an `Option<DidInner::Plc>`, in a sense.
///
/// Can be turned into a `DidInner::Plc` via `try_into` only if the first byte is 0.
#[repr(transparent)]
pub struct OptionDidPlc([u8; 16]);

impl OptionDidPlc {
    pub const INVALID: OptionDidPlc = {
        let mut val = OptionDidPlc([0; 16]);
        val.0[0] = 1;
        val
    };
}

impl TryFrom<OptionDidPlc> for DidInner {
    type Error = ();

    fn try_from(val: OptionDidPlc) -> Result<Self, Self::Error> {
        // Compile-time check that the first byte of `DidInner::Plc` is a discriminant
        // with a value of 0
        const {
            let plc_val = DidInner::Plc([0xff; 15]);
            // SAFETY: DidInner is 16 bytes
            let bytes = unsafe { mem::transmute::<DidInner, [u8; 16]>(plc_val) };
            assert!(bytes[0] == 0, "The discriminant of `DidInner::Plc` should be 0");
        }

        if val.0[0] == 0 {
            // SAFETY: The discriminant of `DidInner::Plc` is 0
            unsafe { Ok(mem::transmute::<OptionDidPlc, DidInner>(val)) }
        } else {
            Err(())
        }
    }
}

/// Validates and decodes a `did:plc:` string, producing a `DidInner::Plc` if successful.
///
/// [`OptionDidPlc::try_into()`] produces an `Ok(DidInner)` if successful.
#[inline]
pub fn decode_plc(plc_str: &[u8; 32]) -> OptionDidPlc {
    if is_x86_feature_detected!("avx2") {
        // SAFETY: avx2 is detected
        unsafe { decode_plc_avx2(plc_str) }
    } else {
        decode_plc_non_avx(plc_str)
    }
}

/// Validates and decodes a `did:plc:` string, producing a `DidInner::Plc` if successful.
///
/// Uses AVX2 SIMD instructions.
///
/// [`OptionDidPlc::try_into()`] produces an `Ok(DidInner)` if successful.
#[target_feature(enable = "avx2")]
#[inline]
fn decode_plc_avx2(plc_str: &[u8; 32]) -> OptionDidPlc {
    // SAFETY: plc_str is 32 bytes (256 bits)
    let data = unsafe { x86_64::_mm256_loadu_si256(plc_str.as_ptr() as _) };

    // For "did:plc:abcdefghijklmnopqrstuvwx", the debugger shows `data` as:
    // [0] = {i64} 4207325706165971300 [0x3a636c703a646964]
    // [1] = {i64} 7523094288207667809 [0x6867666564636261]
    // [2] = {i64} 8101815670912281193 [0x706f6e6d6c6b6a69]
    // [3] = {i64} 8680537053616894577 [0x7877767574737271]
    // The did:plc string is loaded into the SIMD register from LSB to MSB

    // let did_plc_eq_mask = x86_64::_mm256_cmpeq_epi8(
    //     data,
    //     x86_64::_mm256_set_epi64x(
    //         0, // The others can be 0
    //         0,
    //         0,
    //         0x3a636c703a646964, // "did:plc:" with bytes in reverse order
    //     ),
    // );

    // TODO: rewrite to be more readable (macros, const eval)

    let alpha_mask = {
        x86_64::_mm256_andnot_si256(
            x86_64::_mm256_cmpgt_epi8(
                data,
                x86_64::_mm256_set_epi64x(
                    0x7a7a7a7a7a7a7a7a, // "z" repeated
                    0x7a7a7a7a7a7a7a7a,
                    0x7a7a7a7a7a7a7a7a,
                    0x3a636c703a646964, // "did:plc:" with bytes in reverse order
                ),
            ),
            x86_64::_mm256_cmpgt_epi8(
                data,
                x86_64::_mm256_set_epi64x(
                    0x6060606060606060, // "a" - 1 repeated
                    0x6060606060606060,
                    0x6060606060606060,
                    0x39626b6f39636863, // "did:plc:" as above, just 1 bit lower
                ),
            ),
        )
    };
    // let alpha_mask = x86_64::_mm256_cmpgt_epi8(data, x86_64::_mm256_set1_epi8((b'a' - 1) as _));

    let num_mask = {
        x86_64::_mm256_andnot_si256(
            x86_64::_mm256_cmpgt_epi8(
                data,
                x86_64::_mm256_set_epi64x(
                    0x3737373737373737, // "7" repeated
                    0x3737373737373737,
                    0x3737373737373737,
                    0x3a636c703a646964, // "did:plc:" with bytes in reverse order
                ),
            ),
            x86_64::_mm256_cmpgt_epi8(
                data,
                x86_64::_mm256_set_epi64x(
                    0x3131313131313131, // "2" - 1 repeated
                    0x3131313131313131,
                    0x3131313131313131,
                    0x39626b6f39636863, // "did:plc:" as above, just 1 bit lower
                ),
            ),
        )
    };

    let char_to_val = x86_64::_mm256_blendv_epi8(
        x86_64::_mm256_set1_epi8((b'2' - 26) as i8),
        x86_64::_mm256_set1_epi8(b'a' as i8),
        alpha_mask,
    );

    let values = x86_64::_mm256_sub_epi8(data, char_to_val);

    let is_valid = {
        // alpha and num are masks for a..=z and 2..=7 respectively
        // In addition, they also both check for the "did:plc:" prefix
        let alpha = x86_64::_mm256_movemask_epi8(alpha_mask) as u32;
        let num = x86_64::_mm256_movemask_epi8(num_mask) as u32;
        let base32 = alpha | num;

        base32 == !0 // all ones
    };

    // Current register layout:
    // MSB ____________________________________________________________________________________
    // | 000xxxxx | 000wwwww | 000vvvvv | 000uuuuu | 000ttttt | 000sssss | 000rrrrr | 000qqqqq |
    // | 000ppppp | 000ooooo | 000nnnnn | 000mmmmm | 000lllll | 000kkkkk | 000jjjjj | 000iiiii |
    // | 000hhhhh | 000ggggg | 000fffff | 000eeeee | 000ddddd | 000ccccc | 000bbbbb | 000aaaaa |
    // | 00111010 | 01100011 | 01101100 | 01110000 | 00111010 | 01100100 | 01101001 | 01100100 |
    //  ____________________________________________________________________________________ LSB
    // If the identifier was valid base32, all bytes (a-x) have been converted to 5-bit values
    // (If not, the bytes will contain garbage, but `is_valid` will be set to 0)

    // Permute in order to bring some values to the lower half
    let reg1 = x86_64::_mm256_permute4x64_epi64::<0b11100110>(values); // 3, 2, 1, 2

    // Swizzle to allow for u32 bit shifts
    #[rustfmt::skip]
    let reg1 = x86_64::_mm256_shuffle_epi8(
        reg1,
        x86_64::_mm256_set_epi8(
            15, 14,  7,  6,
            13, 12,  5,  4,
            11, 10,  3,  2,
             9,  8,  1,  0,

            15, 14,  7,  6,
            13, 12,  5,  4,
            11, 10,  3,  2,
             9,  8,  1,  0,
        ),
    );

    // Split u16 halves
    let reg2 = x86_64::_mm256_and_si256(reg1, x86_64::_mm256_set1_epi16(0xff00u16 as i16));
    let reg3 = x86_64::_mm256_and_si256(reg1, x86_64::_mm256_set1_epi16(0x00ffu16 as i16));

    // Bit-shift
    let reg2 = x86_64::_mm256_srlv_epi32(
        reg2,
        x86_64::_mm256_set_epi32(
            8, 6, 4, 2, // (x, p), (v, n), (t, l), (r, j)
            8, 6, 4, 2, // (h, p), (f, n), (d, l), (b, j)
        ),
    );
    let reg3 = x86_64::_mm256_sllv_epi32(
        reg3,
        x86_64::_mm256_set_epi32(
            5, 7, 1, 3, // (w, o), (u, m), (s, k), (q, i)
            5, 7, 1, 3, // (g, o), (e, m), (c, k), (a, i)
        ),
    );

    // Shuffle
    #[rustfmt::skip]
    let reg2 = x86_64::_mm256_shuffle_epi8(
        reg2,
        x86_64::_mm256_set_epi8(
            14, 10,  6,  2,  3, 12,  8,  4, // x, v, t, r, r, p, n, l
            -1, -1, -1,  7, -1, -1, -1, -1, // 0, 0, 0, t, 0, 0, 0, 0
             5, -1, -1, -1, -1,  7, -1, -1, // l, 0, 0, 0, 0, d, 0, 0
             0,  1, 14, 10,  6,  2,  3, -1, // j, j, h, f, d, b, b, 0
        ),
    );

    #[rustfmt::skip]
    let reg3 = x86_64::_mm256_shuffle_epi8(
        reg3,
        x86_64::_mm256_set_epi8(
            14, 15, 11,  6,  2, 12, 13,  9, // w, w, u, s, q, o, o, m
            -1, 10, -1, -1, -1, -1,  8, -1, // 0, u, 0, 0, 0, 0, m, 0
            -1, -1, -1, 10, -1, -1, -1, -1, // 0, 0, 0, e, 0, 0, 0, 0
             4,  0, 14, 15, 11,  6,  2, -1, // k, i, g, g, e, c, a, 0
        ),
    );

    // OR-Reduce
    let reduce1 = x86_64::_mm256_or_si256(reg2, reg3);
    let reduce_hi =
        x86_64::_mm256_castsi256_si128(x86_64::_mm256_permute4x64_epi64::<0b1101>(reduce1));
    let reduce_lo =
        x86_64::_mm256_castsi256_si128(x86_64::_mm256_permute4x64_epi64::<0b1000>(reduce1));

    let reduce2 = x86_64::_mm_or_si128(reduce_hi, reduce_lo);

    // Current register layout:
    // MSB ____________________________________________________________________________________
    // | wwwxxxxx | uvvvvvww | ttttuuuu | rrssssst | qqqqqrrr | oooppppp | mnnnnnoo | llllmmmm |
    // | jjkkkkkl | iiiiijjj | ggghhhhh | efffffgg | ddddeeee | bbcccccd | aaaaabbb | 00000000 |
    //  ____________________________________________________________________________________ LSB
    // This is then written in reverse order, so:
    // MSB ____________________________________________________________________________________
    // | 00000000 | aaaaabbb | bbcccccd | ddddeeee | efffffgg | ggghhhhh | iiiiijjj | jjkkkkkl |
    // | llllmmmm | mnnnnnoo | oooppppp | qqqqqrrr | rrssssst | ttttuuuu | uvvvvvww | wwwxxxxx |
    //  ____________________________________________________________________________________ LSB

    let mut out = OptionDidPlc([0; 16]);

    // SAFETY: `out` is 16 bytes (128 bits)
    unsafe { x86_64::_mm_storeu_si128(out.0.as_mut_ptr() as _, reduce2) };

    out.0[0] = if is_valid { 0 } else { 1 };

    out
}

/// Validates and decodes a `did:plc:` string, producing a `DidInner::Plc` if successful.
///
/// Avoids using AVX instructions, but uses BMI2 if able.
///
/// [`OptionDidPlc::try_into()`] produces an `Ok(DidInner)` if successful.
#[inline]
fn decode_plc_non_avx(plc_str: &[u8; 32]) -> OptionDidPlc {
    let Some(ident) = plc_str.strip_prefix(b"did:plc:") else {
        return OptionDidPlc::INVALID;
    };

    if !ident.iter().all(|&b| matches!(b, b'a'..=b'z' | b'2'..=b'7')) {
        return OptionDidPlc::INVALID;
    }

    let mut out = OptionDidPlc([0u8; 16]);

    #[inline]
    fn pack_bytes(ident_bytes: &[u8]) -> u64 {
        // Note: all ident_bytes must be valid base32 chars! ('a'..='z', '2'..='7')
        debug_assert_eq!(ident_bytes.len(), 8);

        let bytes = u64::from_le_bytes([
            ident_bytes[7],
            ident_bytes[6],
            ident_bytes[5],
            ident_bytes[4],
            ident_bytes[3],
            ident_bytes[2],
            ident_bytes[1],
            ident_bytes[0],
        ]);

        // Here we treat the u64 as packed u8 values
        // There are some add/sub ops, but none of them should overflow within their u8
        // All bytes are already validated when this function is used

        // For reference:
        // b'2' =  50 = 0x32
        // b'7' =  55 = 0x37
        // b'a' =  97 = 0x61
        // b'z' = 122 = 0x7a

        // Chars 'a'..='z' have the 0x40 byte set, while '2'..='7' don't
        let alpha_mask = 0x4040404040404040_u64;

        // alpha_flags has a 0x01 byte for every char that is 'a'..='z'
        let alpha_flags = (bytes & alpha_mask) >> 6;

        // Bring values from alpha chars right "behind" numeric chars
        // That is, if a char was b'z', it should end up as (b'2' - 1)
        // This should not underflow any 8-bit part of the 64bit value
        // In other words, b'z' (0x7a) maps to b'2' - 1 (0x31), so -73 (-0x49)
        let values = bytes - alpha_flags * (b'z' - b'2' + 1) as u64;

        // The character values now represent a contiguous range, but it's not yet zero-based
        // b'z' (0x7a) ended up at 0x31, so b'a' (0x61) is now at 0x18
        let values = values - 0x1818181818181818_u64;

        // At this point, the base32 chars should have all been converted appropriately
        // The value ranges 'a'..='z' and '2'..='7' have been made contiguous,
        // and then shifted to make the range start at 0.

        // Finally, the bits need to be packed.
        // Every 8 bits actually represent only 5 bits.

        // Compile-time detection
        if is_x86_feature_detected!("bmi2") {
            // SAFETY: bmi2 is active
            unsafe { x86_64::_pext_u64(values, 0x1f1f1f1f1f1f1f1f) }
        } else {
            let [h, g, f, e, d, c, b, a] = values.to_le_bytes();

            ((a as u64) << 35)
                | ((b as u64) << 30)
                | ((c as u64) << 25)
                | ((d as u64) << 20)
                | ((e as u64) << 15)
                | ((f as u64) << 10)
                | ((g as u64) << 5)
                | (h as u64)
        }
    }

    debug_assert_eq!(ident.len(), 24);

    for i in 0..3 {
        let from = i * 8;
        // The compiler is not convinced that a regular index access is safe
        // This is a minor optimization, and just gets rid of a few calls to slice_index_fail
        // SAFETY: this is in bounds
        // `from` is 0, 8, 16
        // The indexed ranges are 0..8, 8..16, 16..24
        let ident_slice = unsafe { ident.get_unchecked(from..from + 8) };
        let bytes = pack_bytes(ident_slice).to_le_bytes();
        out.0[i * 5 + 1] = bytes[4];
        out.0[i * 5 + 2] = bytes[3];
        out.0[i * 5 + 3] = bytes[2];
        out.0[i * 5 + 4] = bytes[1];
        out.0[i * 5 + 5] = bytes[0];
    }

    out
}

/// Encodes a [`DidInner::Plc`] into a `did:plc:` string.
///
/// Precondition: `val` must be a [`DidInner::Plc`]
#[allow(dead_code)] // while still WIP
pub fn encode_plc(val: DidInner, out: &mut [u8; 32]) {
    debug_assert!(matches!(val, DidInner::Plc(_)), "Input should be `DidInner::Plc`");

    // SAFETY: DidInner is 16 bytes, and known to be the PLC variant
    // The latter is only debug-asserted locally, but the function is not public API
    let bytes: [u8; 16] = unsafe { mem::transmute::<DidInner, [u8; 16]>(val) };

    if is_x86_feature_detected!("avx2") {
        // SAFETY: avx2 is detected
        unsafe {
            encode_plc_avx2(bytes, out);
        }
    } else {
        encode_plc_non_avx(bytes, out);
    }
}

#[target_feature(enable = "avx2")]
#[inline]
fn encode_plc_avx2(bytes_with_discr: [u8; 16], out: &mut [u8; 32]) {
    // SAFETY: bytes_with_discr is 16 bytes (128 bits)
    let data = unsafe { x86_64::_mm_loadu_si128(bytes_with_discr.as_ptr() as _) };
    let data_x2 = x86_64::_mm256_broadcastsi128_si256(data);

    // Data is loaded in little-endian format (so it gets reversed)
    //  __ 0x00 _______________________________________________________________________________
    // | 00000000 | aaaaabbb | bbcccccd | ddddeeee | efffffgg | ggghhhhh | iiiiijjj | jjkkkkkl |
    // | llllmmmm | mnnnnnoo | oooppppp | qqqqqrrr | rrssssst | ttttuuuu | uvvvvvww | wwwxxxxx |
    //  _______________________________________________________________________________ 0x0f __
    // This is loaded into the register as:
    // MSB ____________________________________________________________________________________
    // | wwwxxxxx | uvvvvvww | ttttuuuu | rrssssst | qqqqqrrr | oooppppp | mnnnnnoo | llllmmmm |
    // | jjkkkkkl | iiiiijjj | ggghhhhh | efffffgg | ddddeeee | bbcccccd | aaaaabbb | 00000000 |
    //  ____________________________________________________________________________________ LSB
    // This is subsequently duplicated into the other 128-bit lane

    // Register layout goal:
    // MSB ____________________________________________________________________________________
    // | ...xxxxx | ...wwwww | ...vvvvv | ...uuuuu | ...ttttt | ...sssss | ...rrrrr | ...qqqqq |
    // | ...ppppp | ...ooooo | ...nnnnn | ...mmmmm | ...lllll | ...kkkkk | ...jjjjj | ...iiiii |
    // | ...hhhhh | ...ggggg | ...fffff | ...eeeee | ...ddddd | ...ccccc | ...bbbbb | ...aaaaa |
    // | 00111010 | 01100011 | 01101100 | 01110000 | 00111010 | 01100100 | 01101001 | 01100100 |
    //  ____________________________________________________________________________________ LSB

    // The result will be assembled from two 256-bit registers, which will hold alternating columns:
    // MSB ____________________________________________________________________________________
    // | ...xxxxx | ........ | ...vvvvv | ........ | ...ttttt | ........ | ...rrrrr | ........ |
    // | ...ppppp | ........ | ...nnnnn | ........ | ...lllll | ........ | ...jjjjj | ........ |
    // | ...hhhhh | ........ | ...fffff | ........ | ...ddddd | ........ | ...bbbbb | ........ |
    // | ........ | ........ | ........ | ........ | ........ | ........ | ........ | ........ |
    //  ____________________________________________________________________________________ LSB
    // MSB ____________________________________________________________________________________
    // | ........ | ...wwwww | ........ | ...uuuuu | ........ | ...sssss | ........ | ...qqqqq |
    // | ........ | ...ooooo | ........ | ...mmmmm | ........ | ...kkkkk | ........ | ...iiiii |
    // | ........ | ...ggggg | ........ | ...eeeee | ........ | ...ccccc | ........ | ...aaaaa |
    // | ........ | ........ | ........ | ........ | ........ | ........ | ........ | ........ |
    //  ____________________________________________________________________________________ LSB

    // However, because AVX2 only supports shifts of packed 32-bit integers (and not 16),
    // Values have to first be grouped up appropriately.

    // Steps:
    // MSB ____________________________________________________________________________________
    // | ........ | wwwxxxxx | ........ | oooppppp | rrssssst | ttttuuuu | jjkkkkkl | llllmmmm |
    // | ........ | uvvvvvww | ........ | mnnnnnoo | qqqqqrrr | rrssssst | iiiiijjj | jjkkkkkl |
    // | ........ | ggghhhhh | ........ | ........ | bbcccccd | ddddeeee | ........ | ........ |
    // | ........ | efffffgg | ........ | ........ | aaaaabbb | bbcccccd | ........ | ........ |
    //  ____________________________________________________________________________________ LSB
    // MSB ____________________________________________________________________________________
    // | ........ | ...xxxxx | ........ | ...ppppp | .......t | tttt.... | .......l | llll.... |
    // | ........ | .vvvvv.. | ........ | .nnnnn.. | .....rrr | rr...... | .....jjj | jj...... |
    // | ........ | ...hhhhh | ........ | ........ | .......d | dddd.... | ........ | ........ |
    // | ........ | .fffff.. | ........ | ........ | .....bbb | bb...... | ........ | ........ |
    //  ____________________________________________________________________________________ LSB
    // MSB ____________________________________________________________________________________
    // | ...xxxxx | ........ | ...ppppp | ........ | ...ttttt | ........ | ...lllll | ........ |
    // | ...vvvvv | ........ | ...nnnnn | ........ | ...rrrrr | ........ | ...jjjjj | ........ |
    // | ...hhhhh | ........ | ........ | ........ | ...ddddd | ........ | ........ | ........ |
    // | ...fffff | ........ | ........ | ........ | ...bbbbb | ........ | ........ | ........ |
    //  ____________________________________________________________________________________ LSB
    // MSB ____________________________________________________________________________________
    // | ...xxxxx | ........ | ...vvvvv | ........ | ...ttttt | ........ | ...rrrrr | ........ |
    // | ...ppppp | ........ | ...nnnnn | ........ | ...lllll | ........ | ...jjjjj | ........ |
    // | ...hhhhh | ........ | ...fffff | ........ | ...ddddd | ........ | ...bbbbb | ........ |
    // | ........ | ........ | ........ | ........ | ........ | ........ | ........ | ........ |
    //  ____________________________________________________________________________________ LSB
    #[rustfmt::skip]
    let half1 = x86_64::_mm256_shuffle_epi8(
        data_x2,
        x86_64::_mm256_set_epi8(
            -1, 15, -1, 10, 12, 13,  7,  8,
            -1, 14, -1,  9, 11, 12,  6,  7,

            -1,  5, -1, -1,  2,  3, -1, -1,
            -1,  4, -1, -1,  1,  2, -1, -1,
        )
    );

    #[rustfmt::skip]
    let half1 = x86_64::_mm256_sllv_epi32(half1, x86_64::_mm256_set_epi32(
        8, 4,
        6, 2,

        8, 4,
        6, 2,
    ));

    #[rustfmt::skip]
    let half1 = x86_64::_mm256_shuffle_epi8(
        half1,
        x86_64::_mm256_set_epi8(
            15, -1,  7, -1, 11, -1,  3, -1,
            13, -1,  5, -1,  9, -1,  1, -1,

            15, -1,  7, -1, 11, -1,  3, -1,
            -1, -1, -1, -1, -1, -1, -1, -1,
        )
    );

    // MSB ____________________________________________________________________________________
    // | ......ww | www..... | ......oo | ooo..... | ..sssss. | ........ | ..kkkkk. | ........ |
    // | ....uuuu | u....... | ....mmmm | m....... | qqqqq... | ........ | iiiii... | ........ |
    // | ......gg | ggg..... | ........ | ........ | ..ccccc. | ........ | ........ | ........ |
    // | ....eeee | e....... | ........ | ........ | aaaaa... | ........ | ........ | ........ |
    //  ____________________________________________________________________________________ LSB
    // MSB ____________________________________________________________________________________
    // | ........ | ...wwwww | ........ | ...ooooo | ........ | ...sssss | ........ | ...kkkkk |
    // | ........ | ...uuuuu | ........ | ...mmmmm | ........ | ...qqqqq | ........ | ...iiiii |
    // | ........ | ...ggggg | ........ | ........ | ........ | ...ccccc | ........ | ........ |
    // | ........ | ...eeeee | ........ | ........ | ........ | ...aaaaa | ........ | ........ |
    //  ____________________________________________________________________________________ LSB
    // MSB ____________________________________________________________________________________
    // | ........ | ...wwwww | ........ | ...uuuuu | ........ | ...sssss | ........ | ...qqqqq |
    // | ........ | ...ooooo | ........ | ...mmmmm | ........ | ...kkkkk | ........ | ...iiiii |
    // | ........ | ...ggggg | ........ | ...eeeee | ........ | ...ccccc | ........ | ...aaaaa |
    // | ........ | ........ | ........ | ........ | ........ | ........ | ........ | ........ |
    //  ____________________________________________________________________________________ LSB
    #[rustfmt::skip]
    let half2 = x86_64::_mm256_shuffle_epi8(
        data_x2,
        x86_64::_mm256_set_epi8(
            14, 15,  9, 10, 12, -1,  7, -1,
            13, 14,  8,  9, 11, -1,  6, -1,

             4,  5, -1, -1,  2, -1, -1, -1,
             3,  4, -1, -1,  1, -1, -1, -1,
        ),
    );
    #[rustfmt::skip]
    let half2 = x86_64::_mm256_srlv_epi32(half2, x86_64::_mm256_set_epi32(
        5, 9,
        7, 11,

        5, 9,
        7, 11,
    ));

    #[rustfmt::skip]
    let half2 = x86_64::_mm256_shuffle_epi8(
        half2,
        x86_64::_mm256_set_epi8(
            -1, 14, -1,  6, -1, 10, -1,  2,
            -1, 12, -1,  4, -1,  8, -1,  0,

            -1, 14, -1,  6, -1, 10, -1,  2,
            -1, -1, -1, -1, -1, -1, -1, -1,
        ),
    );

    let combined = x86_64::_mm256_or_si256(half1, half2);
    let combined = x86_64::_mm256_and_si256(combined, x86_64::_mm256_set1_epi16(0x1f1f));

    let alpha_or_num_mask =
        x86_64::_mm256_cmpgt_epi8(combined, x86_64::_mm256_set1_epi8((b'z' - b'a') as i8));

    let add_vec = x86_64::_mm256_blendv_epi8(
        x86_64::_mm256_set1_epi8(b'a' as i8),
        x86_64::_mm256_set1_epi8((b'2' - (b'z' - b'a') - 1) as i8),
        alpha_or_num_mask,
    );

    let chars = x86_64::_mm256_add_epi8(combined, add_vec);

    // SAFETY: `out` is 32 bytes (256 bits)
    unsafe {
        x86_64::_mm256_storeu_si256(out.as_mut_ptr() as _, chars);
    }

    // TODO: is writing the prefix from the vector register faster?
    out[..8].copy_from_slice(b"did:plc:");
    // dbg!(chars);
    // dbg!(out);
}

#[inline]
fn encode_plc_non_avx(bytes_with_discr: [u8; 16], out: &mut [u8; 32]) {
    // Note: bytes_with_discr includes the zero-byte at the start!
    let bytes = &bytes_with_discr[1..];

    fn byte_to_base32(val: u8) -> u8 {
        match val {
            0..26 => val + b'a',
            26..32 => val - 26 + b'2',
            _ => unreachable!(),
        }
    }

    out[..8].copy_from_slice(b"did:plc:");

    for i in 0..3 {
        let bytes_pos = i * 5;
        let packed = usize::from_le_bytes([
            bytes[bytes_pos + 4],
            bytes[bytes_pos + 3],
            bytes[bytes_pos + 2],
            bytes[bytes_pos + 1],
            bytes[bytes_pos],
            0,
            0,
            0,
        ]);

        let a = byte_to_base32((packed >> 35) as u8 & 0x1f);
        let b = byte_to_base32((packed >> 30) as u8 & 0x1f);
        let c = byte_to_base32((packed >> 25) as u8 & 0x1f);
        let d = byte_to_base32((packed >> 20) as u8 & 0x1f);
        let e = byte_to_base32((packed >> 15) as u8 & 0x1f);
        let f = byte_to_base32((packed >> 10) as u8 & 0x1f);
        let g = byte_to_base32((packed >> 5) as u8 & 0x1f);
        let h = byte_to_base32(packed as u8 & 0x1f);

        let start = 8 + i * 8;
        let end = start + 8;
        out[start..end].copy_from_slice(&[a, b, c, d, e, f, g, h]);
    }
}

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

    #[test]
    #[cfg(target_feature = "avx2")]
    fn individual_bytes_decode_ok_avx2() {
        if !is_x86_feature_detected!("avx2") {
            panic!("AVX2 feature not detected");
        }
        test_individual_bytes_decode(|x| unsafe { decode_plc_avx2(x) });
    }

    #[test]
    fn individual_bytes_decode_ok_non_avx() {
        // TODO: test both the BMI and non-BMI impls
        test_individual_bytes_decode(decode_plc_non_avx);
    }

    /// Tests parsing for every base32 character at every individual position.
    ///
    /// All errors are reported together at the end.
    fn test_individual_bytes_decode<F: Fn(&[u8; 32]) -> OptionDidPlc>(decoder: F) {
        let mut did = "did:plc:aaaaaaaaaaaaaaaaaaaaaaaa".to_string();

        let mut bad_results = vec![];

        for i in 8..32 {
            let base32_alphabet = b"abcdefghijklmnopqrstuvwxyz234567";
            // Test every char except 'a'
            for c in &base32_alphabet[1..] {
                unsafe { did.as_bytes_mut()[i] = *c };

                let result: DidInner = decoder(did.as_bytes().as_array().unwrap())
                    .try_into()
                    .unwrap_or_else(|_| panic!("Decoder failed on {did}"));

                let mut expected_bytes =
                    base32::decode(base32::Alphabet::Rfc4648Lower { padding: false }, &did[8..])
                        .unwrap();

                // Prefix 0
                expected_bytes.insert(0, 0);

                let result_bytes = unsafe { mem::transmute::<DidInner, [u8; 16]>(result) };

                if result_bytes != expected_bytes[..] {
                    bad_results.push((did.to_owned(), result_bytes, expected_bytes));
                }
            }

            unsafe { did.as_bytes_mut()[i] = b'a' }; // reset again
        }

        if !bad_results.is_empty() {
            let mut out = format!("{} error(s):\n", bad_results.len());
            out.push_str("           ");

            // Byte indices
            let byte_indices =
                (00..16).map(|i| format!("{i:02x}")).collect::<Vec<_>>().as_slice().join(", ");
            let ref_did = "did:plc:abcdefghijklmnopqrstuvwx";
            out.push('\n');
            for (did, result, expected) in bad_results {
                out.push_str(&format!("Ref DID:  {ref_did}\n"));
                out.push_str(&format!("DID:      {did}\n"));
                out.push_str(&format!("Indices:   {byte_indices}\n"));
                out.push_str(&format!("Result:   {result:02x?}\n"));
                out.push_str(&format!("Expected: {expected:02x?}\n\n"));
            }
            panic!("{out}");
        }
    }
}