corevm-host 0.1.28

use crate::{
	fs, Arg, AudioMode, InputChunk, InputStreams, OutputStream, PageNum, RangeSet, ServiceMessage,
	VideoMode,
};
use alloc::{collections::VecDeque, vec::Vec};
use bytes::Bytes;
use codec::{ConstEncodedLen, Decode, Encode, MaxEncodedLen};
use corevm_types::GuestId;
use jam_types::{
	Hash, SegmentBytes, SegmentTreeRoot, ServiceId, SignedGas, VecMap, MEMO_LEN, SEGMENT_LEN,
};

/// CoreVM-specific [work payload](jam_types::WorkPayload).
///
/// This is the input of the work package.
#[derive(Encode, Decode, Debug, Clone)]
pub struct CoreVmPayload {
	/// Inner VM gas.
	pub gas: SignedGas,
	/// The prior (known) VM state from which the program run should continue.
	///
	/// Equals [`VmState::initial`] on the first program run.
	pub vm_state: VmState,
	/// Execution environment location in the virtual file system.
	pub exec_ref: fs::BlockRef,
}

/// CoreVM-specific extrinsics.
///
/// This is the input of the work package.
#[derive(Debug, Default, Clone)]
pub struct CoreVmExtrinsics {
	/// Imported memory pages; extrinsic 0.
	pub imported_memory_pages: Vec<SegmentBytes>,
	/// Incoming inter-service messages, their sources and their contents; extrinsic 1.
	pub incoming_service_messages: VecDeque<(ServiceMessage, Bytes)>,
	/// Input chunks (video/audio/console); extrinsic 2.
	///
	/// In the extrinsic the number of chunks is not stored, only chunks themselves and their
	/// signatures. Each chunk is signed with [`ExecEnv::input_key`].
	pub input_chunks: Vec<InputChunk>,
}

impl CoreVmExtrinsics {
	/// Max. number of extrinsics a CoreVM-specific work package can have.
	#[doc(hidden)]
	pub const MAX_COUNT: usize = 3;
}

/// CoreVM-specific [work output](jam_types::WorkOutput).
///
/// This is the output of the work package.
#[derive(Encode, Decode, Debug)]
pub struct CoreVmOutput {
	/// The specification of the VM's output data.
	pub vm_output: VmOutput,
	/// The posterior VM state on which the program run was suspended.
	///
	/// Use this state in [`CoreVmPayload`] to continue the execution in another work package.
	pub vm_state: VmState,
	pub old_hash: Hash,
	pub new_hash: Hash,
	/// Memory pages that were imported and read from/written to while refining the package.
	pub touched_imported_pages: VecMap<PageNum, Hash>,
	/// Memory pages that were read from/written to while refining this work package.
	///
	/// Include deallocated pages indicated by zero hashes.
	pub updated_pages: VecMap<PageNum, Hash>,
	/// Execution environment location in the virtual file system.
	pub exec_ref: fs::BlockRef,
	/// Inter-service messages that were processed during the program run.
	pub processed_service_messages: Vec<ServiceMessage>,
	/// Messages that were sent by this guest to some other guests.
	pub outgoing_messages: Vec<(GuestId, Vec<u8>)>,
}

/// The specification of the VM's output data.
///
/// Includes the information that is needed to decode work package exports and also the final state
/// after the VM invocation (remaining gas and the outcome).
///
/// Exports have the following structure:
///
/// `[ memory pages | output stream 0 | ... | output stream N | null segments ]`
#[derive(Encode, Decode, MaxEncodedLen, Debug)]
pub struct VmOutput {
	/// Remainig inner VM gas.
	pub remaining_gas: SignedGas,
	/// The outcome of the inner VM invocation.
	pub outcome: Outcome,
	/// The no. of exported memory pages.
	pub num_memory_pages: u32,
	/// Per-stream output size.
	pub stream_len: [u32; OutputStream::COUNT],
}

impl VmOutput {
	/// Get the size of the stream `i`.
	pub fn stream_len(&self, i: OutputStream) -> u32 {
		self.stream_len[i as usize - 1]
	}

	/// Total length of all streams.
	pub fn all_streams_len(&self) -> u32 {
		self.stream_len.iter().sum()
	}

	/// Returns `true` if all streams are empty.
	pub fn is_empty(&self) -> bool {
		self.stream_len.iter().all(|len| *len == 0)
	}

	/// Get the number of output segments.
	pub fn num_output_segments(&self) -> u32 {
		self.stream_len.iter().sum::<u32>().div_ceil(SEGMENT_LEN as u32)
	}

