bitpack.rs

#![feature(portable_simd)]

/// A multiplier constant which shifts each bit of a byte into the lowest
/// bit of the corresponding byte in u64.  Constant calculation:
/// `(0..8).fold(0, |acc, x| acc | (1 << ((8 + 1) * x)))`
const PACK_BYTES: u64 = 0x8040201008040201;

/// Retain only the lowest bit in each byte of u64, making it a valid [bool; 8].
/// Constant calculation: `(0..8).fold(0, |acc, x| acc | (1 << (8 * x)))`
const BOOL_MASK:  u64 = 0x0101010101010101;

pub fn bitpack(value: [bool; 8]) -> u8 {
    let bytes = unsafe { core::mem::transmute::<[bool; 8], [u8; 8]>(value) };
    let inter = u64::from_le_bytes(bytes).wrapping_mul(PACK_BYTES);
    (inter >> 56) as u8
}

pub fn unbitpack(value: u8) -> [bool; 8] {
    let unpacked = (value as u64).wrapping_mul(PACK_BYTES) >> 7;
    let bytes = (unpacked & BOOL_MASK).to_le_bytes();
    unsafe { core::mem::transmute::<[u8; 8], [bool; 8]>(bytes) }
}

use std::simd::{LaneCount, Simd, SupportedLaneCount};

pub fn bitpack_simd<const L: usize>(value: [[bool; 8]; L]) -> [u8; L]
    where LaneCount<L>: SupportedLaneCount
{
    let arr = value.map(|x| {
        let bytes = unsafe { core::mem::transmute::<[bool; 8], [u8; 8]>(x) };
        u64::from_le_bytes(bytes)
    });
    let lanes = Simd::from_array(arr);
    let bytes = ((lanes * Simd::splat(PACK_BYTES)) >> Simd::splat(56)).cast::<u8>();
    bytes.to_array()
}

pub fn unbitpack_simd<const L: usize>(value: [u8; L]) -> [[bool; 8]; L]
    where LaneCount<L>: SupportedLaneCount
{
    let bytes = Simd::from_array(value);
    let value = bytes.cast::<u64>() * Simd::splat(PACK_BYTES);
    let value = (value >> Simd::splat(7)) & Simd::splat(BOOL_MASK);
    value.to_array().map(|x| {
        unsafe { core::mem::transmute::<[u8; 8], [bool; 8]>(x.to_le_bytes()) }
    })
}

pub fn bitpack_64(value: [bool; 64]) -> u64 {
    let bytes = unsafe { core::mem::transmute::<[bool; 64], [[u8; 8]; 8]>(value) };
    let lanes = Simd::from_array(bytes.map(|x| u64::from_le_bytes(x)));
    let bytes = ((lanes * Simd::splat(PACK_BYTES)) >> Simd::splat(56)).cast::<u8>();
    u64::from_le_bytes(bytes.to_array())
}

pub fn unbitpack_64(value: u64) -> [bool; 64] {
    let bytes = Simd::from_array(value.to_le_bytes());
    let value = bytes.cast::<u64>() * Simd::splat(PACK_BYTES);
    let value = (value >> Simd::splat(7)) & Simd::splat(BOOL_MASK);
    let bytes = value.to_array().map(|p| p.to_le_bytes());
    unsafe { core::mem::transmute::<[[u8; 8]; 8], [bool; 64]>(bytes) }
}

pub fn main() {
    for b in 0..=255_u8 {
        assert!(b == bitpack(unbitpack(b)));
        let big = u64::from_le_bytes([b; 8]);
        assert!(big == bitpack_64(unbitpack_64(big)));
        assert!(big.to_le_bytes() == bitpack_simd(unbitpack_simd(big.to_le_bytes())));
    }
}