//! 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 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  │
        // └────────────────────────┴────────────┴────────┴────────┴─────────┘
        let imm = sign_extend(instruction >> 20, 12);

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

        // 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: {}", instruction)),
                };
                let value = self.read(a.wrapping_add(imm), size)?;
                result = Some(if signed {
                    sign_extend(value, size as u32 * 8)
                } else {
                    value
                });
            }

            // 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 (A): {}", instruction)),
                });
            }
            _ => return Err(anyhow!("illegal instruction (B): {}", 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 = match name {
            "lb" => 0b000,
            "lh" => 0b001,
            "lw" => 0b010,
            "lbu" => 0b100,
            "lhu" => 0b101,
            _ => panic!("unknown load instruction: {name}"),
        };

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

        (imm << 20) | (rs1 << 15) | (funct3 << 12) | (rd << 7) | OPCODE_LOAD
    }

    #[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],
            );
        }
    }

    #[test]
    fn test_ops_load() {
        let test_vectors: [(&str, u32, i32, u32); 8] = [
            ("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],
            );
        }
    }
}