	/// Get the start and end offset of the specified stream.
	///
	/// The offsets are measured from the first exported segment.
	pub fn get_stream_range(&self, stream: OutputStream) -> (u32, u32) {
		let mut start = 0;
		let mut end = 0;
		for i in OutputStream::ALL {
			let len = self.stream_len[i as usize - 1];
			if i == stream {
				end = start + len;
				break;
			}
			start += len;
		}
		let offset = self.num_memory_pages * SEGMENT_LEN as u32;
		(offset + start, offset + end)
	}
}

/// The result code returned by CoreVM service.
///
/// It is similar to `jam_pvm_common::InvokeOutcome` but does not include host-call
/// faults because they are automatically handled by CoreVM.
#[derive(Encode, Decode, MaxEncodedLen, Debug, PartialEq, Eq, Clone, Copy)]
pub enum Outcome {
	/// Completed normally.
	Halt,
	/// Completed with a panic.
	Panic,
	/// Completed with a page fault that couldn't be handled by CoreVM.
	///
	/// Undhandled page fault might occur because _either_ the max. no. of exports is reached _or_
	/// the corresponding page was not imported by the work package. In the former case the
	/// program needs to be resumed in the next work package without any further action from the
	/// builder. In the latter case the page with the specified address has to be imported on the
	/// next program run to continue.
	PageFault { page: PageNum, num_pages: u32 },
	/// Completed by running out of gas.
	OutOfGas,
	/// The inner VM has reached the max. no. of exports while appending new data to the output
	/// stream.
	OutputLimitReached,
	/// The program produced one time-slot worth of video frames or audio samples.
	///
	/// If both streams are active, then both streams need to have one time-slot worth of data to
	/// trigger this outcome.
	TimeLimitReached,
	/// The program is waiting for the specified input streams.
	WaitingForInput(InputStreams),
	/// The inner VM has reached max. bundle size while reading input stream or host/inter-service
	/// message or on page fault.
	InputLimitReached,
}

/// Inner VM state.
#[derive(Encode, Decode, Debug, Clone)]
pub struct VmState {
	/// Registers.
	pub regs: [u64; 13],
	/// Program counter.
	pub program_counter: u64,
	/// Current step.
	pub step: u32,
	/// Mapped heap memory pages' indices.
	///
	/// These pages were allocated by the guest but not necessarily read from/written to.
	///
	/// Note that it's possible to include only a subset of the pages if you're sure that they will
	/// not be touched again during the program run. Howeve,r if such a page is touched the engine
	/// will panic.
	pub mapped_heap_pages: RangeSet,
	/// All memory pages that were read from/written to by the guest but haven't been unmapped
	/// yet.
	///
	/// Note that it's possible to include only a subset of the pages if you're sure that they will
	/// not be touched again during the program run. However, if such a page is touched the
	/// engine will silently initialize it from the default: from zeroes for stack/heap and from
	/// the program's RW data section for RW data. The program might continue without an error
	/// after that, but the output will most likely be incorrect.
	pub resident_pages: RangeSet,
	/// The state of the Linux kernel syscall layer.
	pub kernel: KernelState,
	/// The index of the host-call that needs to be restarted on resuming the program execution.
	#[doc(hidden)]
	pub restart_host_call: Option<u64>,
	/// Video output mode.
	pub video: Option<VideoMode>,
	/// Audio output mode.
	pub audio: Option<AudioMode>,
}

impl VmState {
	/// Create initial state.
	pub const fn initial() -> Self {
		Self {
			regs: [0; 13],
			program_counter: 0,
			mapped_heap_pages: RangeSet::new(),
			resident_pages: RangeSet::new(),
			kernel: KernelState { fds: VecMap::new() },
			restart_host_call: None,
			video: None,
			audio: None,
			step: 0,
		}
	}
}

