path: root/src/mem.rs

// Copyright (c) 2025-2026 taitep
// SPDX-License-Identifier: BSD-2-Clause
//
// This file is part of TRVE (https://gitea.taitep.se/taitep/trve)
// See LICENSE file in the project root for full license text.

use std::sync::{
    Arc,
    atomic::{
        AtomicU8, AtomicU16, AtomicU32, AtomicU64,
        Ordering::{self, Relaxed},
    },
};

use memmap2::MmapMut;

use crate::exceptions::{MemoryException, MemoryExceptionType};

pub type PageNum = usize;

pub const RAM_START: u64 = 0x8000_0000;

#[derive(Clone)]
pub struct MemConfig {
    pub ram: Arc<Ram>,
    pub mmio_root: MmioRoot,
}

impl MemConfig {
    pub fn memory_mapping_type(&self, addr: u64) -> Option<MemoryMappingType> {
        if addr >= RAM_START {
            Some(MemoryMappingType::RAM)
        } else {
            self.mmio_root
                .get_device(addr)
                .map(|_| MemoryMappingType::MMIO)
        }
    }

    pub fn read_dword(&self, addr: u64) -> Result<u64, MemoryException> {
        if addr >= RAM_START {
            self.ram.read_dword(addr - RAM_START)
        } else {
            if !addr.is_multiple_of(8) && self.mmio_root.crosses_boundary(addr, 8) {
                return Err(MemoryException {
                    type_: MemoryExceptionType::AddressMisaligned,
                    addr,
                });
            }
            let (interface, addr) = self.mmio_root.get_device(addr).ok_or(MemoryException {
                type_: MemoryExceptionType::AccessFault,
                addr,
            })?;

            interface.read_dword(addr).map_err(|e| e.with_addr(addr))
        }
    }
    pub fn read_word(&self, addr: u64) -> Result<u32, MemoryException> {
        if addr >= RAM_START {
            self.ram.read_word(addr - RAM_START)
        } else {
            if !addr.is_multiple_of(4) && self.mmio_root.crosses_boundary(addr, 4) {
                return Err(MemoryException {
                    type_: MemoryExceptionType::AddressMisaligned,
                    addr,
                });
            }
            let (interface, addr) = self.mmio_root.get_device(addr).ok_or(MemoryException {
                type_: MemoryExceptionType::AccessFault,
                addr,
            })?;

            interface.read_word(addr).map_err(|e| e.with_addr(addr))
        }
    }
    pub fn read_hword(&self, addr: u64) -> Result<u16, MemoryException> {
        if addr >= RAM_START {
            self.ram.read_hword(addr - RAM_START)
        } else {
            if !addr.is_multiple_of(2) && self.mmio_root.crosses_boundary(addr, 2) {
                return Err(MemoryException {
                    type_: MemoryExceptionType::AddressMisaligned,
                    addr,
                });
            }
            let (interface, addr) = self.mmio_root.get_device(addr).ok_or(MemoryException {
                type_: MemoryExceptionType::AccessFault,
                addr,
            })?;
            interface.read_hword(addr).map_err(|e| e.with_addr(addr))
        }
    }
    pub fn read_byte(&self, addr: u64) -> Result<u8, MemoryException> {
        if addr >= RAM_START {
            self.ram.read_byte(addr - RAM_START)
        } else {
            let (interface, addr) = self.mmio_root.get_device(addr).ok_or(MemoryException {
                type_: MemoryExceptionType::AccessFault,
                addr,
            })?;
            interface.read_byte(addr).map_err(|e| e.with_addr(addr))
        }
    }

