corevm_engine/

mem.rs

use crate::{
	get_work_output_len, hash_raw, AddressKind, InnerVm, OuterVm, OutputBuffers, PageMapper,
	ProgramData, MAX_TOTAL_OUTPUT_BLOB_SIZE,
};
use corevm_host::{OutputStream, PageSegment, Range, RangeSet, PAGE_SIZE};
use jam_pvm_common::ApiError;
use jam_types::{max_exports, Hash, VecMap, VecSet, SEGMENT_LEN};
use log::{debug, trace};

/// Inner VM's memory manager.
///
/// Allocates and deallocates inner VM's memory pages and handles page faults.
pub struct MemoryManager<I: InnerVm> {
	pub(crate) inner_vm: I,
	program: ProgramData,
	/// All memory pages that were read from/written to across all program runs but haven't been
	/// unmapped yet.
	resident_pages: RangeSet,
	/// Memory pages that were imported by the builder.
	imported_pages: VecMap<u64, PageSegment>,
	/// Memory pages that were touched (faulted) during the current program run.
	///
	/// Include deallocated pages.
	touched_pages: VecSet<u64>,
	touched_ro_pages: VecSet<u64>,
	touched_imported_pages: VecSet<u64>,
	/// Should be equal to `WorkItem::export_count`.
	export_count: usize,
	/// Guest program's output streams.
	output: OutputBuffers,
	heap_page_mapper: PageMapper,
	auth_output_len: usize,
}

impl<I: InnerVm> MemoryManager<I> {
	pub fn new<O: OuterVm>(
		inner_vm: I,
		outer_vm: &mut O,
		program: ProgramData,
		heap_pages: RangeSet,
		resident_pages: RangeSet,
		export_count: usize,
		auth_output_len: usize,
	) -> Result<Self, ApiError> {
		let heap_address_range = program.heap_range();
		let heap_page_range =
			Range::new(heap_address_range.start / PAGE_SIZE, heap_address_range.end / PAGE_SIZE);
		let heap_page_mapper = PageMapper::new(heap_pages, heap_page_range);
		assert!(export_count <= max_exports() as usize);
		let mut man = Self {
			inner_vm,
			program,
			resident_pages,
			imported_pages: Default::default(),
			touched_pages: Default::default(),
			touched_ro_pages: Default::default(),
			touched_imported_pages: Default::default(),
			export_count,
			output: Default::default(),
			heap_page_mapper,
			auth_output_len,
		};
		man.import_pages(outer_vm);
		Ok(man)
	}

	pub fn alloc(&mut self, size: u64) -> Option<(u64, u64)> {
		let size = size.checked_next_multiple_of(PAGE_SIZE)?;
		let num_pages = size / PAGE_SIZE;
		let page_range = self.heap_page_mapper.map(num_pages)?;
		let address = page_range.start * PAGE_SIZE;
		let size = (page_range.end - page_range.start) * PAGE_SIZE;
		Some((address, size))
	}

	pub fn dealloc(&mut self, address: u64, size: u64) -> Result<(), ApiError> {
		let start_page = address / PAGE_SIZE;
		let end_page = address.saturating_add(size).div_ceil(PAGE_SIZE);
		self.heap_page_mapper.unmap(start_page, end_page);
		self.resident_pages.remove(&Range::new(start_page, end_page));
		self.inner_vm.void(start_page, end_page - start_page)?;
		trace!("Freed guest memory block {:#x}..{:#x}", address, address + size);
		Ok(())
	}

