//! Implementation of the RV32I ISA. Each opcode has a comment explaining it.
use anyhow::{anyhow, Result};

pub struct RV32 {
    regs: [u32; 32], // registers
    mem: Vec<u8>,    // memory
    pc: u32,
}

impl RV32 {
    pub fn new(mem: Vec<u8>) -> Self {
        Self {
            regs: [0; 32],
            pc: 0,
            mem,
        }
    }
    pub fn set_register(&mut self, index: usize, value: u32) {
        assert!(index < 32, "register index must <32");
        if index != 0 {
            self.regs[index] = value;
        }
    }
    // validate the access
    fn range(&self, address: u32, size: usize) -> Result<std::ops::Range<usize>> {
        if !address.is_multiple_of(size as u32) {
            return Err(anyhow!("misaligned at address: {}", address));
        }
        let start = address as usize;
        let end = start
            .checked_add(size)
            .filter(|&end| end <= self.mem.len())
            .ok_or(anyhow!("access fault, address: {}", address))?;
        Ok(start..end)
    }
    fn read(&self, address: u32, size: usize) -> Result<u32> {
        let range = self.range(address, size)?;
        let mut bytes = [0; 4];
        bytes[..size].copy_from_slice(&self.mem[range]);
        Ok(u32::from_le_bytes(bytes))
    }
    pub fn step(&mut self) -> Result<()> {
        let instruction = self.read(self.pc, 4)?;
        let mut next_pc = self.pc.wrapping_add(4);

        // R-type instruction:
        //                             R-type:
        // 31        25 24        20 19        15 14    12 11     7 6        0
        // ┌───────────┬────────────┬────────────┬────────┬────────┬─────────┐
        // │   funct7  │    rs2     │    rs1     │ funct3 │   rd   │ opcode  │
        // └───────────┴────────────┴────────────┴────────┴────────┴─────────┘
        let rd = ((instruction >> 7) & 31) as usize;
        let funct3 = (instruction >> 12) & 7;
        let rs1 = ((instruction >> 15) & 31) as usize;
        let rs2 = ((instruction >> 20) & 31) as usize;
        let funct7 = instruction >> 25;
        //                             I-type:
        // 31                      20 19        15 14    12 11    7 6        0
        // ┌────────────────────────┬────────────┬────────┬────────┬─────────┐
        // │     imm[11:0]          │    rs1     │ funct3 │   rd   │ opcode  │
        // └────────────────────────┴────────────┴────────┴────────┴─────────┘
        //                             S-type:
        // 31        25 24        20 19        15 14    12 11     7 6        0
        // ┌───────────┬────────────┬────────────┬────────┬────────┬─────────┐
        // │ imm[11:5] │    rs2     │    rs1     │ funct3 │imm[4:0]│ opcode  │
        // └───────────┴────────────┴────────────┴────────┴────────┴─────────┘
        //                             B-type:
        // 31    30       25 24   20 19        15 14    12 11     8 7 6      0
        // ┌─────┬─────────┬────────┬────────────┬────────┬────────┬──┬──────┐
        // │imm12│imm[10:5]│ rs2    │    rs1     │ funct3 │imm[4:1]│11│opcode│
        // └─────┴─────────┴────────┴────────────┴────────┴────────┴──┴──────┘
        //
        // Note: we keep `imm` as the I-type imm, and in the S-type we compute
        // the S-type respective imm from 31-25 and 11-7 positions, and for
        // B-type we obtain imm from the positions 31, 31-25, 11-8, 7, and
        // imm[0]=0.
        let imm = sign_extend(instruction >> 20, 12);

        let a = self.regs[rs1];
        let b = self.regs[rs2];
        let mut result: Option<u32> = None;

        // 0x7f=0111_1111, match the lowest 7 bits to read the opcode
        match instruction & 0x7f {
            // I-type op'ns, loads:
            0x03 => {
                let (size, signed) = match funct3 {
                    // LB: load a byte at rs1 + immediate and sign-extend it
                    // into rd
                    0 => (1, true),
                    // LH: load a halfword at rs1 + immediate and sign-extend it
                    // into rd
                    1 => (2, true),
                    // LW: load a 32-bit word at rs1 + immediate into rd
                    2 => (4, false),
                    // LBU: load a byte at rs1 + immediate and zero-extend it
                    // into rd
                    4 => (1, false),
                    // LHU: load a halfword at rs1 + immediate and zero-extend
                    // it into rd
                    5 => (2, false),
                    _ => {
                        return Err(anyhow!(
                            "illegal instruction in I-type (load): {}",
                            instruction
                        ))
                    }
                };
                let value = self.read(a.wrapping_add(imm), size)?;
                result = Some(if signed {
                    sign_extend(value, size as u32 * 8)
                } else {
                    value
                });
            }
            // I-type arithmetic & logical & shifts op'ns:
            0x13 => {
                result = Some(match funct3 {
                    // ADDI: add the sign-extended immediate to rs1, wrapping to
                    // 32 bits
                    0 => a.wrapping_add(imm),
                    // SLTI: signed comparison. Set rd to one if rs1 is less
                    // than the immediate as signed values
                    2 => u32::from((a as i32) < (imm as i32)),
                    // SLTIU: unsigned comparison. Compare rs1 and the
                    // sign-extended immediate as unsigned values
                    3 => u32::from(a < imm),
                    // XORI: write the bitwise XOR of rs1 and the sign-extended
                    // immediate
                    4 => a ^ imm,
                    // ORI: write the bitwise OR of rs1 and the sign-extended
                    // immediate
                    6 => a | imm,
                    // ANDI: write the bitwise AND of rs1 and the sign-extended
                    // immediate
                    7 => a & imm,
                    // SLLI: shift rs1 left by the five-bit immediate shift
                    // amount
                    1 if funct7 == 0 => a << rs2,
                    // SRLI: logical right shift. Shift rs1 right by the
                    // immediate amount, filling with zeros
                    5 if funct7 == 0 => a >> rs2,
                    // SRAI: arithmetic right shift. Shift rs1 right by the
                    // immediate amount, preserving its sign
                    5 if funct7 == 0x20 => ((a as i32) >> rs2) as u32,
                    _ => {
                        return Err(anyhow!(
                            "illegal instruction in I-type (arithmetic/logical): {}",
                            instruction
                        ))
                    }
                });
            }
            // S-type op'ns:
            0x23 => {
                let size = match funct3 {
                    // SB: store the low byte of rs2 at rs1 + immediate
                    0 => 1,
                    // SH: store the low halfword of rs2 at rs1 + immediate
                    1 => 2,
                    // SW: store the full 32-bit value of rs2 at rs1 + immediate
                    2 => 4,
                    _ => return Err(anyhow!("illegal instruction in S-type: {}", instruction)),
                };
                // for the S-type imm, get it from 25-31 and 7-11 positions
                let s_type_imm = ((instruction >> 25) << 5) | ((instruction >> 7) & 31);
                let offset = sign_extend(s_type_imm, 12);
                let range = self.range(a.wrapping_add(offset), size)?;
                self.mem[range].copy_from_slice(&b.to_le_bytes()[..size]);
            }

            // R-type (register-register) integer operations
            0x33 => {
                result = Some(match (funct7, funct3) {
                    // ADD: add rs1 and rs2, wrapping to 32 bits
                    (0, 0) => a.wrapping_add(b),
                    // SUB: subtract rs2 from rs1, wrapping to 32 bits
                    (0x20, 0) => a.wrapping_sub(b),
                    // SLL: shift rs1 left by the low five bits of rs2
                    (0, 1) => a << (b & 31),
                    // SLT: set rd to one if rs1 is less than rs2 as signed
                    // values, else zero
                    (0, 2) => u32::from((a as i32) < (b as i32)),
                    // SLTU: set rd to one if rs1 is less than rs2 as unsigned
                    // values, else zero
                    (0, 3) => u32::from(a < b),
                    // XOR: write the bitwise XOR of rs1 and rs2 into rd
                    (0, 4) => a ^ b,
                    // SRL: shift rs1 right by the low five bits of rs2, filling
                    // with zeros
                    (0, 5) => a >> (b & 31),
                    // SRA: shift rs1 right by the low five bits of rs2,
                    // preserving its sign
                    (0x20, 5) => ((a as i32) >> (b & 31)) as u32,
                    // OR: write the bitwise OR of rs1 and rs2 into rd
                    (0, 6) => a | b,
                    // AND: write the bitwise AND of rs1 and rs2 into rd
                    (0, 7) => a & b,

                    // Note: MUL,DIV,REM are not part of RV32I, but they are in
                    // the M extension.

                    // MUL: write the low 32 bits of the product of rs1 and rs2
                    (1, 0) => a.wrapping_mul(b),
                    // DIV: signed quotient toward zero; zero divisors return
                    // all ones
                    (1, 4) => {
                        if b == 0 {
                            u32::MAX
                        } else {
                            (a as i32).wrapping_div(b as i32) as u32
                        }
                    }
                    // REM: signed remainder; zero divisors return the dividend
                    (1, 6) => {
                        if b == 0 {
                            a
                        } else {
                            (a as i32).wrapping_rem(b as i32) as u32
                        }
                    }

                    _ => return Err(anyhow!("illegal instruction in R-type: {}", instruction)),
                });
            }
            // B-type operations
            0x63 => {
                let taken = match funct3 {
                    // BEQ: branch by the signed PC-relative offset if the
                    // registers are equal
                    0 => a == b,
                    // BNE: branch if the registers are unequal
                    1 => a != b,
                    // BLT: branch if rs1 is less than rs2 using signed comparison
                    4 => (a as i32) < (b as i32),
                    // BGE: branch if rs1 is at least rs2 using signed comparison
                    5 => (a as i32) >= (b as i32),
                    // BLTU: branch if rs1 is less than rs2 using unsigned comparison
                    6 => a < b,
                    // BGEU: branch if rs1 is at least rs2 using unsigned comparison
                    7 => a >= b,
                    _ => return Err(anyhow!("illegal instruction in B-type: {}", instruction)),
                };
                if taken {
                    // offset: reconstruct imm for B-type from positions 31,
                    // 31-25, 11-8, 7, and imm[0]=0
                    let b_type_imm = ((instruction >> 31) << 12)
                        | (((instruction >> 7) & 1) << 11)
                        | (((instruction >> 25) & 0x3f) << 5)
                        | (((instruction >> 8) & 15) << 1);
                    next_pc = self.pc.wrapping_add(sign_extend(b_type_imm, 13));
                }
            }

            _ => return Err(anyhow!("illegal instruction: {}", instruction)),
        }

        if next_pc & 3 != 0 {
            return Err(anyhow!("misaligned, address: {}", next_pc,));
        }
        if let Some(value) = result {
            self.set_register(rd, value);
        }
        self.pc = next_pc;
        Ok(())
    }
}