    pub fn write_dword(&self, addr: u64, value: u64) -> Result<(), MemoryException> {
        if addr >= RAM_START {
            self.ram.write_dword(addr - RAM_START, value)
        } else {
            if !addr.is_multiple_of(8) && self.mmio_root.crosses_boundary(addr, 8) {
                return Err(MemoryException {
                    type_: MemoryExceptionType::AddressMisaligned,
                    addr,
                });
            }
            let (interface, addr) = self.mmio_root.get_device(addr).ok_or(MemoryException {
                type_: MemoryExceptionType::AccessFault,
                addr,
            })?;
            interface
                .write_dword(addr, value)
                .map_err(|e| e.with_addr(addr))
        }
    }
    pub fn write_word(&self, addr: u64, value: u32) -> Result<(), MemoryException> {
        if addr >= RAM_START {
            self.ram.write_word(addr - RAM_START, value)
        } else {
            if !addr.is_multiple_of(4) && self.mmio_root.crosses_boundary(addr, 4) {
                return Err(MemoryException {
                    type_: MemoryExceptionType::AddressMisaligned,
                    addr,
                });
            }
            let (interface, addr) = self.mmio_root.get_device(addr).ok_or(MemoryException {
                type_: MemoryExceptionType::AccessFault,
                addr,
            })?;
            interface
                .write_word(addr, value)
                .map_err(|e| e.with_addr(addr))
        }
    }
    pub fn write_hword(&self, addr: u64, value: u16) -> Result<(), MemoryException> {
        if addr >= RAM_START {
            self.ram.write_hword(addr - RAM_START, value)
        } else {
            if !addr.is_multiple_of(2) && self.mmio_root.crosses_boundary(addr, 2) {
                return Err(MemoryException {
                    type_: MemoryExceptionType::AddressMisaligned,
                    addr,
                });
            }
            let (interface, addr) = self.mmio_root.get_device(addr).ok_or(MemoryException {
                type_: MemoryExceptionType::AccessFault,
                addr,
            })?;
            interface
                .write_hword(addr, value)
                .map_err(|e| e.with_addr(addr))
        }
    }
    pub fn write_byte(&self, addr: u64, value: u8) -> Result<(), MemoryException> {
        if addr >= RAM_START {
            self.ram.write_byte(addr - RAM_START, value)
        } else {
            let (interface, addr) = self.mmio_root.get_device(addr).ok_or(MemoryException {
                type_: MemoryExceptionType::AccessFault,
                addr,
            })?;
            interface
                .write_byte(addr, value)
                .map_err(|e| e.with_addr(addr))
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MemoryMappingType {
    MMIO,
    RAM,
}

pub struct Ram {
    buf: MmapMut,
    version_counters: Arc<[AtomicU32]>,
}

#[cfg(target_endian = "big")]
compile_error!("Current RAM implementation requires a little-endian host.");

impl Ram {
    pub fn try_new(size: usize) -> Result<Self, std::io::Error> {
        if !size.is_multiple_of(8) {
            return Err(std::io::Error::other("ram size must be a multiple of 8"));
        }
        Ok(Self {
            buf: MmapMut::map_anon(size)?,
            // SAFETY: We do not care about the initial version counts. Wrapping is fine. Only
            // equality is ever checked for, not magnitude.
            version_counters: unsafe {
                Arc::new_uninit_slice(size.div_ceil(Self::VERSION_CHUNK_SIZE)).assume_init()
            },
        })
    }

    const VERSION_CHUNK_SIZE: usize = 64;

    pub fn buf_mut(&mut self) -> &mut [u8] {
        self.buf.as_mut()
    }

    /// # Safety
    /// Safe if the size of the memory in bytes is divisible by the size of T
    /// Assuming try_new is used, RAM size is guaranteed to be a multiple of 8
    /// meaning anything with size 1, 2, 4, or 8 bytes is valid.
    /// It must also be known that the contents of RAM are made up of naturally
    /// aligned valid instances of T.
    #[inline]
    unsafe fn buf_transmuted<T>(&self) -> &[T] {
        debug_assert!(self.buf.len().is_multiple_of(std::mem::size_of::<T>()));
        unsafe {
            std::slice::from_raw_parts(
                self.buf.as_ptr() as *const T,
                self.buf.len() / std::mem::size_of::<T>(),
            )
        }
    }

    #[inline]
    pub fn buf_atomic(&self) -> &[AtomicU8] {
        unsafe { std::slice::from_raw_parts(self.buf.as_ptr() as *const AtomicU8, self.buf.len()) }
    }

    #[inline]
    pub fn read_dword(&self, addr: u64) -> Result<u64, MemoryException> {
        if !addr.is_multiple_of(8) {
            let high_word_addr = addr.wrapping_add(4);

            let low_word = self.read_word(addr)?;
            let high_word = self.read_word(high_word_addr)?;

            return Ok((low_word as u64) | (high_word as u64) << 32);
        }

        let index = (addr / 8) as usize;
        Ok(unsafe {
            self.buf_transmuted::<AtomicU64>()
                .get(index)
                .ok_or(MemoryException {
                    type_: MemoryExceptionType::AccessFault,
                    addr,
                })
        }?
        .load(Relaxed))
    }
    #[inline]
    pub fn read_word(&self, addr: u64) -> Result<u32, MemoryException> {
        if !addr.is_multiple_of(4) {
            let high_hword_addr = addr.wrapping_add(2);

            let low_hword = self.read_hword(addr)?;
            let high_hword = self.read_hword(high_hword_addr)?;

            return Ok((low_hword as u32) | (high_hword as u32) << 16);
        }

        let index = (addr / 4) as usize;
        Ok(unsafe {
            self.buf_transmuted::<AtomicU32>()
                .get(index)
                .ok_or(MemoryException {
                    type_: MemoryExceptionType::AccessFault,
                    addr,
                })
        }?
        .load(Relaxed))
    }
    #[inline]
    pub fn read_hword(&self, addr: u64) -> Result<u16, MemoryException> {
        if !addr.is_multiple_of(2) {
            let high_byte_addr = addr.wrapping_add(1);

            let low_byte = self.read_byte(addr)?;
            let high_byte = self.read_byte(high_byte_addr)?;

            return Ok((low_byte as u16) | (high_byte as u16) << 8);
        }

        let index = (addr / 2) as usize;
        Ok(unsafe {
            self.buf_transmuted::<AtomicU16>()
                .get(index)
                .ok_or(MemoryException {
                    type_: MemoryExceptionType::AccessFault,
                    addr,
                })
        }?
        .load(Relaxed))
    }
    #[inline]
    pub fn read_byte(&self, addr: u64) -> Result<u8, MemoryException> {
        Ok(self
            .buf_atomic()
            .get(addr as usize)
            .ok_or(MemoryException {
                type_: MemoryExceptionType::AccessFault,
                addr,
            })?
            .load(Relaxed))
    }

    #[inline]
    pub fn write_dword(&self, addr: u64, value: u64) -> Result<(), MemoryException> {
        if !addr.is_multiple_of(8) {
            let low_word = value as u32;
            let high_word = (value >> 32) as u32;

            let high_word_address = addr.wrapping_add(4);

            self.write_word(addr, low_word)?;
            self.write_word(high_word_address, high_word)?;
            return Ok(());
        }

        self.claim_addr_even(addr)
            .ok_or_else(|| MemoryExceptionType::AccessFault.with_addr(addr))?;

        let index = (addr / 8) as usize;
        unsafe {
            self.buf_transmuted::<AtomicU64>()
                .get(index)
                .ok_or(MemoryException {
                    type_: MemoryExceptionType::AccessFault,
                    addr,
                })
        }?
        .store(value, Relaxed);
        Ok(())
    }
    #[inline]
    pub fn write_word(&self, addr: u64, value: u32) -> Result<(), MemoryException> {
        if !addr.is_multiple_of(4) {
            let low_hword = value as u16;
            let high_hword = (value >> 16) as u16;

            let high_hword_address = addr.wrapping_add(2);

            self.write_hword(addr, low_hword)?;
            self.write_hword(high_hword_address, high_hword)?;
            return Ok(());
        }

        self.claim_addr_even(addr)
            .ok_or_else(|| MemoryExceptionType::AccessFault.with_addr(addr))?;

        let index = (addr / 4) as usize;
        unsafe {
            self.buf_transmuted::<AtomicU32>()
                .get(index)
                .ok_or(MemoryException {
                    type_: MemoryExceptionType::AccessFault,
                    addr,
                })
        }?
        .store(value, Relaxed);
        Ok(())
    }
    #[inline]
    pub fn write_hword(&self, addr: u64, value: u16) -> Result<(), MemoryException> {
        if !addr.is_multiple_of(2) {
            let low_byte = value as u8;
            let high_byte = (value >> 8) as u8;

            let high_byte_address = addr.wrapping_add(1);

            self.write_byte(addr, low_byte)?;
            self.write_byte(high_byte_address, high_byte)?;
            return Ok(());
        }

        self.claim_addr_even(addr)
            .ok_or_else(|| MemoryExceptionType::AccessFault.with_addr(addr))?;

        let index = (addr / 2) as usize;
        unsafe {
            self.buf_transmuted::<AtomicU16>()
                .get(index)
                .ok_or(MemoryException {
                    type_: MemoryExceptionType::AccessFault,
                    addr,
                })
        }?
        .store(value, Relaxed);
        Ok(())
    }
    #[inline]
    pub fn write_byte(&self, addr: u64, value: u8) -> Result<(), MemoryException> {
        self.claim_addr_even(addr)
            .ok_or_else(|| MemoryExceptionType::AccessFault.with_addr(addr))?;
        self.buf_atomic()
            .get(addr as usize)
            .ok_or(MemoryException {
                type_: MemoryExceptionType::AccessFault,
                addr,
            })?
            .store(value, Relaxed);
        Ok(())
    }

    pub fn claim_addr_even<'a>(&'a self, addr: u64) -> Option<RamVersionClaim<'a>> {
        let chunk_id = addr as usize / Self::VERSION_CHUNK_SIZE;
        let chunk_counter = self.version_counters.get(chunk_id)?;
        Some(RamVersionClaim::claim_even(&chunk_counter))
    }
}

pub struct RamVersionClaim<'a> {
    version_counter: &'a AtomicU32,
    intial_version: u32,
}