	/// Handle page fault at the specified `address`.
	pub fn touch(&mut self, address: u64) -> Result<(), PageFaultError> {
		let Some(class) = self.program.classify_address(address) else {
			// Memory map doesn't have this address.
			trace!("Page fault at out-of-range address {:#x}", address);
			return Err(ApiError::OutOfBounds.into());
		};
		// We don't add RO data pages to `self.touched_pages` because they don't need to be
		// exported.
		let page = address / PAGE_SIZE;
		if class == AddressKind::RoData {
			if self.touched_ro_pages.contains(&address) {
				return Ok(());
			}
			trace!("Page fault at {page}/{address:#x}/{class:?}");
			self.inner_vm.poke_ro_data_page(address, &self.program)?;
			self.touched_ro_pages.insert(address);
			return Ok(());
		}
		// RW data, stack and heap need to be exported.
		if self
			.num_exported_pages()
			.checked_add(self.output.segment_count())
			.unwrap_or(usize::MAX) >=
			self.export_count
		{
			// The page will not fit into work package exports.
			trace!(
				"Page fault at {page}/{address:#x}/{class:?}: \
                The page will not fit into work package exports"
			);
			return Err(PageFaultError::PageFault { page, num_pages: 1 });
		}
		let resident_page_inserted = if self.resident_pages.contains_index(page) {
			// We know this page.
			let Some(segment) = self.imported_pages.get(&address) else {
				// This page was allocated by some previous work package,
				// however, the builder haven't imported this page to the current work package.
				// We error out to let the builder import this page for the next program run.
				trace!("Page fault at {page}/{address:#x}/{class:?} (missing page)");
				return Err(PageFaultError::PageFault { page, num_pages: 1 });
			};
			trace!("Page fault at {page}/{address:#x}/{class:?} (resident page)");
			self.inner_vm.zero_poke(segment.page(), address)?;
			self.touched_imported_pages.insert(address);
			false
		} else {
			// We don't know this page. Initialize the page based on the address range.
			match class {
				AddressKind::RoData => {
					// RO pages are handled above.
					unreachable!()
				},
				AddressKind::RwData => {
					trace!("Page fault at {page}/{address:#x}/{class:?}");
					self.inner_vm.poke_rw_data_page(address, &self.program)?;
				},
				AddressKind::Stack => {
					trace!("Page fault at {page}/{address:#x}/{class:?}");
					self.inner_vm.zero_stack_page(address, &self.program)?;
				},
				AddressKind::Heap => {
					if !self.heap_page_mapper.is_mapped(address / PAGE_SIZE) {
						panic!("Page fault at {page}/{address:#x}/{class:?} (unmapped)");
					}
					trace!("Page fault at {page}/{address:#x}/{class:?}");
					self.inner_vm.zero_heap_page(address, &self.program)?;
				},
			}
			self.resident_pages.insert(Range::new(page, page + 1));
			true
		};
		let touched_pages_inserted = self.touched_pages.insert(address);
		// Check that we're still within the maximum bounds of the work output.
		if (resident_page_inserted || touched_pages_inserted) &&
			self.check_work_output_len().is_err()
		{
			// The page will not fit into the work output.
			trace!("Page fault at {page}/{address:#x}/{class:?}: The page will not fit into work output");
			// Revert all insertions that we made.
			if resident_page_inserted {
				self.resident_pages.remove(&Range::new(page, page + 1));
			}
			if touched_pages_inserted {
				self.touched_pages.remove(&address);
			}
			return Err(PageFaultError::PageFault { page, num_pages: 1 });
		}
		Ok(())
	}

	fn check_work_output_len(&self) -> Result<(), ()> {
		let work_output_len = get_work_output_len(
			&self.heap_page_mapper.pages,
			&self.resident_pages,
			self.touched_imported_pages.len(),
			self.num_exported_pages(),
		);
		if work_output_len + self.auth_output_len >= MAX_TOTAL_OUTPUT_BLOB_SIZE {
			return Err(());
		}
		Ok(())
	}

	fn num_exported_pages(&self) -> usize {
		self.touched_pages
			.iter()
			.filter(|address| {
				// Filter out deallocated pages.
				let page = *address / PAGE_SIZE;
				self.resident_pages.contains_index(page)
			})
			.count()
	}

