corevm-engine 0.1.28

use crate::{Engine, InnerVm, OuterVm};
use alloc::{sync::Arc, vec::Vec};
use bytes::Bytes;
use core::sync::atomic::{AtomicBool, Ordering};
use corevm_host::{fs, PageNum, PAGE_SIZE};
use jam_pvm_common::{ApiError, InvokeOutcome};
use jam_types::{SegmentBytes, SignedGas, SEGMENT_LEN};
use polkavm::{InterruptKind, ProgramCounter, RawInstance, Reg};

/// An engine that simulates outer and inner VM host-calls.
///
/// The main use case is to predict which memory pages a CoreVM service will access in the next
/// run.
pub type Simulator<R> = Engine<OuterVmSimulator<R>>;

/// An implementation of the host-calls related to outer VMs.
///
/// Forwards arguments to the corresponding PolkaVM calls most of the time.
/// The implementation should behave exactly like `RefineCallContext` from `jam-node` for the
/// simulation to be correct.
pub struct OuterVmSimulator<R: fs::ReadBlock> {
	engine: Arc<polkavm::Engine>,
	fs: R,
	exports: Vec<SegmentBytes>,
	stopped: Arc<AtomicBool>,
}

impl<R: fs::ReadBlock> OuterVmSimulator<R> {
	/// Create new outer VM simulator.
	pub fn new(
		engine: Arc<polkavm::Engine>,
		fs: R,
		stopped: Arc<AtomicBool>,
	) -> Result<Self, polkavm::Error> {
		Ok(Self { engine, fs, exports: Default::default(), stopped })
	}

	/// Reset exports.
	pub fn reset(&mut self) {
		self.exports.clear();
	}

	pub fn exports(&self) -> &[SegmentBytes] {
		self.exports.as_slice()
	}

	pub fn into_exports(self) -> Vec<SegmentBytes> {
		self.exports
	}

	pub fn exports_mut(&mut self) -> &mut Vec<SegmentBytes> {
		&mut self.exports
	}

	pub fn fs_mut(&mut self) -> &mut R {
		&mut self.fs
	}
}

impl<R: fs::ReadBlock> OuterVm for OuterVmSimulator<R> {
	type InnerVm = InnerVmSimulator;

	fn export(&mut self, segment: &[u8; SEGMENT_LEN]) -> Result<(), ApiError> {
		self.exports.push(segment.into());
		Ok(())
	}

	fn read_file_block(&mut self, block_ref: &fs::BlockRef) -> Option<Bytes> {
		self.fs.read_block(block_ref).ok()
	}

	fn machine(&mut self, code: &[u8], program_counter: u64) -> Result<InnerVmSimulator, ApiError> {
		let mut parts = polkavm::ProgramParts::empty(polkavm::program::InstructionSetKind::JamV1);
		parts.code_and_jump_table = code.into();
		let pvm_blob = polkavm::ProgramBlob::from_parts(parts).map_err(api_error)?;
		let mut config = polkavm::ModuleConfig::default();
		config.set_gas_metering(Some(polkavm::GasMeteringKind::Sync));
		config.set_dynamic_paging(true);
		let module =
			polkavm::Module::from_blob(&self.engine, &config, pvm_blob).map_err(api_error)?;
		let mut instance = module.instantiate().map_err(api_error)?;
		instance.set_next_program_counter(ProgramCounter(
			program_counter.try_into().map_err(api_error)?,
		));
		Ok(InnerVmSimulator { instance, stopped: self.stopped.clone() })
	}
}

fn api_error(e: impl core::fmt::Display) -> ApiError {
	log::debug!("Failed to create inner VM: {e}");
	ApiError::ActionInvalid
}

/// An implementation of the host-calls related to inner VMs.
///
/// Forwards arguments to the corresponding PolkaVM calls most of the time.
/// The implementation should behave exactly like `RefineCallContext` from `jam-node` for the
/// simulation to be correct.
pub struct InnerVmSimulator {
	instance: RawInstance,
	stopped: Arc<AtomicBool>,
}

impl InnerVm for InnerVmSimulator {
	fn void(&mut self, page: PageNum, num_pages: u32) -> Result<(), ApiError> {
		let (address, length) = to_address(page.0, num_pages).inspect_err(|_| {
			log::debug!("Overflow in `void`: page={page}, num_pages={num_pages}");
		})?;
		self.instance.free_pages(address, length).unwrap_or_else(|e| {
			panic!("Failed to free memory region {address:#x}..{:#x}: {e}", address + length)
		});
		Ok(())
	}