impl<'a> RamVersionClaim<'a> {
    pub fn claim_even(counter: &'a AtomicU32) -> RamVersionClaim<'a> {
        loop {
            let current_version = counter.load(Ordering::Acquire);
            if !current_version.is_multiple_of(2) {
                continue;
            }

            // Attempt to increment and therefore successfully claim the version
            let res = counter.compare_exchange(
                current_version,
                current_version.wrapping_add(1),
                Ordering::AcqRel,
                Ordering::Acquire,
            );

            if let Ok(initial_version) = res {
                return RamVersionClaim {
                    version_counter: counter,
                    intial_version: initial_version,
                };
            }
        }
    }
}

impl<'a> Drop for RamVersionClaim<'a> {
    fn drop(&mut self) {
        self.version_counter
            .store(self.intial_version.wrapping_add(2), Ordering::Release);
    }
}

pub const MMIO_SECOND_LEVEL_PAGE_SIZE: usize = 64 * 1024;
pub const MMIO_ROOT_PAGE_SIZE: usize = MMIO_SECOND_LEVEL_PAGE_SIZE * 64;

const MMIO_ROOT_ENTRIES: usize = RAM_START as usize / MMIO_ROOT_PAGE_SIZE;
const MMIO_SECOND_LEVEL_ENTRIES: usize = MMIO_ROOT_PAGE_SIZE / MMIO_SECOND_LEVEL_PAGE_SIZE;

#[derive(Clone)]
pub struct MmioRoot(Box<[Option<MmioSecondLevel>; MMIO_ROOT_ENTRIES]>);

impl MmioRoot {
    pub fn insert(&mut self, base_addr: u64, interface: Arc<dyn MemDeviceInterface>) {
        assert!(base_addr.is_multiple_of(MMIO_SECOND_LEVEL_PAGE_SIZE as u64));
        assert!(base_addr < RAM_START);

        let page_id = base_addr as usize / MMIO_SECOND_LEVEL_PAGE_SIZE;
        let root_page_id = page_id / MMIO_SECOND_LEVEL_ENTRIES;
        let second_level_page_id = page_id % MMIO_SECOND_LEVEL_ENTRIES;

        let second_level = self.0[root_page_id].get_or_insert_default();

        if let MmioSecondLevel::SubTable(t) = second_level {
            t[second_level_page_id] = Some(interface);
        }
    }

