DSA/dsa/emulator/src/processor/processor.rs

use std::{
    hint::unlikely,
    ops::{Add, AddAssign},
    sync::{
        Arc,
        atomic::Ordering,
        mpsc::{self, TryRecvError},
    },
    thread,
    time::{Duration, Instant},
};

use common::{
    instructions::{Instruction, Opcode},
    register::Register,
};

use crate::{
    Page,
    config::{IoMapping, MemoryMap, RegionType},
    io::{IoAccess, IoDevice, MappedDevice},
    memory::{mmu::MMU, ram::RandomAccessMemory},
    processor::{interrupts::Interrupt, state::SharedState},
};

pub struct Emulator<Mem: RandomAccessMemory> {
    internal_state: ProcessorSnapshot,
    shared_state: Arc<SharedState>,

    // Interrupts
    interrupts: mpsc::Receiver<u8>,
    pending_fault: Option<Interrupt>,

    // memory
    mmu: MMU,
    mainstore: Mem,

    // IO
    mmio: Vec<MappedDevice>,
    mmio_region: Option<(u32, u32)>, // start end end of segment

    // optionals - not necessarily set on boot.
    memory_map: Option<MemoryMap>,

    // Config params
    __mmap_configured: bool,
    __io_configured: bool,
}

unsafe impl<Mem: RandomAccessMemory> Send for Emulator<Mem> {}
impl<Mem: RandomAccessMemory> Emulator<Mem> {
    pub fn new(mem: Mem) -> Self {
        let (sender, receiver) = mpsc::channel::<u8>();

        Self {
            interrupts: receiver,
            pending_fault: None,

            internal_state: ProcessorSnapshot::default(),
            shared_state: Arc::new(SharedState::new(sender)),

            memory_map: None,
            mmu: MMU::new(),
            mainstore: mem,

            // mmio
            mmio: Vec::new(),
            mmio_region: None,

            // Config params
            __io_configured: false,
            __mmap_configured: false,
        }
    }

    pub fn with_device(mut self, device: impl IoDevice + 'static) -> Self {
        self.mmio.push(MappedDevice {
            base: 0,
            device: Box::new(device),
        });
        self
    }

    pub fn state_handle(&self) -> Arc<SharedState> {
        self.shared_state.clone()
    }

    #[inline]
    pub fn mmu_mut(&mut self) -> &mut MMU {
        &mut self.mmu
    }

    #[inline]
    pub fn memory_mut(&mut self) -> &mut impl RandomAccessMemory {
        &mut self.mainstore
    }

    #[must_use]
    #[inline]
    pub fn reg(&self, reg: Register) -> u32 {
        if reg as u8 == Register::Zero as u8 {
            return 0;
        }
        debug_assert!((reg as usize) < ProcessorSnapshot::REG_COUNT);

        unsafe { *self.internal_state.registers.get_unchecked(reg as usize) }
    }

    #[inline]
    pub fn mut_reg(&mut self, reg: Register) -> &mut u32 {
        debug_assert!((reg as usize) < ProcessorSnapshot::REG_COUNT);

        unsafe {
            self.internal_state
                .registers
                .get_unchecked_mut(reg as usize)
        }
    }

    #[cold]
    pub fn apply_memory_map(mut self, map: MemoryMap) -> Self {
        self.mmio_region = map
            .regions
            .iter()
            .find(|r| matches!(r.region_type, RegionType::MMIO))
            .map(|r| (r.base, r.base + r.size));

        map.apply(&mut self);
        self.memory_map = Some(map);
        self.__mmap_configured = true;
        self
    }

    pub fn apply_io_map(mut self, map: Vec<IoMapping>) -> Result<Self, String> {
        if !self.__mmap_configured {
            return Err("You must map memory before applying I/O mappings".to_string());
        }

        for entry in map {
            if let Some(idx) = self.mmio.iter().position(|d| d.device.id() == entry.device) {
                self.mmio[idx].base = entry.base;
            } else {
                eprintln!("WARN: no device registgered for {:?}", entry.device);
            }
        }

        self.__io_configured = true;
        Ok(self)
    }

