corevm-engine 0.1.28

use alloc::vec::Vec;
use codec::{Encode, Output as _};
use corevm_host::{ConsoleChunks, OutputStream};
use jam_types::SEGMENT_LEN;

/// CoreVM output buffers (service counterpart).
///
/// The service uses this struct to collect output from the guest.
///
/// The [builder counterpart](corevm_host::OutputBuffers).
#[derive(Default)]
pub struct OutputBuffers {
	console_buffers: [ConsoleChunks; 2],
	buffers: [Vec<u8>; 2],
}

impl OutputBuffers {
	/// Pre-allocate `len` bytes in output buffer `i`.
	///
	/// Returns the mutable slice that references the pre-allocated space.
	///
	/// `time_offset` is the offset from the start of the slot in milliseconds. Only used for
	/// stdout/stderr buffers.
	#[must_use]
	pub fn pre_allocate(&mut self, time_offset: u64, i: OutputStream, len: usize) -> &mut [u8] {
		if len == 0 {
			return &mut [];
		}
		match i {
			OutputStream::Stdout | OutputStream::Stderr =>
				self.console_buffers[i as usize].pre_allocate(time_offset, len),
			_ => {
				let buf = &mut self.buffers[i as usize - 2];
				let offset = buf.len();
				buf.resize(offset + len, 0_u8);
				&mut buf[offset..]
			},
		}
	}

	/// Append `slice` to output stream `i`.
	///
	/// `time_offset` is the offset from the start of the slot in milliseconds. Only used for
	/// stdout/stderr buffers.
	pub fn append(&mut self, time_offset: u64, i: OutputStream, slice: &[u8]) {
		if slice.is_empty() {
			return;
		}
		match i {
			OutputStream::Stdout | OutputStream::Stderr =>
				self.console_buffers[i as usize].append(time_offset, slice.into()),
			_ => self.buffers[i as usize - 2].extend_from_slice(slice),
		}
	}

	/// Export all buffers as segments.
	///
	/// Normally `export` should simply call `jam_pvm_common::export_slice`.
	///
	/// Use `Self::new_segments_after` to ensure that max. no. of exported segments is not
	/// surpassed.
	pub fn export_segments<F, E>(&mut self, export: F) -> Result<(), E>
	where
		F: FnMut(&RawSegment) -> Result<(), E>,
	{
		let mut output = SegmentedOutput::new(export);
		for chunks in self.console_buffers.iter() {
			chunks.encode_to(&mut output);
			if let Some(e) = output.error {
				return Err(e);
			}
		}
		for buf in self.buffers.iter() {
			output.write(buf);
			if let Some(e) = output.error {
				return Err(e);
			}
		}
		output.finish()?;
		Ok(())
	}

	/// Returns the size of each stream.
	pub fn stream_len(&self) -> [u32; OutputStream::COUNT] {
		[
			self.console_buffers[0].encoded_size() as u32,
			self.console_buffers[1].encoded_size() as u32,
			self.buffers[0].len() as u32,
			self.buffers[1].len() as u32,
		]
	}

	/// Returns the no. of segments the buffers would require on export.
	pub fn segment_count(&self) -> usize {
		self.total_len().div_ceil(SEGMENT_LEN)
	}

	/// Returns the total length of the buffers.
	///
	/// Equals the no. of bytes the buffers will require on export.
	pub fn total_len(&self) -> usize {
		self.stream_len().iter().sum::<u32>() as usize
	}

	fn buffer_len(&self, i: OutputStream) -> usize {
		use OutputStream::*;
		match i {
			Stdout | Stderr => self.console_buffers[i as usize].encoded_size(),
			Video | Audio => self.buffers[i as usize - 2].len(),
		}
	}

	fn buffer_len_after(&self, time_offset: u64, i: OutputStream, len: usize) -> Option<usize> {
		use OutputStream::*;
		match i {
			Stdout | Stderr =>
				self.console_buffers[i as usize].encoded_size_after(time_offset, len),
			Video | Audio => self.buffers[i as usize - 2].len().checked_add(len),
		}
	}

	/// Returns the no. of additional segments the buffers will require on export after `len` bytes
	/// is appended to buffer `i`.
	pub fn new_segments_after(
		&self,
		time_offset: u64,
		i: OutputStream,
		len: usize,
	) -> Option<usize> {
		let old_len = self.buffer_len(i);
		let new_len = self.buffer_len_after(time_offset, i, len)?;
		let old_total_len = self.total_len();
		let old_segments = old_total_len.div_ceil(SEGMENT_LEN);
		let new_total_len = (old_total_len - old_len).checked_add(new_len)?;
		let new_segments = new_total_len.div_ceil(SEGMENT_LEN);
		Some(new_segments - old_segments)
	}
}