    pub fn insert_full(&mut self, base_addr: u64, interface: Arc<dyn MemDeviceInterface>) {
        assert!(base_addr.is_multiple_of(MMIO_ROOT_PAGE_SIZE as u64));
        assert!(base_addr < RAM_START);

        let page_id = base_addr as usize / MMIO_ROOT_PAGE_SIZE;

        self.0[page_id] = Some(MmioSecondLevel::Interface(interface));
    }

    fn get_device(&self, addr: u64) -> Option<(Arc<dyn MemDeviceInterface>, u64)> {
        debug_assert!(addr < RAM_START);

        let page_id = addr as usize / MMIO_SECOND_LEVEL_PAGE_SIZE;
        let root_page_id = page_id / MMIO_SECOND_LEVEL_ENTRIES;

        self.0[root_page_id]
            .as_ref()
            .and_then(|s| s.get_device(addr % MMIO_ROOT_PAGE_SIZE as u64))
    }

    fn crosses_boundary(&self, addr: u64, size: u64) -> bool {
        if addr >= RAM_START {
            return false;
        }

        if addr + size > RAM_START {
            return true;
        }

        let page_id = addr as usize / MMIO_SECOND_LEVEL_PAGE_SIZE;
        let root_page_id = page_id / MMIO_SECOND_LEVEL_ENTRIES;
        let end = addr + size - 1;

        match self.0[root_page_id].as_ref() {
            Some(s) => match s {
                MmioSecondLevel::SubTable(_) => {
                    let end_page_id = end as usize / MMIO_SECOND_LEVEL_PAGE_SIZE;
                    page_id != end_page_id
                }
                MmioSecondLevel::Interface(_) => {
                    let end_root_page_id = end as usize / MMIO_ROOT_PAGE_SIZE;
                    root_page_id != end_root_page_id
                }
            },
            None => false,
        }
    }
}

impl Default for MmioRoot {
    fn default() -> Self {
        Self(Box::new([(); MMIO_ROOT_ENTRIES].map(|_| None)))
    }
}

#[derive(Clone)]
enum MmioSecondLevel {
    SubTable(Box<[Option<Arc<dyn MemDeviceInterface>>; MMIO_SECOND_LEVEL_ENTRIES]>),
    Interface(Arc<dyn MemDeviceInterface>),
}

impl MmioSecondLevel {
    fn get_device(&self, addr: u64) -> Option<(Arc<dyn MemDeviceInterface>, u64)> {
        let page_id = addr as usize / MMIO_SECOND_LEVEL_PAGE_SIZE;
        match self {
            Self::SubTable(t) => t[page_id]
                .as_ref()
                .map(|i| (i.clone(), addr % MMIO_SECOND_LEVEL_PAGE_SIZE as u64)),

            Self::Interface(i) => Some((i.clone(), addr)),
        }
    }
}

impl Default for MmioSecondLevel {
    fn default() -> Self {
        Self::SubTable(Box::new([(); MMIO_SECOND_LEVEL_ENTRIES].map(|_| None)))
    }
}

#[allow(unused_variables)]
pub trait MemDeviceInterface {
    fn write_dword(&self, addr: u64, value: u64) -> Result<(), MemoryExceptionType> {
        Err(MemoryExceptionType::AccessFault)
    }
    fn write_word(&self, addr: u64, value: u32) -> Result<(), MemoryExceptionType> {
        Err(MemoryExceptionType::AccessFault)
    }
    fn write_hword(&self, addr: u64, value: u16) -> Result<(), MemoryExceptionType> {
        Err(MemoryExceptionType::AccessFault)
    }
    fn write_byte(&self, addr: u64, value: u8) -> Result<(), MemoryExceptionType> {
        Err(MemoryExceptionType::AccessFault)
    }

    fn read_dword(&self, addr: u64) -> Result<u64, MemoryExceptionType> {
        Err(MemoryExceptionType::AccessFault)
    }
    fn read_word(&self, addr: u64) -> Result<u32, MemoryExceptionType> {
        Err(MemoryExceptionType::AccessFault)
    }
    fn read_hword(&self, addr: u64) -> Result<u16, MemoryExceptionType> {
        Err(MemoryExceptionType::AccessFault)
    }
    fn read_byte(&self, addr: u64) -> Result<u8, MemoryExceptionType> {
        Err(MemoryExceptionType::AccessFault)
    }
}