corevm-engine 0.1.28

use corevm_host::{Range, RangeSet, PAGE_SIZE};
use log::{debug, trace};

/// Maps memory blocks of arbitrary sizes onto the heap address range.
#[derive(Debug)]
pub struct PageMapper {
	/// Heap page range.
	heap_page_range: Range,
	/// Mapped page indices (i.e. address / 4096).
	pub pages: RangeSet,
}

impl PageMapper {
	pub fn new(pages: RangeSet, heap_page_range: Range) -> Self {
		// Sanity checks.
		if let Some(range) = pages.enclosing_range() {
			let next_page_index = range.end;
			assert!(
				(heap_page_range.start..=heap_page_range.end).contains(&next_page_index),
				"Heap page range {heap_page_range:?} doesn't contain the next page index {next_page_index}"
			);
		}
		Self { heap_page_range, pages }
	}

	pub fn map(&mut self, num_pages: u32) -> Option<Range> {
		if num_pages == 0 {
			return None;
		}
		let Some(page_range) = self.find_unmapped_page_range(num_pages) else {
			debug!("Failed to map {num_pages} guest memory page(s): no next address");
			return None;
		};
		self.pages.insert(page_range.clone());
		trace!(
			"Mapped guest memory pages {:#x}..{:#x}",
			page_range.start * PAGE_SIZE,
			page_range.end * PAGE_SIZE,
		);
		Some(page_range)
	}

	pub fn unmap(&mut self, start_page: u32, end_page: u32) {
		self.pages.remove(&Range::new(start_page, end_page));
	}

	pub fn is_mapped(&self, page: u32) -> bool {
		self.pages.as_ref().iter().any(|range| range.contains(page))
	}

	fn find_unmapped_page_range(&mut self, num_pages: u32) -> Option<Range> {
		let start_page = match self.pages.enclosing_range() {
			Some(range) => range.end,
			None => self.heap_page_range.start,
		};
		let end_page = start_page.checked_add(num_pages)?;
		if end_page > self.heap_page_range.end {
			return None;
		}
		Some(Range::new(start_page, end_page))
	}
}

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

	#[test]
	fn map_works() {
		let mut mapper = PageMapper::new(Default::default(), Range::new(1, 10));
		assert_eq!(None, mapper.map(0));
		assert_eq!(None, mapper.map(10));
		assert_eq!(None, mapper.map(u32::MAX));
		assert_eq!(Some(Range::new(1, 2)), mapper.map(1));
		mapper.unmap(0, u32::MAX);
		assert_eq!(0, mapper.pages.as_ref().len(), "{mapper:?}");
		assert_eq!(Some(Range::new(1, 10)), mapper.map(9));
		mapper.unmap(0, u32::MAX);
		assert_eq!(0, mapper.pages.as_ref().len(), "{mapper:?}");
		assert_eq!(Some(Range::new(1, 2)), mapper.map(1));
	}
}