Generic memory controller and memory chip implementations in Rust

memory.rs

use num_traits::PrimInt;
use rangemap::RangeInclusiveMap;
use std::collections::HashMap;
use std::ops::Range;
use std::ops::{Index, IndexMut};

/// Represents a memory controller with arbitrary address bus width
/// that routes read/write requests to the appropriate memory chip.
pub struct MemoryController<T> {
    /// A one-to-one map of string identifiers to memory chips.
    pub chips: HashMap<&'static str, MemoryChip>,
    /// A many-to-one map of memory address ranges to connected chips.
    pub map: RangeInclusiveMap<T, &'static str>,
    /// A user-defined function to handle out-of-bounds byte reads.
    pub on_read_out_of_bounds: Box<dyn Fn(T) -> u8>,
}

impl<T: PrimInt + Into<usize> + rangemap::StepLite> MemoryController<T> {
    pub fn new() -> Self {
        MemoryController {
            chips: HashMap::new(),
            map: RangeInclusiveMap::new(),
            on_read_out_of_bounds: Box::new(|_| 0),
        }
    }

    /// Constructs a memory controller with a single memory chip.
    /// The chip is assumed to cover the entire memory map.
    pub fn from_chip(chip_id: &'static str, chip: MemoryChip) -> Self {
        let mut controller = MemoryController::new();
        controller.chips.insert(chip_id, chip);
        controller.map.insert(T::zero()..=T::max_value(), chip_id);

        controller
    }

    pub fn read_byte(&self, addr: T) -> u8 {
        if let Some(chip) = self.map.get(&addr).and_then(|&id| self.chips.get(id)) {
            chip[addr.into()]
        } else {
            (self.on_read_out_of_bounds)(addr)
        }
    }

    pub fn write_byte(&mut self, addr: T, val: u8) {
        if let Some(id) = self.map.get(&addr) {
            if let Some(chip) = self.chips.get_mut(id) {
                chip[addr.into()] = val;
            }
        }
    }
}

impl<T: PrimInt + Into<usize> + rangemap::StepLite> Default for MemoryController<T> {
    fn default() -> Self {
        Self::new()
    }
}

/// Represents a memory chip of arbitrary address bus width.
/// Chip size must be a power of 2 for proper address bit masking.
pub struct MemoryChip {
    data: Vec<u8>,
}

impl MemoryChip {
    pub fn new(size: usize) -> Self {
        if size & (size - 1) != 0 {
            panic!("memory chip size must be a non-zero power of 2");
        }

        MemoryChip {
            data: vec![0; size],
        }
    }
}

impl Index<usize> for MemoryChip {
    type Output = u8;

    fn index(&self, idx: usize) -> &Self::Output {
        let mask = self.data.len() - 1;
        &self.data[idx & mask]
    }
}

impl IndexMut<usize> for MemoryChip {
    fn index_mut(&mut self, idx: usize) -> &mut Self::Output {
        let mask = self.data.len() - 1;
        &mut self.data[idx & mask]
    }
}

impl Index<Range<usize>> for MemoryChip {
    type Output = [u8];

    fn index(&self, range: Range<usize>) -> &Self::Output {
        let mask = self.data.len() - 1;
        &self.data[(range.start & mask)..(range.end & mask)]
    }
}

