Basic memory allocator for a virtual machine in Rust

vm.rs

use std::cmp::Reverse;
use std::collections::BinaryHeap;

struct Mem {
    data: Vec<Option<Vec<u32>>>,
    free_pq: BinaryHeap<Reverse<u32>>,
    len: u32, // FIXME: this is not actually needed for now, we can just use data.len()
}

impl Mem {
    fn new() -> Self {
        Mem {
            data: Vec::new(),
            free_pq: BinaryHeap::new(),
            len: 0,
        }
    }

    fn len(&self) -> u32 {
        self.len
    }

    fn alloc(&mut self, size: u32) -> u32 {
        match self.free_pq.pop() {
            Some(Reverse(addr)) => {
                let v = vec![0; size as usize];
                self.data[addr as usize] = Some(v);
                addr
            }

            None => {
                let v = vec![0; size as usize];
                self.data.push(Some(v));
                self.len += 1;
                self.len - 1
            }
        }
    }

    fn free(&mut self, addr: u32) {
        match self.data.get_mut(addr as usize) {
            Some(Some(v)) => {
                v.clear(); // FIXME: is that needed?
                self.data[addr as usize] = None;
                self.free_pq.push(Reverse(addr));
            }
            Some(None) => panic!(
                "free: attempt to free address {} which is already free",
                addr
            ),
            None => panic!("free: attempt to free unallocated address {}", addr),
        }
    }

    fn read(&self, addr: u32, offset: u32) -> &u32 {
        match &self.data.get(addr as usize) {
            Some(Some(v)) => match v.get(offset as usize) {
                Some(val) => val,
                None => panic!(
                    "read: offset {} is out of bounds for address {} (len: {})",
                    offset,
                    addr,
                    v.len()
                ),
            },
            Some(None) => panic!("read: address {} has been deallocated", addr),
            None => panic!("read: address {} has not been allocated", addr),
        }
    }

    fn write(&mut self, addr: u32, offset: u32, val: u32) {
        match self.data.get_mut(addr as usize) {
            Some(Some(v)) => {
                if (offset as usize) < v.len() {
                    v[offset as usize] = val;
                } else {
                    panic!(
                        "write: offset {} is out of bounds for address {} (len: {})",
                        offset,
                        addr,
                        v.len()
                    );
                }
            }
            Some(None) => panic!("write: address {} has been deallocated", addr),
            None => panic!("write: address {} has not been allocated", addr),
        }
    }
}

// FIXME: convert this function to tests
fn main() {
    let mut mem = Mem::new();

    let m0 = mem.alloc(10);
    let m1 = mem.alloc(20);
    let m2 = mem.alloc(30);
    assert_eq!(mem.len(), 3);
    assert_eq!(m0, 0);
    assert_eq!(m1, 1);
    assert_eq!(m2, 2);

    mem.free(m0);
    mem.free(m1);
    mem.free(m2);

    // should panic:
    // assert_eq!(mem.read(m0, 1), &0);

    let m3 = mem.alloc(40);
    assert_eq!(m3, 0);

    let m4 = mem.alloc(50);
    assert_eq!(m4, 1);

    assert_eq!(mem.len(), 3);

    assert_eq!(mem.read(0, 1), &0);

    // should panic:
    // assert_eq!(mem.read(0, 100), &0);
    // assert_eq!(mem.read(20, 1), &0);

    mem.write(0, 1, 384);
    assert_eq!(mem.read(0, 1), &384);
    mem.free(0);
}