vu128: Efficient variable-length integers

// Copyright (c) John Millikin <john@john-millikin.com>
// SPDX-License-Identifier: 0BSD

macro_rules! encode_uNN {
	($name:ident, $t:ident, $max_len:literal, $len_mask:literal) => {
		#[inline]
		pub fn $name(buf: &mut [u8; $max_len], x: $t) -> usize {
			if x < 0xF0 {
				buf[0] = x as u8;
				return 1;
			}
			unsafe {
				(buf as *mut [u8; $max_len])
					.cast::<u8>().add(1).cast::<$t>()
					.write_unaligned(x.to_le());
			}
			let len = ((x.leading_zeros() >> 3) as u8) ^ $len_mask;
			buf[0] = 0xF0 | len;
			usize::from(len + 2)
		}
	};
}

macro_rules! decode_uNN {
	($name:ident, $t:ident, $max_len:literal, $len_mask:literal) => {
		#[inline]
		pub fn $name(buf: &[u8; $max_len]) -> ($t, usize) {
			if buf[0] < 0xF0 {
				return ($t::from(buf[0]), 1);
			}
			let x = $t::from_le(unsafe {
				(buf as *const [u8; $max_len])
					.cast::<u8>().add(1).cast::<$t>()
					.read_unaligned()
			});
			let len = buf[0] & 0x0F;
			let mask = $t::MAX >> ((len & $len_mask) ^ $len_mask);
			(x & mask, usize::from(len + 2))
		}
	};
}

encode_uNN!(encode_u16, u16, 3, 0x01);
encode_uNN!(encode_u32, u32, 5, 0x03);
encode_uNN!(encode_u64, u64, 9, 0x07);
encode_uNN!(encode_u128, u128, 17, 0x0F);

decode_uNN!(decode_u16, u16, 3, 0x01);
decode_uNN!(decode_u32, u32, 5, 0x03);
decode_uNN!(decode_u64, u64, 9, 0x07);
decode_uNN!(decode_u128, u128, 17, 0x0F);

#[inline]
pub fn encode_u8(buf: &mut [u8; 2], x: u8) -> usize {
	if x < 0xF0 {
		buf[0] = x;
		return 1;
	}
	buf[0] = 0xF0;
	buf[1] = x;
	2
}

#[inline]
pub fn decode_u8(buf: &[u8; 2]) -> (u8, usize) {
	if buf[0] < 0xF0 {
		return (buf[0], 1);
	}
	(buf[1], 2)
}

encode_u32:
        cmp     esi, 240
        jae     .LBB1_1
        mov     eax, 1
        mov     byte ptr [rdi], sil
        ret
.LBB1_1:
        mov     dword ptr [rdi + 1], esi
        lzcnt   esi, esi
        shr     esi, 3
        mov     al, 5
        sub     al, sil
        xor     sil, -13
        movzx   eax, al
        mov     byte ptr [rdi], sil
        ret

decode_u32:
        movzx   eax, byte ptr [rdi]
        cmp     eax, 240
        jae     .LBB0_1
        mov     edx, 1
        ret
.LBB0_1:
        mov     ecx, eax
        not     cl
        and     cl, 3
        mov     edx, 32
        sub     edx, ecx
        bzhi    ecx, dword ptr [rdi + 1], edx
        and     al, 15
        add     al, 2
        movzx   edx, al
        mov     eax, ecx
        ret

macro_rules! encode_iNN {
	(
		$name:ident,
		$encode_fn:ident,
		$iNN:ident,
		$uNN:ident,
		$max_len:literal,
		$zigzag_shift:literal
	) => {
		#[inline]
		pub fn $name(buf: &mut [u8; $max_len], x: $iNN) -> usize {
			let zigzag = ((x >> $zigzag_shift) as $uNN) ^ ((x << 1) as $uNN);
			$encode_fn(buf, zigzag)
		}
	};
}

macro_rules! decode_iNN {
	(
		$name:ident,
		$decode_fn:ident,
		$iNN:ident,
		$uNN:ident,
		$max_len:literal
	) => {
		#[inline]
		pub fn $name(buf: &[u8; $max_len]) -> ($iNN, usize) {
			let (zz, len) = $decode_fn(buf);
			let x = ((zz >> 1) as $iNN) ^ (-((zz & 1) as $iNN));
			(x, len)
		}
	};
}

encode_iNN!(encode_i8, encode_u8, i8, u8, 2, 7);
encode_iNN!(encode_i16, encode_u16, i16, u16, 3, 15);
encode_iNN!(encode_i32, encode_u32, i32, u32, 5, 31);
encode_iNN!(encode_i64, encode_u64, i64, u64, 9, 63);
encode_iNN!(encode_i128, encode_u128, i128, u128, 17, 127);

decode_iNN!(decode_i8, decode_u8, i8, u8, 2);
decode_iNN!(decode_i16, decode_u16, i16, u16, 3);
decode_iNN!(decode_i32, decode_u32, i32, u32, 5);
decode_iNN!(decode_i64, decode_u64, i64, u64, 9);
decode_iNN!(decode_i128, decode_u128, i128, u128, 17);

#[inline]
pub fn encode_f32(buf: &mut [u8; 5], x: f32) -> usize {
	encode_u32(buf, x.to_bits().swap_bytes())
}

#[inline]
pub fn encode_f64(buf: &mut [u8; 9], x: f64) -> usize {
	encode_u64(buf, x.to_bits().swap_bytes())
}

#[inline]
pub fn decode_f32(buf: &[u8; 5]) -> (f32, usize) {
	let (swapped, len) = decode_u32(buf);
	(f32::from_bits(swapped.swap_bytes()), len)
}

#[inline]
pub fn decode_f64(buf: &[u8; 9]) -> (f64, usize) {
	let (swapped, len) = decode_u64(buf);
	(f64::from_bits(swapped.swap_bytes()), len)
}