Bitfield benchmark

#![allow(
    clippy::cast_sign_loss,
    clippy::cast_possible_truncation,
    clippy::cast_lossless
)]

use std::{ops::RangeInclusive, time::Instant};

use bitvec::order::Lsb0;
use bitvec::view::BitView;
use rand::{thread_rng, Rng};

fn main() {
    benchmark();
}

fn benchmark() {
    let gen_signed =
        |nbits: usize| -> i32 { thread_rng().gen_range(-(1 << (nbits - 1))..1 << (nbits - 1)) };
    let gen_unsigned = |nbits: usize| -> u32 { thread_rng().gen_range(0..1 << nbits) };
    let data = (0..1_000_000)
        .map(|_| {
            (
                gen_unsigned(7) as u8,
                gen_signed(13) as i16,
                gen_unsigned(3) as u8,
                gen_unsigned(5) as u8,
                gen_unsigned(5) as u8,
            )
        })
        .collect::<Vec<_>>();

    let mut merge_results = vec![0; data.len()];
    let start = Instant::now();
    for (i, (opcode, imm, funct3, rs1, rs2)) in data.iter().copied().enumerate() {
        merge_results[i] = riscv_b_instruction_merge(opcode, imm, funct3, rs1, rs2);
    }
    let end = Instant::now();
    println!("merge: {:?}", end - start);

    let mut bitvec_results = vec![0; data.len()];
    let start = Instant::now();
    for (i, (opcode, imm, funct3, rs1, rs2)) in data.iter().copied().enumerate() {
        bitvec_results[i] = riscv_b_instruction_bitvec(opcode, imm, funct3, rs1, rs2);
    }
    let end = Instant::now();
    println!("bitvec: {:?}", end - start);

    let mut manual_results = vec![0; data.len()];
    let start = Instant::now();
    for (i, (opcode, imm, funct3, rs1, rs2)) in data.iter().copied().enumerate() {
        manual_results[i] = riscv_b_instruction_manual(opcode, imm, funct3, rs1, rs2);
    }
    let end = Instant::now();
    println!("manual: {:?}", end - start);

    for ((merge_result, bitvec_result), manual_result) in merge_results
        .into_iter()
        .zip(bitvec_results.into_iter())
        .zip(manual_results.into_iter())
    {
        assert!(
            merge_result == bitvec_result,
            "\nmerge\t{:032b}\nbitvec\t{:032b}\ndiff\t{}\n",
            merge_result,
            bitvec_result,
            diff(merge_result, bitvec_result)
        );
        assert!(
            merge_result == manual_result,
            "\nmerge\t{:032b}\nmanual\t{:032b}\ndiff\t{}\n",
            merge_result,
            manual_result,
            diff(merge_result, manual_result)
        );
    }
}

fn diff(a: u32, b: u32) -> String {
    let diff = a ^ b;
    let mut result = String::with_capacity(32);
    for i in 0..32 {
        let index = (diff >> (31 - i)) & 1;
        result.push(['_', '^'][index as usize]);
    }
    result
}

// #[inline(always)]
#[allow(clippy::inline_always)]
fn riscv_b_instruction_merge(opcode: u8, imm: i16, funct3: u8, rs1: u8, rs2: u8) -> u32 {
    merge_bitfields([
        (0..=6, opcode as u32, 0..=6),
        (7..=7, imm as u32, 11..=11),
        (8..=11, imm as u32, 1..=4),
        (12..=14, funct3 as u32, 0..=2),
        (15..=19, rs1 as u32, 0..=4),
        (20..=24, rs2 as u32, 0..=4),
        (25..=30, imm as u32, 5..=10),
        (31..=31, imm as u32, 12..=12),
    ])
}

// #[inline(always)]
#[allow(clippy::inline_always)]
fn riscv_b_instruction_manual(opcode: u8, imm: i16, funct3: u8, rs1: u8, rs2: u8) -> u32 {
    (opcode as u32) & (0b0111_1111)
        | ((((imm >> 11) as u32) & 1) << 7)
        | ((((imm >> 1) as u32) & 0b1111) << 8)
        | ((funct3 as u32 & 0b111) << 12)
        | ((rs1 as u32 & 0b11111) << 15)
        | ((rs2 as u32 & 0b11111) << 20)
        | ((((imm >> 5) as u32) & 0b11_1111) << 25)
        | (((((imm >> 12) as u32) >> 12) & 1) << 31)
}

// #[inline(always)]
#[allow(clippy::inline_always)]
fn riscv_b_instruction_bitvec(opcode: u8, imm: i16, funct3: u8, rs1: u8, rs2: u8) -> u32 {
    let opcode = opcode as u32;
    let opcode = opcode.view_bits::<Lsb0>();
    let imm = imm as u32;
    let imm = imm.view_bits::<Lsb0>();
    let funct3 = funct3 as u32;
    let funct3 = funct3.view_bits::<Lsb0>();
    let rs1 = rs1 as u32;
    let rs1 = rs1.view_bits::<Lsb0>();
    let rs2 = rs2 as u32;
    let rs2 = rs2.view_bits::<Lsb0>();

    let mut instruction = 0;
    let bits = instruction.view_bits_mut::<Lsb0>();
    bits[0..7].copy_from_bitslice(&opcode[0..7]);
    bits[7..8].copy_from_bitslice(&imm[11..12]);
    bits[8..12].copy_from_bitslice(&imm[1..5]);
    bits[12..15].copy_from_bitslice(&funct3[0..3]);
    bits[15..20].copy_from_bitslice(&rs1[0..5]);
    bits[20..25].copy_from_bitslice(&rs2[0..5]);
    bits[25..31].copy_from_bitslice(&imm[5..11]);
    bits[31..32].copy_from_bitslice(&imm[12..13]);
    instruction
}

/// Combines `bitfields` into a single value. Each bit field is a tuple of:
/// - dst range
/// - value
/// - src range
// FIXME bitfields as array instead of reference
// #[inline(always)]
#[allow(clippy::inline_always)]
pub(crate) const fn merge_bitfields<const N: usize>(
    bitfields: [(RangeInclusive<usize>, u32, RangeInclusive<usize>); N],
) -> u32 {
    let mut dst_bits_visited = 0u32;
    let mut dst = 0;
    let mut i = 0;
    while i < bitfields.len() {
        let (dst_range, src, src_range) = &bitfields[i];
        assert!(
            *src_range.end() < 32 && *dst_range.end() < 32,
            "bit range crosses 32-bit boundary"
        );
        assert!(
            (*dst_range.end() - *dst_range.start()) == *src_range.end() - *src_range.start(),
            "bit range lengths do not match"
        );

        // Copy the bitfield
        dst |= ((*src & !(0xFFFF_FFFF_u32 << *src_range.end() << 1)) >> *src_range.start())
            << *dst_range.start();

        // Check for overlaps
        let dst_mask =
            (0xFFFF_FFFF_u32 << *dst_range.end() << 1) ^ (0xFFFF_FFFF_u32 << *dst_range.start());
        assert!(
            dst_bits_visited & dst_mask == 0,
            "bit field overlap detected",
        );
        dst_bits_visited |= dst_mask;

        i += 1;
    }
    dst
}