	/// Append data to the guest program's output stream.
	pub fn append_output(
		&mut self,
		stream: OutputStream,
		inner_src: u64,
		len: u64,
	) -> Result<(), AppendOutputError> {
		if len == 0 {
			// Succeed regardless of the validity of `inner_src` to minimize potential page faults.
			return Ok(());
		}
		let Some(new_segment_count) = self.output.segment_count_after(stream, len as usize) else {
			// The `len` is too large for appending to the current buffer.
			trace!("Output limit reached while appending {len} byte(s) to {stream:?}");
			return Err(AppendOutputError::OutputLimitReached);
		};
		if self.num_exported_pages().saturating_add(new_segment_count) >= self.export_count {
			trace!("Output limit reached while appending {len} byte(s) to {stream:?}");
			// TODO @ivan This error can trigger in the middle of the execution because we've run
			// out of exports, however, at the end of the execution the program might deallocate
			// some pages and free space can be reused for the output.
			return Err(AppendOutputError::OutputLimitReached);
		}
		if inner_src == 0 {
			// Confirm that the operation can be carried out successfully when repeated with a
			// valid address.
			trace!("Confirm {len} byte(s) can be appended to {stream:?} without reaching the output limit");
			return Ok(());
		}
		// Touch the pages first because they might have been swapped out in the previous run.
		let page_aligned_inner_src = inner_src - inner_src % PAGE_SIZE;
		let page_aligned_inner_src_end = inner_src
			.checked_add(len)
			.and_then(|x| x.checked_next_multiple_of(PAGE_SIZE))
			.ok_or(AppendOutputError::ApiError(ApiError::OutOfBounds))?;
		let num_pages = (page_aligned_inner_src_end - page_aligned_inner_src) / PAGE_SIZE;
		for address in
			(page_aligned_inner_src..page_aligned_inner_src_end).step_by(PAGE_SIZE as usize)
		{
			let page = address / PAGE_SIZE;
			if self.touched_pages.contains(&address) && self.resident_pages.contains_index(page) {
				// We already handled this page fault.
				continue;
			}
			self.touch(address).map_err(|e| match e {
				// Replace normal page fault with the one that includes all the pages that we
				// need to move forward.
				PageFaultError::PageFault { .. } => AppendOutputError::PageFault {
					page: page_aligned_inner_src / PAGE_SIZE,
					num_pages,
				},
				e => e.into(),
			})?;
		}
		let buf = self.output.pre_allocate(stream, len as usize);
		self.inner_vm.peek_into(buf, inner_src)?;
		trace!(
			"Appended {len} bytes to {stream:?}: src={inner_src:#x}, utf8={:?}",
			if matches!(stream, OutputStream::Stdout | OutputStream::Stderr) {
				core::str::from_utf8(buf)
			} else {
				Ok("")
			}
		);
		Ok(())
	}

	/// Exports the modified pages as segments.
	pub fn export<O: OuterVm>(
		&mut self,
		outer_vm: &mut O,
	) -> Result<
		(RangeSet, RangeSet, VecMap<u64, Hash>, VecMap<u64, Hash>, u32, [u32; OutputStream::COUNT]),
		ApiError,
	> {
		let stream_len = self.output.stream_len();
		let output_len = self.output.total_len() as u32;
		// Export all touched pages.
		let mut touched_pages: VecMap<u64, Hash> = VecMap::new();
		let mut num_exported_pages = 0;
		for address in self.touched_pages.iter() {
			let page = address / PAGE_SIZE;
			let hash = if self.resident_pages.contains_index(page) {
				let mut segment = PageSegment::zero(*address);
				self.inner_vm.peek_into(segment.page_mut(), *address)?;
				outer_vm.export(segment.as_ref())?;
				num_exported_pages += 1;
				trace!("Exported regular page {:#x}", address);
				hash_raw(segment.as_ref())
			} else {
				trace!("Not exporting freed regular page {:#x}", address);
				// The page was freed.
				[0; 32]
			};
			// Update page hash.
			touched_pages.insert(*address, hash);
		}
		// Export output as segments.
		self.output.export_segments(|segment| outer_vm.export(&segment[..]))?;
		let num_output_segments = output_len.div_ceil(SEGMENT_LEN as u32) as usize;
		// We need to make the total no. of exported pages equal `WorkItem::export_count` for
		// the work output to be considered valid. Hence we export zero-filled pages with null
		// addresses.
		let total_segment_count = num_exported_pages + num_output_segments;
		for _ in total_segment_count..self.export_count {
			let segment = PageSegment::null();
			outer_vm.export(segment.as_ref())?;
		}
		debug!(
			"Exported page(s): {:?}",
			touched_pages
				.iter()
				.map(|(address, _)| {
					let page = address / PAGE_SIZE;
					Range::new(page, page + 1)
				})
				.collect::<RangeSet>()
		);
		debug!(
			"Exported {} segments: {} memory page(s), {} output segment(s), {} null segment(s)",
			self.export_count,
			num_exported_pages,
			output_len.div_ceil(SEGMENT_LEN as u32),
			self.export_count.saturating_sub(total_segment_count)
		);
		let resident_pages = core::mem::take(&mut self.resident_pages);
		let heap_pages = core::mem::take(&mut self.heap_page_mapper.pages);
		let touched_imported_pages = self
			.imported_pages
			.iter()
			.filter(|(address, _segment)| self.touched_imported_pages.contains(address))
			.map(|(address, segment)| (*address, hash_raw(segment.as_ref())))
			.collect();
		Ok((
			heap_pages,
			resident_pages,
			touched_imported_pages,
			touched_pages,
			num_exported_pages as u32,
			stream_len,
		))
	}