/// This macro provides `Self::COUNT_WITHOUT_PARAMS` constant that specifies the total number of
/// variants without parameters.
macro_rules! enum_with_count {
    (
    $(#[$($attr:tt)*])*
    pub enum $enum:ident {
        $(
            $(#[$($variant_attr:tt)*])*
            $variant:ident$(($($field:ident),+))?,
        )*
    }
    ) => {
        $(#[$($attr)*])*
        pub enum $enum {
            $(
                $(#[$($variant_attr)*])*
                $variant$(($($field),+))?,
            )*
        }

        impl $enum {
            /// Total number of all variants without parameters.
            #[doc(hidden)]
            pub const COUNT_WITHOUT_PARAMS: usize = {
                const fn one(_: &'static str) -> usize { 1 }
                const fn one_if_has_params(_: &'static str, num_params: usize) -> usize {
                    if num_params == 0 { 1 } else { 0 }
                }
                0
                $(
                    + one_if_has_params(
                        stringify!($variant),
                        0 $(+ one(concat!($(stringify!($field),)+)))?
                    )
                )*
            };
        }
    };
}

enum_with_count! {
	/// Storage keys used by CoreVM.
	#[derive(Encode, Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
	#[non_exhaustive]
	pub enum StorageKey {
		/// Inner VM gas.
		Gas,
		/// The hash of this CoreVM instance.
		///
		/// Computed with `corevm_engine::compute_state_hash`.
		StateHash,
		/// The last known metadata of a memory page with the specified page number.
		///
		/// An instance of [`PageInfo`](crate::PageInfo).
		PageInfo(PageNum),
		/// The output data specification of the VM.
		///
		/// An instance of [`VmSpec`](crate::VmSpec).
		VmSpec,
		/// The current video output mode.
		///
		/// An instance of [`VideoMode`](crate::VideoMode).
		VideoMode,
		/// The current audio output mode.
		///
		/// An instance of [`AudioMode`](crate::AudioMode).
		AudioMode,
		/// Execution environment location in the virtual file system.
		///
		/// The environment ([`ExecEnv`]) itself is stored as a file.
		ExecEnvRef,
		/// The owner of this CoreVM instance; they can reset the code.
		Owner,
		/// The pages that are currently in the storage.
		///
		/// Used internally by `accumulate` to verify `refine` output.
		#[doc(hidden)]
		StoredPages,
		/// Incoming inter-service messages.
		///
		/// An instance of `VecDeque<ServiceMessage>`.
		IncomingServiceMessages,
		/// Outgoing inter-service messages.
		///
		/// An instance of [`MessageQueue<OutgoingServiceMessage>`](crate::MessageQueue).
		OutgoingServiceMessages,
		/// Outgoing inter-service message with the specified index.
		OutgoingMessage(u64),
		/// Incoming inter-service message.
		IncomingMessage(ServiceMessage),
	}
}

/// Memory page metadata.
#[derive(Encode, Decode, Debug)]
pub struct PageInfo {
	/// The hash of the page segment.
	pub hash: Hash,
	/// Segment-root of the package exports where the page contents are stored.
	pub exports_root: SegmentTreeRoot,
	/// The index of the exported segment where the page contents are stored.
	pub export_index: u16,
}

/// Same as [`CoreVmOutput`] but only includes the data necessary to continue program execution.
#[derive(Encode, Decode, Debug)]
pub struct VmSpec {
	/// Segment-root of the package exports where the actual output is stored.
	pub exports_root: SegmentTreeRoot,
	/// The specification of the VM's output data.
	pub output: VmOutput,
	pub state: VmState,
}

/// Guest execution environment.
#[derive(Encode, Decode, Clone, Debug)]
pub struct ExecEnv {
	/// Program location in the virtual file system.
	pub program: fs::BlockRef,
	/// Root directory location in the virtual file system.
	pub root_dir: fs::BlockRef,
	// TODO @ivan add current working directory
	/// Command-line arguments.
	pub args: Vec<Arg>,
	/// Environment variables.
	pub env: Vec<Arg>,
	/// Video input mode.
	pub video_input: Option<VideoMode>,
	/// Audio input mode.
	pub audio_input: Option<AudioMode>,
	/// Ed25519 public key that is used to verify input chunks' signatures.
	pub input_key: Option<[u8; 32]>,
}

/// Instruction to send to the CoreVM service.
///
/// This should fit into [`Memo`](jam_types::Memo) in the encoded form.
#[derive(Encode, Decode, Debug)]
pub enum CoreVmInstruction {
	/// Reset the VM.
	Reset {
		/// Inner VM gas.
		gas: SignedGas,
		/// Reference to a file where the environment ([`ExecEnv`]) is stored.
		exec_ref: fs::BlockRef,
	},
	/// Set the owner of this CoreVM instance.
	SetOwner(ServiceId),
	/// Push a message from another service to the current service's message queue.
	PushServiceMessage(ServiceMessage),
	/// Destroy the VM.
	///
	/// This instruction clears the storage and zombifies the service.
	Destroy {
		/// A service that should eject the CoreVM service after it is destroyed.
		ejector: ServiceId,
	},
}

impl MaxEncodedLen for CoreVmInstruction {
	fn max_encoded_len() -> usize {
		MEMO_LEN
	}
}

/// The state of the Linux kernel syscall layer.
#[derive(Encode, Decode, Clone, Debug, Default)]
pub struct KernelState {
	/// Opened file descriptors.
	pub fds: VecMap<u32, KernelFd>,
}

/// In-kernel file descriptor state.
#[derive(Encode, Decode, MaxEncodedLen, Clone, Debug)]
pub struct KernelFd {
	/// File reference in the virtual FS.
	pub block_ref: fs::BlockRef,
	/// File read/write offset.
	pub position: u64,
}

impl ConstEncodedLen for KernelFd {}