	fn zero(&mut self, page: PageNum, num_pages: u32) -> Result<(), ApiError> {
		let (address, length) = to_address(page.0, num_pages).inspect_err(|_| {
			log::debug!("Overflow in `zero`: page={page}, num_pages={num_pages}");
		})?;
		self.instance
			.zero_memory_with_memory_protection(
				address,
				length,
				polkavm::MemoryProtection::ReadWrite,
			)
			.unwrap_or_else(|e| {
				panic!("Failed to zero memory region {address:#x}..{:#x}: {e}", address + length)
			});
		Ok(())
	}

	fn poke(&mut self, outer_src: &[u8], inner_dst: u32) -> Result<(), ApiError> {
		let len: u32 = outer_src.len().try_into().map_err(|_| {
			log::debug!("Overflow in `poke`: len={}", outer_src.len());
			ApiError::OutOfBounds
		})?;
		match self.instance.write_memory(inner_dst, outer_src) {
			Ok(()) => Ok(()),
			Err(polkavm::MemoryAccessError::OutOfRangeAccess { .. }) => {
				log::debug!("`poke`: {inner_dst:#x}..{:#x} is not writable", inner_dst + len);
				Err(ApiError::OutOfBounds)
			},
			Err(error) => {
				panic!(
					"Failed to write memory region {inner_dst:#x}..{:#x}: {error}",
					inner_dst + len
				);
			},
		}
	}

	fn peek_into(&mut self, outer_dst: &mut [u8], inner_src: u32) -> Result<(), ApiError> {
		let len: u32 = outer_dst.len().try_into().map_err(|_| {
			log::debug!("Overflow in `peek`: len={}", outer_dst.len());
			ApiError::OutOfBounds
		})?;
		self.instance.read_memory_into(inner_src, outer_dst).map_err(|e| {
			log::debug!("Failed to read memory region {inner_src:#x}..{:#x}: {e}", inner_src + len);
			ApiError::OutOfBounds
		})?;
		Ok(())
	}

	fn expunge(self) -> Result<u64, ApiError> {
		Ok(self.instance.next_program_counter().map_or(0, |pc| pc.0.into()))
	}

	fn invoke(
		&mut self,
		gas: SignedGas,
		mut regs: [u64; 13],
	) -> Result<(InvokeOutcome, SignedGas, [u64; 13]), ApiError> {
		if self.stopped.load(Ordering::Relaxed) {
			// This outcome shouldn't be used anywhere.
			return Ok((InvokeOutcome::Panic, 0, [0; 13]));
		}
		self.instance.set_gas(gas);
		for (reg, value) in Reg::ALL.into_iter().zip(regs) {
			self.instance.set_reg(reg, value)
		}
		let result = self.instance.run().expect("PolkaVM runtime failure");
		let gas = self.instance.gas();
		for (reg, value) in Reg::ALL.into_iter().zip(regs.iter_mut()) {
			*value = self.instance.reg(reg)
		}
		let outcome = match result {
			InterruptKind::Finished => InvokeOutcome::Halt,
			InterruptKind::Ecalli(index) =>
				InvokeOutcome::HostCallFault(index as i32 as i64 as u64),
			InterruptKind::Trap => {
				log::warn!("Panicked at {:?}", self.instance.program_counter());
				InvokeOutcome::Panic
			},
			InterruptKind::NotEnoughGas => InvokeOutcome::OutOfGas,
			InterruptKind::Segfault(segfault) =>
				InvokeOutcome::PageFault(segfault.page_address as _),
			// This must be explicitly enabled to be emitted.
			InterruptKind::Step => unreachable!(),
		};
		Ok((outcome, gas, regs))
	}
}

fn check_page_range(page: u32, num_pages: u32) -> bool {
	(MIN_PAGE..=MAX_PAGE).contains(&page) &&
		num_pages <= MAX_PAGE &&
		page.saturating_add(num_pages) <= MAX_PAGE
}

fn to_address(page: u32, num_pages: u32) -> Result<(u32, u32), ApiError> {
	if !check_page_range(page, num_pages) {
		return Err(ApiError::ActionInvalid);
	}
	Ok((page * PAGE_SIZE, num_pages * PAGE_SIZE))
}

const MIN_PAGE: u32 = 16;
const MAX_PAGE: u32 = u32::MAX / PAGE_SIZE;