	fn import_pages<O: OuterVm>(&mut self, outer_vm: &mut O) {
		let mut i = 0;
		while let Some(page_segment) = outer_vm.import(i) {
			let address = page_segment.address();
			let page = address / PAGE_SIZE;
			let Some(class) = self.program.classify_address(address) else {
				panic!("Failed to import segment {i}: Invalid page {page}/{address:#x}");
			};
			if !self.resident_pages.contains_index(page) {
				panic!(
					"Failed to import segment {i}: Unknown page {page}/{address:#x}/{class:?}, known pages {:?}",
					self.resident_pages
				);
			}
			trace!("Imported page {page}/{address:#x}/{class:?}");
			self.imported_pages.insert(address, page_segment);
			i += 1;
		}
		debug!(
			"Imported page(s): {:?}",
			self.imported_pages
				.iter()
				.map(|(address, _)| {
					let page = address / PAGE_SIZE;
					Range::new(page, page + 1)
				})
				.collect::<RangeSet>()
		);
	}
}

#[derive(Debug)]
pub enum PageFaultError {
	ApiError(ApiError),
	PageFault { page: u64, num_pages: u64 },
}

impl From<ApiError> for PageFaultError {
	fn from(e: ApiError) -> Self {
		Self::ApiError(e)
	}
}

#[derive(Debug)]
pub enum AppendOutputError {
	ApiError(ApiError),
	PageFault { page: u64, num_pages: u64 },
	OutputLimitReached,
}

impl From<ApiError> for AppendOutputError {
	fn from(e: ApiError) -> Self {
		Self::ApiError(e)
	}
}