    #[cold]
    fn update(&mut self) {
        self.shared_state
            .proc
            .store(Arc::new(self.internal_state.clone()));
    }

    #[cold]
    fn boot(&mut self) -> Result<(), String> {
        if !(self.__io_configured && self.__mmap_configured) {
            return Err("Processor not configured".to_string());
        }

        if let Some(map) = &self.memory_map {
            MemoryMap::identity_map(&mut self.mmu, map.regions.get(0).unwrap());
        }

        self.internal_state.running = true;
        Ok(())
    }

    #[cold]
    fn shutdown(&mut self) {
        self.internal_state.running = false;
    }

    #[cold]
    pub fn idle_wait(&mut self) {
        self.internal_state.running = false;

        // Wait for an interrupt or state update to continue.
        loop {
            // Check for interrupts.
            if let Ok(code) = self.interrupts.recv_timeout(Duration::from_millis(100)) {
                self.interrupt(Interrupt::Software(code));
                break;
            }

            // // If we've received a request to continue running. DEPRECATED (we can run through interrupts)
            // if self.shared_state.running.load(Ordering::Relaxed) {
            //     panic!("3");
            //     break;
            // }

            // UI requested a state update.
            if self.shared_state.update_req.load(Ordering::Relaxed) {
                self.update();
            }
        }

        self.internal_state.running = true;
    }

    pub fn run(&mut self) -> Result<(), String> {
        self.boot()?;

        let mut time = Instant::now();

        'emu: loop {
            // Update UI thread (roughly every 512k cycles targeting 60 UPS)
            if unlikely(self.internal_state.clock & 0x7FFFF == 0) {
                if self.shared_state.update_req.load(Ordering::Relaxed) {
                    self.update();
                }
            }

            // Check for commands or hardware Interrupts every 32k cycles
            if self.internal_state.clock % 0x7FFF == 0 {
                if self.shared_state.running.load(Ordering::Relaxed) == false {
                    self.idle_wait();
                }

                match self.interrupts.try_recv() {
                    Ok(code) => self.interrupt(Interrupt::Software(code)),
                    Err(TryRecvError::Disconnected) => break 'emu,
                    Err(TryRecvError::Empty) => {}
                }
            }

            // if we got a halt instruction, wait
            if unlikely(!self.internal_state.running) {
                let mips = (self.internal_state.clock as u128 / time.elapsed().as_micros());
                println!(
                    "TIME TAKEN: {:?}, clock: {}, {}MIPS",
                    time.elapsed(),
                    self.internal_state.clock,
                    mips
                );
                time = Instant::now();

                // temporary while i figure out a better solution
                // exit when we hit halt (not useful for a real os ofc)
                break 'emu;

                self.idle_wait();
            }

            let pc = self.reg(Register::Pcx);
            let instruction = self.mem_read_word(pc);

            if cfg!(debug_assertions) {
                println!(
                    "Clock: {} Executing {:?} PCX: {}, reg: {:?}",
                    self.internal_state.clock,
                    Instruction(instruction),
                    pc,
                    self.internal_state.registers
                );
                thread::sleep(Duration::from_micros(10));
            }

            self.mut_reg(Register::Pcx).add_assign(4);
            self.execute(Instruction(instruction));

            // Check if executing the interrupt caused a fault.
            if let Some(fault) = self.pending_fault {
                self.pending_fault = None;
                println!("WARN fault: {:?}", fault);
                self.interrupt(fault);
            }

            // always increment clock.
            self.internal_state.clock += 1;
        }

        self.shutdown();