impl IndexMut<Range<usize>> for MemoryChip {
    fn index_mut(&mut self, range: Range<usize>) -> &mut Self::Output {
        let mask = self.data.len() - 1;
        &mut self.data[(range.start & mask)..(range.end & mask)]
    }
}

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

    #[test]
    fn read_mapped_memory() {
        let mut controller = MemoryController::<u16>::new();
        controller.chips.insert("RAM", MemoryChip::new(0x10000));
        controller.map.insert(0x0000..=0xFFFF, "RAM");

        controller.chips.get_mut("RAM").unwrap()[0x0000] = 0x0F;
        controller.chips.get_mut("RAM").unwrap()[0x5555] = 0x55;
        controller.chips.get_mut("RAM").unwrap()[0xFFFF] = 0xF0;

        assert_eq!(0x0F, controller.read_byte(0x0000));
        assert_eq!(0x55, controller.read_byte(0x5555));
        assert_eq!(0xF0, controller.read_byte(0xFFFF));
    }

    #[test]
    fn write_mapped_memory() {
        let mut controller = MemoryController::<u16>::new();
        controller.chips.insert("RAM", MemoryChip::new(0x10000));
        controller.map.insert(0x0000..=0xFFFF, "RAM");

        controller.write_byte(0x0000, 0x0F);
        controller.write_byte(0x5555, 0x55);
        controller.write_byte(0xFFFF, 0xF0);

        assert_eq!(0x0F, controller.chips.get("RAM").unwrap()[0x0000]);
        assert_eq!(0x55, controller.chips.get("RAM").unwrap()[0x5555]);
        assert_eq!(0xF0, controller.chips.get("RAM").unwrap()[0xFFFF]);
    }

    #[test]
    fn read_write_banked_memory() {
        let mut controller = MemoryController::<u16>::new();
        controller.chips.insert("Bank 1", MemoryChip::new(0x10000));
        controller.chips.insert("Bank 2", MemoryChip::new(0x10000));
        controller.map.insert(0x0000..=0xFFFF, "Bank 1");

        // Write to bank 1
        controller.write_byte(0x5555, 0x55);

        // Switch to bank 2
        controller.map.insert(0x0000..=0xFFFF, "Bank 2");
        assert_eq!(0x00, controller.read_byte(0x5555));

        // Write to bank 2
        controller.write_byte(0x5555, 0x66);

        // Switch to bank 1
        controller.map.insert(0x0000..=0xFFFF, "Bank 1");
        assert_eq!(0x55, controller.read_byte(0x5555));

        // Switch to bank 2
        controller.map.insert(0x0000..=0xFFFF, "Bank 2");
        assert_eq!(0x66, controller.read_byte(0x5555));
    }

    #[test]
    fn read_write_overlapping_memory() {
        let mut controller = MemoryController::<u16>::new();
        controller.chips.insert("Chip 1", MemoryChip::new(0x10000));
        controller.chips.insert("Chip 2", MemoryChip::new(0x1000));
        controller.map.insert(0x0000..=0xFFFF, "Chip 1");

        // Write to chip 1
        controller.write_byte(0x0000, 0x0F);
        controller.write_byte(0x5555, 0x55);
        controller.write_byte(0xFFFF, 0xF0);

        // Connect chip 2
        controller.map.insert(0x5000..=0x5FFF, "Chip 2");
        assert_eq!(0x0F, controller.read_byte(0x0000));
        assert_eq!(0x00, controller.read_byte(0x5555));
        assert_eq!(0xF0, controller.read_byte(0xFFFF));

        // Write to bank 2
        controller.write_byte(0x5555, 0x66);
        assert_eq!(0x0F, controller.read_byte(0x0000));
        assert_eq!(0x66, controller.read_byte(0x5555));
        assert_eq!(0xF0, controller.read_byte(0xFFFF));

        // Disconnect chip 2
        controller.map.insert(0x0000..=0xFFFF, "Chip 1");
        assert_eq!(0x0F, controller.read_byte(0x0000));
        assert_eq!(0x55, controller.read_byte(0x5555));
        assert_eq!(0xF0, controller.read_byte(0xFFFF));
    }

    #[test]
    fn read_out_of_bounds() {
        let mut controller = MemoryController::<u16>::new();
        controller.chips.insert("RAM", MemoryChip::new(0x1000));
        controller.map.insert(0x0000..=0x0FFF, "RAM");

        // Check the default out-of-bounds read behavior.
        assert_eq!(0, controller.read_byte(0xFFFF));

        // Define open bus behavior for out-of-bounds reads.
        controller.on_read_out_of_bounds = Box::new(|addr| addr as u8);
        assert_eq!(0xFF, controller.read_byte(0xFFFF));
    }

    #[test]
    fn write_out_of_bounds() {
        let mut controller = MemoryController::<u16>::new();
        controller.chips.insert("RAM", MemoryChip::new(0x1000));
        controller.map.insert(0x0000..=0x0FFF, "RAM");

        // Out-of-bounds write behavior is a no-op.
        controller.write_byte(0xFFFF, 0xFF);
        assert_eq!(controller.chips.get("RAM").unwrap().data, vec![0; 0x1000]);
    }

    #[test]
    fn read_write_missing_chip() {
        let mut controller = MemoryController::<u16>::new();
        controller.map.insert(0x0000..=0xFFFF, "RAM");

        // Reads and writes on a missing chip should behave as out-of-bounds.
        controller.on_read_out_of_bounds = Box::new(|addr| addr as u8);
        controller.write_byte(0xFFFF, 0x55);
        assert_eq!(0xFF, controller.read_byte(0xFFFF));
    }
}