/// Writer that chunks output into segments and exports them on write.
///
/// Uses `export` lambda (normally [`export_slice`](jam_pvm_common::export_slice)) to export
/// segments.
///
/// Use `finish` to write the last segment if any.
struct SegmentedOutput<E, F: FnMut(&RawSegment) -> Result<(), E>> {
	segment: RawSegment,
	offset: usize,
	export: F,
	error: Option<E>,
}

impl<E, F: FnMut(&RawSegment) -> Result<(), E>> SegmentedOutput<E, F> {
	fn new(export: F) -> Self {
		Self { segment: [0; SEGMENT_LEN], offset: 0, export, error: None }
	}

	fn finish(&mut self) -> Result<(), E> {
		if self.offset == 0 {
			// Nothing to export.
			return Ok(());
		}
		// Zeroize the remaining data.
		for i in self.offset..SEGMENT_LEN {
			self.segment[i] = 0;
		}
		(self.export)(&self.segment)?;
		self.offset = 0;
		Ok(())
	}
}

impl<E, F: FnMut(&RawSegment) -> Result<(), E>> codec::Output for SegmentedOutput<E, F> {
	fn write(&mut self, mut bytes: &[u8]) {
		if self.error.is_some() {
			return;
		}
		while !bytes.is_empty() {
			let n = (SEGMENT_LEN - self.offset).min(bytes.len());
			self.segment[self.offset..self.offset + n].copy_from_slice(&bytes[..n]);
			self.offset += n;
			bytes = &bytes[n..];
			if self.offset == SEGMENT_LEN {
				if let Err(e) = (self.export)(&self.segment) {
					self.error = Some(e);
					break;
				}
				self.offset = 0;
			}
		}
	}
}

type RawSegment = [u8; SEGMENT_LEN];

#[cfg(test)]
mod tests {
	use super::*;
	use core::convert::Infallible;
	use corevm_host::ConsoleChunk;
	use jam_types::SegmentBytes;
	use rand::{seq::IndexedRandom, Rng, RngCore};

	#[test]
	fn output_buffers_work() {
		let mut rng = rand::rng();
		for _ in 0..1000 {
			let mut guest_buffers = OutputBuffers::default();
			for stream in OutputStream::ALL {
				let num_bytes: u32 = rng.random_range(0..=100);
				let slice = guest_buffers.pre_allocate(0, stream, num_bytes as usize);
				rng.fill_bytes(slice);
			}
			let stream_len = guest_buffers.stream_len();
			let mut segments: Vec<SegmentBytes> = Vec::new();
			guest_buffers
				.export_segments(|s| {
					segments.push(s.into());
					Ok::<(), Infallible>(())
				})
				.unwrap();
			let host_buffers =
				corevm_host::OutputBuffers::from_segments(&segments[..], &stream_len).unwrap();
			for i in [OutputStream::Video, OutputStream::Audio] {
				let guest_bytes = &guest_buffers.buffers[i as usize - 2];
				assert_eq!(guest_bytes, host_buffers.get_encoded_buf(i));
			}
			for i in [OutputStream::Stdout, OutputStream::Stderr] {
				let guest_bytes = &guest_buffers.console_buffers[i as usize];
				assert_eq!(
					guest_bytes
						.iter()
						.map(|ConsoleChunk { time_offset, buf }| (time_offset, buf))
						.collect::<Vec<_>>(),
					host_buffers
						.get_console_buf(i)
						.iter()
						.map(|ConsoleChunk { time_offset, buf }| (time_offset, buf))
						.collect::<Vec<_>>(),
				);
			}
		}
	}

	#[test]
	fn new_segments_after_works() {
		let mut rng = rand::rng();
		for _ in 0..1000 {
			let num_appends = rng.random_range(0..=20);
			let mut guest_buffers = OutputBuffers::default();
			for _ in 0..num_appends {
				let stream = *OutputStream::ALL.choose(&mut rng).unwrap();
				let num_bytes: usize = rng.random_range(0..=100);
				let segment_count = guest_buffers.segment_count();
				let expected_segment_count =
					segment_count + guest_buffers.new_segments_after(0, stream, num_bytes).unwrap();
				let slice = guest_buffers.pre_allocate(0, stream, num_bytes);
				rng.fill_bytes(slice);
				let mut actual_segment_count = 0;
				guest_buffers
					.export_segments(|_| {
						actual_segment_count += 1;
						Ok::<(), Infallible>(())
					})
					.unwrap();
				assert_eq!(
					expected_segment_count, actual_segment_count,
					"expected segment count = {expected_segment_count}, \
                    actual segment count = {actual_segment_count}"
				);
			}
		}
	}
}