// sign-extend a smaller signed number to 32 bits
fn sign_extend(value: u32, bits: u32) -> u32 {
    ((value << (32 - bits)) as i32 >> (32 - bits)) as u32
}

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

    const OPCODE_R: u32 = 0b0110011;
    const OPCODE_LOAD: u32 = 0b0000011;

    // encodes R-type instruction
    fn encode_r(name: &str, rd: u32, rs1: u32, rs2: u32) -> u32 {
        let (funct7, funct3) = match name {
            "add" => (0b0000000, 0b000),
            "sub" => (0b0100000, 0b000),
            "sll" => (0b0000000, 0b001),
            "slt" => (0b0000000, 0b010),
            "sltu" => (0b0000000, 0b011),
            "xor" => (0b0000000, 0b100),
            "srl" => (0b0000000, 0b101),
            "sra" => (0b0100000, 0b101),
            "or" => (0b0000000, 0b110),
            "and" => (0b0000000, 0b111),
            // M extension
            "mul" => (0b0000001, 0b000),
            "div" => (0b0000001, 0b100),
            "rem" => (0b0000001, 0b110),
            _ => panic!("unknown R-type instruction: {name}"),
        };

        (funct7 << 25) | (rs2 << 20) | (rs1 << 15) | (funct3 << 12) | (rd << 7) | OPCODE_R
    }
    fn encode_i(name: &str, rd: u32, rs1: u32, imm: i32) -> u32 {
        let (funct3, opcode) = match name {
            // loads: opcode = 0x03
            "lb" => (0b000, OPCODE_LOAD),
            "lh" => (0b001, OPCODE_LOAD),
            "lw" => (0b010, OPCODE_LOAD),
            "lbu" => (0b100, OPCODE_LOAD),
            "lhu" => (0b101, OPCODE_LOAD),
            // arithmetic & logic & shifts: opcode = 0x13
            "addi" => (0b000, 0x13),
            "slti" => (0b010, 0x13),
            "sltiu" => (0b011, 0x13),
            "xori" => (0b100, 0x13),
            "ori" => (0b110, 0x13),
            "andi" => (0b111, 0x13),
            "slli" => (0b001, 0x13),
            "srli" => (0b101, 0x13),
            "srai" => (0b101, 0x13),

            _ => panic!("unknown I-type instruction: {name}"),
        };
        let funct7 = if name == "srai" {
            0b010_0000
        } else {
            0b000_0000
        };

        let imm = (imm as u32) & 0xfff;

        (imm << 20) | (funct7 << 25) | (rs1 << 15) | (funct3 << 12) | (rd << 7) | opcode
    }
    fn encode_s(name: &str, rs1: u32, rs2: u32, imm: i32) -> u32 {
        let (opcode, funct3) = match name {
            "sb" => (0b0100011, 0b000),
            "sh" => (0b0100011, 0b001),
            "sw" => (0b0100011, 0b010),
            _ => panic!("unknown store instruction: {name}"),
        };

        let imm = imm as u32;

        let imm_11_5 = (imm >> 5) & 0x7f;
        let imm_4_0 = imm & 0x1f;

        (imm_11_5 << 25) | (rs2 << 20) | (rs1 << 15) | (funct3 << 12) | (imm_4_0 << 7) | opcode
    }
    fn encode_b(name: &str, rs1: u32, rs2: u32, offset: i32) -> u32 {
        let funct3 = match name {
            "beq" => 0b000,
            "bne" => 0b001,
            "blt" => 0b100,
            "bge" => 0b101,
            "bltu" => 0b110,
            "bgeu" => 0b111,
            _ => panic!("unknown B-type instruction: {name}"),
        };

        assert!(
            offset % 2 == 0,
            "B-type branch offset must be 2-byte aligned"
        );

        assert!(
            (-4096..=4094).contains(&offset),
            "B-type branch offset out of range: {offset}"
        );

        let imm = offset as u32;

        let imm12 = (imm >> 12) & 0x1;
        let imm11 = (imm >> 11) & 0x1;
        let imm10_5 = (imm >> 5) & 0x3f;
        let imm4_1 = (imm >> 1) & 0xf;

        (imm12 << 31)
            | (imm10_5 << 25)
            | (rs2 << 20)
            | (rs1 << 15)
            | (funct3 << 12)
            | (imm4_1 << 8)
            | (imm11 << 7)
            | 0b1100011
    }

    // R-type ops test
    #[test]
    fn test_ops_r_type() {
        let test_vectors = vec![
            // (op, rs1, rs2, rd)
            ("add", 10, 20, 30),
            ("add", u32::MAX, 1, 0), // overflow
            ("sub", 30, 20, 10),
            ("sub", 0, 1, u32::MAX), // overflow
            ("sll", 1, 4, 16),
            ("sll", 1, 31, 0x8000_0000), // SLL maximum shift
            ("sll", 1, 32, 1),           // SLL shift wrapping
            ("slt", (-10i32) as u32, 20, 1),
            ("sltu", 10, 20, 1),
            ("slt", 0xffff_ffff, 0, 1), // signed comparison: -1 < 0
            ("slt", 0xffff_ffff, 0, 1), // signed vs unsigned distinction
            ("sltu", 0xffff_ffff, 0, 0),
            ("xor", 0b1010, 0b1100, 0b0110),
            ("srl", 0x8000_0000, 4, 0x0800_0000),
            ("srl", 0x8000_0000, 1, 0x4000_0000), // logical shift should zero-fill
            ("sra", 0x8000_0000, 4, 0xf800_0000),
            ("sra", 0x8000_0000, 1, 0xc000_0000), // arithmetic shift should sign-fill
            ("sra", 0xffff_ffff, 31, 0xffff_ffff), // arithmetic shift of -1 stays -1
            ("or", 0b1010, 0b1100, 0b1110),
            ("and", 0b1010, 0b1100, 0b1000),
        ];
        test_ops_r_types_opt(test_vectors);
    }
    #[test]
    fn test_ops_m_extension() {
        let test_vectors = vec![
            ("mul", 10, 20, 200),
            ("mul", 0, 123, 0),
            ("mul", 1, u32::MAX, u32::MAX),
            ("mul", u32::MAX, 2, u32::MAX - 1), // result wraps to lower 32 bits
            ("mul", (-1i32) as u32, 10, (-10i32) as u32), // -1 * 10 = -10
            ("mul", (-1i32) as u32, (-1i32) as u32, 1), // -1 * -1 = 1
            ("mul", 0x8000_0000, 2, 0),         // INT_MIN * 2 = 0 modulo 2^32
            ("mul", 0x8000_0000, 0xffff_ffff, 0x8000_0000), // INT_MIN * -1 also wraps to INT_MIN
            ("div", 20, 5, 4),
            ("div", 21, 5, 4), // truncates toward zero
            ("div", 0, 10, 0),
            ("div", 10, 1, 10),
            ("div", 10, (-1i32) as u32, (-10i32) as u32),
            ("div", (-20i32) as u32, 5, (-4i32) as u32), // negative dividend
            ("div", 20, (-5i32) as u32, (-4i32) as u32), // negative divisor
            ("div", (-20i32) as u32, (-5i32) as u32, 4), // negative / negative
            ("div", (-21i32) as u32, 5, (-4i32) as u32), // truncation toward zero: -21 / 5 = -4
            // RISC-V special case:
            ("div", 0x8000_0000, 0xffff_ffff, 0x8000_0000), // INT_MIN / -1 = INT_MIN (overflow)
            ("div", 10, 0, u32::MAX), // division by zero returns -1 for signed DIV
            ("div", 0x8000_0000, 0, u32::MAX),
            ("rem", 20, 5, 0),
            ("rem", 21, 5, 1),
            ("rem", 0, 10, 0),
            ("rem", 10, 1, 0),
            ("rem", 23, 5, 3),                           // positive remainder
            ("rem", (-23i32) as u32, 5, (-3i32) as u32), // remainder follows the sign of the dividend
            ("rem", 23, (-5i32) as u32, 3), // negative divisor does not change the sign of the remainder
            ("rem", (-23i32) as u32, (-5i32) as u32, (-3i32) as u32), // negative / negative
            ("rem", 3, 5, 3),               // dividend smaller than divisor
            // RISC-V special case:
            ("rem", 0x8000_0000, 0xffff_ffff, 0), // INT_MIN % -1 = 0.
            ("rem", 123, 0, 123),                 // remainder by zero returns the dividend.
            ("rem", u32::MAX, 0, u32::MAX),
            ("rem", 0x8000_0000, 0, 0x8000_0000),
        ];
        test_ops_r_types_opt(test_vectors);
    }
    fn test_ops_r_types_opt(test_vectors: Vec<(&str, u32, u32, u32)>) {
        for (name, rs1_value, rs2_value, expected) in test_vectors {
            let rd = 3;
            let rs1 = 1;
            let rs2 = 2;

            let instruction = encode_r(name, rd, rs1, rs2);

            let mut cpu = RV32::new(instruction.to_le_bytes().to_vec());

            cpu.regs[rs1 as usize] = rs1_value;
            cpu.regs[rs2 as usize] = rs2_value;

            cpu.step().unwrap();

            assert_eq!(
                cpu.regs[rd as usize], expected,
                "{name} failed:
                 rs1={rs1_value:#010x}, rs2={rs2_value:#010x}, 
                 expected={expected:#010x}, got={:#010x}",
                cpu.regs[rd as usize],
            );
        }
    }

    // I-type ops test
    #[test]
    fn test_ops_load() {
        let test_vectors: Vec<(&str, u32, i32, u32)> = vec![
            ("lb", 0x0000_0080, 0, 0xffff_ff80),
            ("lbu", 0x0000_0080, 0, 0x0000_0080),
            ("lb", 0x0000_00ff, 0, 0xffff_ffff),
            ("lbu", 0x0000_00ff, 0, 0x0000_00ff),
            ("lh", 0x0000_8000, 0, 0xffff_8000),
            ("lhu", 0x0000_8000, 0, 0x0000_8000),
            ("lh", 0x0000_ffff, 0, 0xffff_ffff),
            ("lhu", 0x0000_ffff, 0, 0x0000_ffff),
        ];

        for (name, memory_value, address, expected) in test_vectors {
            let rd = 1;
            let rs1 = 2;

            let instruction = encode_i(name, rd, rs1, address);

            let mut memory = instruction.to_le_bytes().to_vec();
            let data_address = memory.len() as u32;
            memory.extend_from_slice(&memory_value.to_le_bytes());
            let mut cpu = RV32::new(memory);

            cpu.regs[rs1 as usize] = data_address;

            cpu.step().unwrap();

            assert_eq!(
                cpu.regs[rd as usize], expected,
                "{name} failed: memory={memory_value:#010x},
                 address={address}, \
                 expected={expected:#010x},
                 got={:#010x}",
                cpu.regs[rd as usize],
            );
        }
    }
    #[test]
    fn test_ops_i_type_arithmetic() {
        // almost same test vectors as for `test_ops_r_type`, same names with
        // `i` as suffix
        let test_vectors: Vec<(&str, u32, i32, u32)> = vec![
            // (name, rs1, immediate, expected rd)
            ("addi", 10, 20, 30),
            ("addi", u32::MAX, 1, 0), // overflow
            ("slli", 1, 4, 16),
            ("slli", 1, 31, 0x8000_0000), // SLL maximum shift
            ("slti", (-10i32) as u32, 20, 1),
            ("sltiu", 10, 20, 1),
            ("slti", 0xffff_ffff, 0, 1), // signed comparison: -1 < 0
            ("slti", 0xffff_ffff, 0, 1), // signed vs unsigned distinction
            ("sltiu", 0xffff_ffff, 0, 0),
            ("xori", 0b1010, 0b1100, 0b0110),
            ("srli", 0x8000_0000, 4, 0x0800_0000),
            ("srli", 0x8000_0000, 1, 0x4000_0000), // logical shift should zero-fill
            ("srai", 0x8000_0000, 4, 0xf800_0000),
            ("srai", 0x8000_0000, 1, 0xc000_0000), // arithmetic shift should sign-fill
            ("srai", 0xffff_ffff, 31, 0xffff_ffff), // arithmetic shift of -1 stays -1
            ("ori", 0b1010, 0b1100, 0b1110),
            ("andi", 0b1010, 0b1100, 0b1000),
        ];

        for (name, rs1_value, imm, expected) in test_vectors {
            dbg!(name);
            let rd = 1;
            let rs1 = 2;

            let instruction = encode_i(name, rd, rs1, imm);
            let mut cpu = RV32::new(instruction.to_le_bytes().to_vec());

            cpu.regs[rs1 as usize] = rs1_value;

            cpu.step().unwrap();

            assert_eq!(
                cpu.regs[rd as usize], expected,
                "{name} failed: rs1={rs1_value:#010x}, \
             imm={imm}, \
             expected={expected:#010x}, \
             got={:#010x}",
                cpu.regs[rd as usize],
            );
        }
    }

    // S-type ops test
    #[test]
    fn test_ops_store() {
        let test_vectors: Vec<(&str, u32, i32)> = vec![
            ("sb", 0x0000_0080, 0),
            ("sb", 0x0000_00ff, 0),
            ("sb", 0x1234_5678, 0),
            ("sh", 0x0000_8000, 0),
            ("sh", 0x0000_ffff, 0),
            ("sw", 0x1234_5678, 0),
        ];

        for (name, value, address) in test_vectors {
            let rs1 = 1;
            let rs2 = 2;

            let instruction = encode_s(name, rs1, rs2, address);

            let mut memory = instruction.to_le_bytes().to_vec();
            let data_address = memory.len() as u32;

            // add space for the value being stored
            memory.extend_from_slice(&[0u8; 4]);

            let mut cpu = RV32::new(memory);

            cpu.regs[rs1 as usize] = data_address;
            cpu.regs[rs2 as usize] = value;

            cpu.step().unwrap();

            let stored = u32::from_le_bytes(
                cpu.mem[data_address as usize..data_address as usize + 4]
                    .try_into()
                    .unwrap(),
            );

            let expected = match name {
                "sb" => value & 0xff,
                "sh" => value & 0xffff,
                "sw" => value,
                _ => unreachable!(),
            };

            assert_eq!(
                stored, expected,
                "{name} failed: value={value:#010x},
                 address={address}, \
                 expected={expected:#010x},
                 got={stored:#010x}",
            );
        }
    }

    // B-type ops test
    #[test]
    fn test_ops_branch() {
        let test_vectors: Vec<(&str, u32, u32, i32, bool)> = vec![
            // (name, rs1 value, rs2 value, offset, taken)
            ("beq", 10, 10, 8, true),
            ("beq", 10, 20, 8, false),
            ("bne", 10, 20, 8, true),
            ("bne", 10, 10, 8, false),
            ("blt", 10, 20, 8, true),
            ("blt", 20, 10, 8, false),
            ("bge", 20, 10, 8, true),
            ("bge", 10, 20, 8, false),
            ("bltu", 10, 20, 8, true),
            ("bltu", 20, 10, 8, false),
            ("bgeu", 20, 10, 8, true),
            ("bgeu", 10, 20, 8, false),
            ("beq", 10, 10, -4, true), // negative branch offsets
            ("bne", 10, 20, -4, true),
            ("blt", 0xffff_ffff, 0, 8, true), // signed vs unsigned comparison
            ("bltu", 0xffff_ffff, 0, 8, false),
            ("bge", 0xffff_ffff, 0, 8, false),
            ("bgeu", 0xffff_ffff, 0, 8, true),
        ];

        for (name, rs1_value, rs2_value, offset, taken) in test_vectors {
            let rs1 = 1;
            let rs2 = 2;

            let instruction = encode_b(name, rs1, rs2, offset);
            let mut cpu = RV32::new(instruction.to_le_bytes().to_vec());

            cpu.regs[rs1 as usize] = rs1_value;
            cpu.regs[rs2 as usize] = rs2_value;

            let initial_pc = cpu.pc;

            cpu.step().unwrap();

            let expected_pc = if taken {
                initial_pc.wrapping_add(offset as u32)
            } else {
                initial_pc.wrapping_add(4)
            };

            assert_eq!(
                cpu.pc, expected_pc,
                "{name} failed: rs1={rs1_value:#010x}, \
                 rs2={rs2_value:#010x}, offset={offset}, \
                 expected_pc={expected_pc:#010x}, \
                 got={:#010x}",
                cpu.pc,
            );
        }
    }
}