impl From<PageFaultError> for AppendOutputError {
	fn from(e: PageFaultError) -> Self {
		match e {
			PageFaultError::ApiError(e) => Self::ApiError(e),
			PageFaultError::PageFault { page, num_pages } => Self::PageFault { page, num_pages },
		}
	}
}

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

	use alloc::{vec, vec::Vec};
	use jam_pvm_common::InvokeOutcome;
	use jam_types::{ServiceId, SignedGas};
	use polkavm::MemoryMapBuilder;
	use rand::{seq::IndexedRandom, Rng, RngCore};

	#[test]
	fn touch_and_append_output_works() {
		let mut rng = rand::rng();
		let (mut mem, _) = MemoryManager::with_data(
			{
				let len = rng.random_range(0..=2 * PAGE_SIZE);
				let mut data = vec![0_u8; len as usize];
				rng.fill_bytes(&mut data[..]);
				data
			},
			{
				let len = rng.random_range(0..=2 * PAGE_SIZE);
				let mut data = vec![0_u8; len as usize];
				rng.fill_bytes(&mut data[..]);
				data
			},
		);
		let ro_page_range = {
			let r = mem.program.ro_data_range();
			let ro_start_page = r.start / PAGE_SIZE;
			let ro_end_page = r.end / PAGE_SIZE;
			ro_start_page..ro_end_page
		};
		let ro_data_page_range = {
			let r = mem.program.ro_data_range();
			let data_len = mem.program.ro_data.len() as u64;
			let ro_start_page = r.start / PAGE_SIZE;
			let ro_end_page = (r.start + data_len).div_ceil(PAGE_SIZE);
			ro_start_page..ro_end_page
		};
		let rw_page_range = {
			let r = mem.program.rw_data_range();
			let rw_start_page = r.start / PAGE_SIZE;
			let rw_end_page = r.end / PAGE_SIZE;
			rw_start_page..rw_end_page
		};
		let rw_data_page_range = {
			let r = mem.program.rw_data_range();
			let data_len = mem.program.rw_data.len() as u64;
			let rw_start_page = r.start / PAGE_SIZE;
			let rw_end_page = (r.start + data_len).div_ceil(PAGE_SIZE);
			rw_start_page..rw_end_page
		};
		let stack_page_range = {
			let r = mem.program.stack_range();
			let stack_start_page = r.start / PAGE_SIZE;
			let stack_end_page = r.end / PAGE_SIZE;
			stack_start_page..stack_end_page
		};
		let heap_page_range = {
			let r = mem.program.heap_range();
			let heap_start_page = r.start / PAGE_SIZE;
			let heap_end_page = r.end / PAGE_SIZE;
			heap_start_page..heap_end_page
		};
		assert_eq!(None, mem.program.classify_address(0));
		assert!(matches!(mem.touch(0), Err(PageFaultError::ApiError(ApiError::OutOfBounds))));
		// Check that `touch` inits the pages and stores them in `resident_pages` and
		// `touched_pages`.
		//
		// - Stack and heap pages are always zeroed.
		// - RO/RW data pages are either initialized with the static data from PolkaVM program or
		//   are zeroed if the address goes beyond the static data.
		macro_rules! check_touch {
			($range: expr, readonly = $read_only: expr, alloc = $alloc: expr, data = $data: expr$(,)*) => {{
				let range = $range;
				let data = $data;
				let (page, address) = if $alloc {
					let (address, _size) = mem.alloc(PAGE_SIZE as u64).unwrap();
					let page = address / PAGE_SIZE as u64;
					(page, address)
				} else {
					let page = rng.random_range(range.clone());
					let address = page * PAGE_SIZE as u64;
					(page, address)
				};
				mem.resident_pages = Default::default();
				mem.touched_pages = Default::default();
				mem.inner_vm.calls.clear();
				assert!(mem.touch(address).is_ok());
				if data.is_empty() {
					assert_eq!(
						[Call::Zero { page, num_pages: 1 }].as_slice(),
						mem.inner_vm.calls.as_slice()
					);
				} else {
					let start = (page - range.start) as usize * PAGE_SIZE as usize;
					let end = (start + PAGE_SIZE as usize).min(data.len());
					let outer_src = data[start..end].to_vec();
					assert_eq!(
						[
							Call::Zero { page, num_pages: 1 },
							Call::Poke { outer_src, inner_dst: address }
						]
						.as_slice(),
						mem.inner_vm.calls.as_slice()
					);
				}
				if !$read_only {
					assert_eq!(
						[Range::new(page, page + 1)].as_slice(),
						mem.resident_pages.as_slice()
					);
					assert_eq!(VecSet::from_iter([address]), mem.touched_pages);
				}
			}};
		}
		for _ in 0..1000 {
			// RO pages with no data.
			check_touch!(
				ro_data_page_range.end..ro_page_range.end,
				readonly = true,
				alloc = false,
				data = [0_u8; 0],
			);
			// RO pages with data.
			check_touch!(
				ro_data_page_range.start..ro_data_page_range.end,
				readonly = true,
				alloc = false,
				data = mem.program.ro_data.clone(),
			);
			// RW pages with no data.
			check_touch!(
				rw_data_page_range.end..rw_page_range.end,
				readonly = false,
				alloc = false,
				data = [0_u8; 0],
			);
			// RW pages with data.
			check_touch!(
				rw_data_page_range.start..rw_data_page_range.end,
				readonly = false,
				alloc = false,
				data = mem.program.rw_data.clone(),
			);
			// Stack pages.
			check_touch!(
				stack_page_range.start..stack_page_range.end,
				readonly = false,
				alloc = false,
				data = [0_u8; 0],
			);
			// Heap pages.
			check_touch!(
				heap_page_range.start..heap_page_range.end,
				readonly = false,
				alloc = true,
				data = [0_u8; 0],
			);
			// Append output.
			{
				mem.resident_pages = Default::default();
				mem.touched_pages = Default::default();
				mem.touched_ro_pages = Default::default();
				mem.inner_vm.calls.clear();
				let stream = *OutputStream::ALL.choose(&mut rng).unwrap();
				let len = rng.random_range(1..=PAGE_SIZE);
				let (address, _size) = mem.alloc(len).unwrap();
				let page_aligned_address = address - address % PAGE_SIZE;
				let start_page = page_aligned_address / PAGE_SIZE;
				let end_page = (address + len).next_multiple_of(PAGE_SIZE) / PAGE_SIZE;
				assert!(mem.append_output(stream, address, len).is_ok());
				let mut calls = Vec::new();
				if len != 0 {
					for page in start_page..end_page {
						calls.push(Call::Zero { page, num_pages: 1 });
					}
					calls.push(Call::PeekInto { inner_src: address, len });
				}
				assert_eq!(calls.as_slice(), mem.inner_vm.calls.as_slice());
			}
		}
	}

	#[test]
	fn page_limit_on_touch_works() {
		let (mut mem, _) = MemoryManager::new_for_tests();
		let mut num_touched = 0;
		for _ in 0..max_exports() {
			let (address, _size) = mem.alloc(PAGE_SIZE).unwrap();
			let result = mem.touch(address);
			assert!(mem.touched_pages.len() <= max_exports() as usize);
			assert_eq!(Ok(()), mem.check_work_output_len());
			if result.is_err() {
				break;
			}
			num_touched += 1;
		}
		assert_ne!(0, num_touched);
	}

	#[test]
	fn page_limit_on_output_works() {
		let _ = env_logger::try_init();
		let (mut mem, mut outer_vm) = MemoryManager::new_for_tests();
		mem.inner_vm.randomize_memory = true;
		assert!(
			mem.program.heap_range().end - mem.program.heap_range().start > max_exports() as u64
		);
		let mut rng = rand::rng();
		let (address, _len) = mem.alloc(PAGE_SIZE).unwrap();
		loop {
			let stream = *OutputStream::ALL.choose(&mut rng).unwrap();
			let len = rng.random_range(0..=PAGE_SIZE);
			let result = mem.append_output(stream, address, len);
			assert_eq!(Ok(()), mem.check_work_output_len());
			match result {
				Ok(()) => {},
				Err(AppendOutputError::OutputLimitReached) => {
					break;
				},
				Err(e) => {
					panic!("Unexpected error {e:?}");
				},
			}
		}
		for _ in 0..10 {
			let stream = *OutputStream::ALL.choose(&mut rng).unwrap();
			assert!(mem.append_output(stream, 0, 0).is_ok());
		}
		assert!(mem.touched_pages.len() <= max_exports() as usize);
		assert_eq!(Ok(()), mem.check_work_output_len());
		let (_, _, _, _, num_memory_pages, stream_len) = mem.export(&mut outer_vm).unwrap();
		let output_len: u32 = stream_len.iter().sum();
		let num_output_segments = output_len.div_ceil(SEGMENT_LEN as u32) as usize;
		assert_eq!(num_memory_pages as usize + num_output_segments, outer_vm.real_export_count());
		assert!(outer_vm.exports.len() <= max_exports() as usize);
	}

	#[test]
	fn append_output_edge_cases() {
		let (mut mem, _) = MemoryManager::new_for_tests();
		let mut rng = rand::rng();
		let stream = *OutputStream::ALL.choose(&mut rng).unwrap();
		let result = mem.append_output(stream, mem.program.heap_range().start, u64::MAX);
		assert!(
			matches!(result, Err(AppendOutputError::OutputLimitReached)),
			"Result = {result:?}"
		);
		for kind in AddressKind::ALL {
			let address = mem.program.address_range(kind).end;
			if mem.program.classify_address(address).is_some() {
				// Blocks are laid out in memory one after another.
				continue;
			}
			assert!(
				matches!(
					mem.append_output(stream, address, 1),
					Err(AppendOutputError::ApiError(ApiError::OutOfBounds)),
				),
				"address kind = {kind:?}",
			);
		}
	}

	impl MemoryManager<TestInnerVm> {
		fn new_for_tests() -> (Self, TestOuterVm) {
			Self::with_data(Default::default(), Default::default())
		}

		fn with_data(ro_data: Vec<u8>, rw_data: Vec<u8>) -> (Self, TestOuterVm) {
			let inner_vm = TestInnerVm::default();
			let mut outer_vm = TestOuterVm::default();
			let memory_map = MemoryMapBuilder::new(PAGE_SIZE as u32)
				.ro_data_size((PAGE_SIZE * 10) as u32)
				.rw_data_size((PAGE_SIZE * 10) as u32)
				.stack_size((PAGE_SIZE * 10) as u32)
				.build()
				.expect("Failed to build memory map");
			let program_data = ProgramData::new(&memory_map, ro_data.into(), rw_data.into());
			let man = Self::new(
				inner_vm,
				&mut outer_vm,
				program_data,
				Default::default(),
				Default::default(),
				max_exports() as usize,
				0,
			)
			.unwrap();
			(man, outer_vm)
		}
	}

	#[derive(Default)]
	struct TestInnerVm {
		calls: Vec<Call>,
		randomize_memory: bool,
	}

	impl InnerVm for TestInnerVm {
		fn void(&mut self, page: u64, num_pages: u64) -> Result<(), ApiError> {
			self.calls.push(Call::Void { page, num_pages });
			Ok(())
		}

		fn zero(&mut self, page: u64, num_pages: u64) -> Result<(), ApiError> {
			self.calls.push(Call::Zero { page, num_pages });
			Ok(())
		}

		fn poke(&mut self, outer_src: &[u8], inner_dst: u64) -> Result<(), ApiError> {
			self.calls.push(Call::Poke { outer_src: outer_src.to_vec(), inner_dst });
			Ok(())
		}

		fn peek_into(&mut self, outer_dst: &mut [u8], inner_src: u64) -> Result<(), ApiError> {
			self.calls.push(Call::PeekInto { len: outer_dst.len() as u64, inner_src });
			if self.randomize_memory {
				rand::rng().fill_bytes(outer_dst);
			}
			Ok(())
		}

		fn expunge(self) -> Result<u64, ApiError> {
			unreachable!()
		}

		fn invoke(
			&mut self,
			_gas: SignedGas,
			_regs: [u64; 13],
		) -> Result<(InvokeOutcome, SignedGas, [u64; 13]), ApiError> {
			unreachable!()
		}
	}

	#[derive(Debug, PartialEq, Eq)]
	enum Call {
		Void { page: u64, num_pages: u64 },
		Zero { page: u64, num_pages: u64 },
		Poke { outer_src: Vec<u8>, inner_dst: u64 },
		PeekInto { len: u64, inner_src: u64 },
	}

	#[derive(Default)]
	struct TestOuterVm {
		imports: Vec<PageSegment>,
		exports: Vec<PageSegment>,
	}

	impl TestOuterVm {
		/// The number of exports without counting the trailing null segments.
		fn real_export_count(&self) -> usize {
			let num_null_segments = self
				.exports
				.iter()
				.rev()
				.take_while(|segment| segment.as_ref().iter().all(|&b| b == 0))
				.count();
			self.exports.len() - num_null_segments
		}
	}

	impl OuterVm for TestOuterVm {
		type InnerVm = TestInnerVm;

		fn import(&mut self, i: usize) -> Option<PageSegment> {
			self.imports.get(i).cloned()
		}

		fn export(&mut self, segment: &[u8]) -> Result<(), ApiError> {
			self.exports.push(PageSegment::new(segment.to_vec().try_into().unwrap()));
			Ok(())
		}

		fn lookup(&mut self, _service_id: ServiceId, _hash: &Hash) -> Option<Vec<u8>> {
			unreachable!()
		}

		fn get_export_count(&mut self) -> u16 {
			unreachable!()
		}

		fn get_auth_output_len(&mut self) -> u32 {
			unreachable!()
		}

		fn machine(
			&mut self,
			_code: &[u8],
			_program_counter: u64,
		) -> Result<TestInnerVm, ApiError> {
			unreachable!()
		}
	}
}