        Ok(())
    }

    #[inline]
    fn interrupt(&mut self, int: Interrupt) {
        let idt = self.reg(Register::Idr);

        *self.mut_reg(Register::Spr) -= 4;
        let spr = self.reg(Register::Spr);
        let pcx = self.reg(Register::Pcx);
        self.mem_write_word(spr, pcx);

        *self.mut_reg(Register::Pcx) = self.mem_read_word(idt + int.code() as u32 * 4)
    }

    #[inline]
    fn execute(&mut self, word: Instruction) {
        // This needs to be unsafe as we're using word.xxxx_uc() functions for decoding which are unsafe
        // as they perform unchecked transmute operations (performance critical)
        unsafe {
            match Opcode::from_u8(word.opcode()) {
                // Nothing
                Some(Opcode::Nop) => {}

                // move
                Some(Opcode::Mov) => {
                    *self.mut_reg(word.dest_uc()) = self.reg(word.src1_uc());
                }
                Some(Opcode::CMov) => {
                    if self.reg(word.misc_uc()) != 0 {
                        *self.mut_reg(word.dest_uc()) = self.reg(word.src1_uc());
                    }
                }

                // load
                Some(Opcode::Ldbs) => todo!(),
                Some(Opcode::Ldhs) => todo!(),
                Some(Opcode::Ldb) => {
                    *self.mut_reg(word.dest_uc()) = u32::from(
                        self.mem_read_byte(self.reg(word.src1_uc()) + word.imm16() as u32),
                    )
                }
                Some(Opcode::Ldh) => {
                    *self.mut_reg(word.dest_uc()) = u32::from(
                        self.mem_read_word(self.reg(word.src1_uc()) + word.imm16() as u32) >> 16,
                    )
                }
                Some(Opcode::Ldw) => {
                    *self.mut_reg(word.dest_uc()) = u32::from(
                        self.mem_read_word(self.reg(word.src1_uc()) + word.imm16() as u32),
                    )
                }

                // store
                Some(Opcode::Stb) => {
                    self.mem_write_byte(
                        self.reg(word.dest_uc()) + word.imm16() as u32,
                        self.reg(word.src1_uc()) as u8,
                    );
                }
                Some(Opcode::Sth) => {
                    self.mem_write_byte(
                        self.reg(word.dest_uc()) + word.imm16() as u32,
                        (self.reg(word.src1_uc()) as u16 >> 8) as u8,
                    );
                    self.mem_write_byte(
                        self.reg(word.dest_uc()) + word.imm16() as u32 + 1,
                        self.reg(word.src1_uc()) as u8,
                    );
                }
                Some(Opcode::Stw) => {
                    self.mem_write_word(
                        self.reg(word.dest_uc()) + word.imm16() as u32,
                        self.reg(word.src1_uc()),
                    );
                }

                // load immediate
                Some(Opcode::Lli) => {
                    *self.mut_reg(word.dest_uc()) = word.imm16() as u32;
                }
                Some(Opcode::Lui) => {
                    *self.mut_reg(word.dest_uc()) =
                        (word.imm16() as u32) << 16 | self.reg(word.dest_uc());
                }

                // Comparison
                Some(Opcode::Ieq) => {
                    *self.mut_reg(word.dest_uc()) =
                        (self.reg(word.src1_uc()) == self.reg(word.src2_uc())) as u32;
                }
                Some(Opcode::Ine) => {
                    *self.mut_reg(word.dest_uc()) =
                        (self.reg(word.src1_uc()) != self.reg(word.src2_uc())) as u32;
                }
                Some(Opcode::Ilt) => {
                    *self.mut_reg(word.dest_uc()) =
                        (self.reg(word.src1_uc()) < self.reg(word.src2_uc())) as u32;
                }
                Some(Opcode::Ile) => {
                    *self.mut_reg(word.dest_uc()) =
                        (self.reg(word.src1_uc()) <= self.reg(word.src2_uc())) as u32;
                }
                Some(Opcode::Igt) => {
                    *self.mut_reg(word.dest_uc()) =
                        (self.reg(word.src1_uc()) > self.reg(word.src2_uc())) as u32;
                }
                Some(Opcode::Ige) => {
                    *self.mut_reg(word.dest_uc()) =
                        (self.reg(word.src1_uc()) >= self.reg(word.src2_uc())) as u32;
                }

                // Jump
                Some(Opcode::Jmp) => {
                    *self.mut_reg(Register::Pcx) = self.reg(word.dest_uc()) + word.imm16() as u32
                }

                Some(Opcode::Jez) => {
                    if self.reg(word.src1_uc()) == 0 {
                        *self.mut_reg(Register::Pcx) =
                            self.reg(word.dest_uc()) + word.imm16() as u32
                    }
                }
                Some(Opcode::Jnz) => {
                    if self.reg(word.src1_uc()) != 0 {
                        *self.mut_reg(Register::Pcx) =
                            self.reg(word.dest_uc()) + word.imm16() as u32
                    }
                }
                Some(Opcode::Jnc) => todo!(),
                Some(Opcode::Jic) => todo!(),

                // Bitwise
                Some(Opcode::And) => {
                    *self.mut_reg(word.dest_uc()) =
                        self.reg(word.src1_uc()) & self.reg(word.src2_uc());
                }
                Some(Opcode::Nand) => {
                    *self.mut_reg(word.dest_uc()) =
                        !(self.reg(word.src1_uc()) & self.reg(word.src2_uc()));
                }
                Some(Opcode::Or) => {
                    *self.mut_reg(word.dest_uc()) =
                        self.reg(word.src1_uc()) | self.reg(word.src2_uc());
                }
                Some(Opcode::Nor) => {
                    *self.mut_reg(word.dest_uc()) =
                        !(self.reg(word.src1_uc()) | self.reg(word.src2_uc()));
                }
                Some(Opcode::Xor) => {
                    *self.mut_reg(word.dest_uc()) =
                        self.reg(word.src1_uc()) ^ self.reg(word.src2_uc());
                }
                Some(Opcode::Xnor) => {
                    *self.mut_reg(word.dest_uc()) =
                        !(self.reg(word.src1_uc()) ^ self.reg(word.src2_uc()));
                }
                Some(Opcode::Not) => {
                    *self.mut_reg(word.dest_uc()) = !self.reg(word.src1_uc());
                }

                // Arithmetic
                Some(Opcode::Add) => {
                    *self.mut_reg(word.dest_uc()) =
                        self.reg(word.src1_uc()) + self.reg(word.src2_uc());
                }

                Some(Opcode::Sub) => {
                    *self.mut_reg(word.dest_uc()) =
                        self.reg(word.src1_uc()) - self.reg(word.src2_uc());
                }

                Some(Opcode::Shl) => {
                    *self.mut_reg(word.dest_uc()) = self.reg(word.src1_uc())
                        << (self.reg(word.src2_uc()) + word.shamt() as u32);
                }

                Some(Opcode::Shr) => {
                    *self.mut_reg(word.dest_uc()) = self.reg(word.src1_uc())
                        >> (self.reg(word.src2_uc()) + word.shamt() as u32);
                }

                Some(Opcode::Addi) => {
                    *self.mut_reg(word.dest_uc()) = self.reg(word.src1_uc()) + word.imm16() as u32;
                }

                Some(Opcode::Subi) => {
                    *self.mut_reg(word.dest_uc()) = self.reg(word.src1_uc()) - word.imm16() as u32;
                }

                // Utility
                Some(Opcode::Push) => {
                    self.push(word.src1_uc());
                }
                Some(Opcode::Pop) => {
                    self.pop(word.dest_uc());
                }

                // Function
                Some(Opcode::Call) => {
                    self.push(Register::Pcx);
                    *self.mut_reg(Register::Pcx) = self.reg(word.dest_uc()) + word.imm16() as u32
                }
                Some(Opcode::Ret) => {
                    self.pop(Register::Ret);
                    *self.mut_reg(Register::Pcx) = self.reg(Register::Ret);
                }

                Some(Opcode::Int) => {
                    self.interrupt(Interrupt::Software(word.imm16() as u8));
                }
                Some(Opcode::Hlt) => self.internal_state.running = false,
                Some(Opcode::IRet) => todo!(),
                None => {}
            }
        }
    }

    #[inline]
    fn push(&mut self, reg: Register) {
        *self.mut_reg(Register::Spr) -= 4;
        self.mem_write_word(self.reg(Register::Spr), self.reg(reg));
    }

    #[inline]
    fn pop(&mut self, reg: Register) {
        *self.mut_reg(reg) = self.mem_read_word(self.reg(Register::Spr));
        *self.mut_reg(Register::Spr) += 4;
    }

    #[inline]
    fn mem_read_byte(&mut self, addr: u32) -> u8 {
        if unlikely(self.is_mmio(addr)) {
            return self.io_read_byte(addr);
        }
        self.mainstore.read_byte(addr)
    }

    #[inline]
    fn mem_write_byte(&mut self, addr: u32, val: u8) {
        if unlikely(self.is_mmio(addr)) {
            self.io_write_byte(addr, val);
            return;
        }
        self.mainstore.write_byte(addr, val);
    }

    #[inline]
    fn mem_read_word(&mut self, addr: u32) -> u32 {
        if unlikely(self.is_mmio(addr)) {
            return self.io_read_word(addr);
        }
        self.mainstore.read_word(addr)
    }

    #[inline]
    fn mem_write_word(&mut self, addr: u32, val: u32) {
        if unlikely(self.is_mmio(addr)) {
            self.io_write_word(addr, val);
            return;
        }
        self.mainstore.write_word(addr, val);
    }

    #[inline]
    fn mem_read_page(&mut self, addr: u32) -> &Page {
        if unlikely(self.is_mmio(addr)) {
            // pages spanning MMIO don't really make sense
            // treat as fault
            self.pending_fault = Some(Interrupt::ProtectionFault);
            return &Page::ZERO;
        }
        self.mainstore.read_page(addr)
    }

    #[inline]
    fn mem_write_page(&mut self, addr: u32, val: &Page) {
        if unlikely(self.is_mmio(addr)) {
            self.pending_fault = Some(Interrupt::ProtectionFault);
            return;
        }
        self.mainstore.write_page(addr, val);
    }

    // single MMIO check reused by all of the above
    #[inline]
    fn is_mmio(&self, addr: u32) -> bool {
        self.mmio_region
            .map(|(base, end)| addr >= base && addr < end)
            .unwrap_or(false)
    }

    fn fault(&mut self, interrupt: Interrupt) {
        self.pending_fault = Some(interrupt);
    }

    #[inline]
    fn get_device(&mut self, addr: u32) -> Option<(&mut Box<dyn IoDevice>, u32)> {
        // Find the device that covers the address.
        self.mmio
            .iter_mut()
            .find(|d| d.base <= addr && addr < d.base + d.device.size())
            .map(|d| (&mut d.device, addr - d.base))
    }
    // --- cold IO paths ---

    #[cold]
    fn io_read_byte(&mut self, addr: u32) -> u8 {
        if let Some((device, offset)) = self.get_device(addr) {
            device.read_byte(offset).unwrap_or_else(|fault| {
                self.fault(fault);
                0
            })
        } else {
            self.fault(Interrupt::UnmappedIo);
            0
        }
    }

    #[cold]
    fn io_read_word(&mut self, addr: u32) -> u32 {
        if let Some((device, offset)) = self.get_device(addr) {
            device.read_word(offset).unwrap_or_else(|fault| {
                self.fault(fault);
                0
            })
        } else {
            self.fault(Interrupt::UnmappedIo);

            0
        }
    }

    #[cold]
    fn io_write_byte(&mut self, addr: u32, val: u8) {
        if let Some((dev, offset)) = self.get_device(addr) {
            dev.write_byte(offset, val).unwrap_or_else(|fault| {
                self.fault(fault);
            });
        } else {
            self.fault(Interrupt::UnmappedIo);
        }
    }

    #[cold]
    fn io_write_word(&mut self, addr: u32, val: u32) {
        if let Some((dev, offset)) = self.get_device(addr) {
            dev.write_word(offset, val).unwrap_or_else(|fault| {
                self.fault(fault);
            });
        } else {
            self.fault(Interrupt::UnmappedIo);
        }
    }
}

#[derive(Debug, Clone)]
pub struct ProcessorSnapshot {
    pub running: bool,
    pub clock: usize,
    pub registers: [u32; Self::REG_COUNT],
}

impl Default for ProcessorSnapshot {
    fn default() -> Self {
        Self {
            running: false,
            clock: 0,
            registers: [0; Self::REG_COUNT],
        }
    }
}

impl ProcessorSnapshot {
    const REG_COUNT: usize = 28;
}