corevm_engine/

output.rs

extern crate alloc;

use alloc::vec::Vec;
use corevm_host::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, Debug)]
pub struct OutputBuffers {
	buffers: [Vec<u8>; OutputStream::COUNT],
}

impl OutputBuffers {
	/// Pre-allocate `len` bytes in output buffer `i`.
	///
	/// Returns the mutable slice that references the pre-allocated space.
	#[must_use]
	pub fn pre_allocate(&mut self, i: OutputStream, len: usize) -> &mut [u8] {
		if len == 0 {
			return &mut [];
		}
		let buf = self.get_mut(i);
		let offset = buf.len();
		buf.resize(offset + len, 0_u8);
		&mut buf[offset..]
	}

	/// Export all buffers as segments.
	///
	/// Normally `export` should simply call [`export_slice`](jam_pvm_common::export_slice).
	///
	/// Use `Self::segment_count_after` to ensure that max. no. of exported segments is not
	/// surpassed.
	pub fn export_segments<F, E>(&mut self, mut export: F) -> Result<(), E>
	where
		F: FnMut(&RawSegment) -> Result<(), E>,
	{
		let mut output = SegmentedOutput::new();
		for i in OutputStream::ALL {
			output.write(self.get(i), &mut export)?;
		}
		output.finish(&mut export)?;
		Ok(())
	}

	/// Returns the size of each stream.
	pub fn stream_len(&self) -> [u32; OutputStream::COUNT] {
		[
			self.buffers[0].len() as u32,
			self.buffers[1].len() as u32,
			self.buffers[2].len() as u32,
			self.buffers[3].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 no. of segments the buffers will require on export after `len` bytes is appended
	/// to buffer `i`.
	pub fn segment_count_after(&self, i: OutputStream, len: usize) -> Option<usize> {
		self.total_len_after(i, len).map(|len| 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.buffers.iter().map(|buf| buf.len()).sum()
	}

	/// Returns the total encoded length of the buffers after `len` bytes is appended to buffer `i`.
	fn total_len_after(&self, i: OutputStream, len: usize) -> Option<usize> {
		let buf = self.get(i);
		let old_len = buf.len();
		let new_len = old_len.checked_add(len)?;
		(self.total_len() - old_len).checked_add(new_len)
	}

	/// Get immutable slice of buffer `i`.
	fn get(&self, i: OutputStream) -> &[u8] {
		&self.buffers[i as usize - 1]
	}

	/// Get mutable slice of buffer `i`.
	fn get_mut(&mut self, i: OutputStream) -> &mut Vec<u8> {
		&mut self.buffers[i as usize - 1]
	}
}

/// 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 {
	segment: RawSegment,
	offset: usize,
}

impl SegmentedOutput {
	fn new() -> Self {
		Self { segment: [0; SEGMENT_LEN], offset: 0 }
	}

	fn write<F, E>(&mut self, mut bytes: &[u8], export: &mut F) -> Result<(), E>
	where
		F: FnMut(&RawSegment) -> Result<(), E>,
	{
		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 {
				export(&self.segment)?;
				self.offset = 0;
			}
		}
		Ok(())
	}

	fn finish<F, E>(&mut self, export: &mut F) -> Result<(), E>
	where
		F: FnMut(&RawSegment) -> 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;
		}
		export(&self.segment)?;
		self.offset = 0;
		Ok(())
	}
}

type RawSegment = [u8; SEGMENT_LEN];

#[cfg(test)]
mod tests {
	use super::*;
	use core::convert::Infallible;
	use jam_types::Segment;
	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();
			let mut stream_len = [0; OutputStream::COUNT];
			for stream in OutputStream::ALL {
				let num_bytes: u32 = rng.random_range(0..=100);
				let slice = guest_buffers.pre_allocate(stream, num_bytes as usize);
				rng.fill_bytes(slice);
				stream_len[stream as usize - 1] = num_bytes;
			}
			let mut segments: Vec<Segment> = Vec::new();
			guest_buffers
				.export_segments(|s| {
					segments.push(s.to_vec().try_into().unwrap());
					Ok::<(), Infallible>(())
				})
				.unwrap();
			let mut host_buffers = corevm_host::OutputBuffers::default();
			host_buffers.append_from_segments(&segments[..], &stream_len);
			for i in OutputStream::ALL {
				let guest_bytes = guest_buffers.get(i);
				assert_eq!(guest_bytes, host_buffers.get(i));
			}
		}
	}

	#[test]
	fn segment_count_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 expected_total_len = guest_buffers.total_len_after(stream, num_bytes);
				let expected_segment_count = guest_buffers.segment_count_after(stream, num_bytes);
				let slice = guest_buffers.pre_allocate(stream, num_bytes);
				rng.fill_bytes(slice);
				assert_eq!(expected_total_len, Some(guest_buffers.total_len()));
				let mut actual_segment_count = 0;
				guest_buffers
					.export_segments(|_| {
						actual_segment_count += 1;
						Ok::<(), Infallible>(())
					})
					.unwrap();
				assert_eq!(
					expected_segment_count,
					Some(actual_segment_count),
					"expected segment count = {expected_segment_count:?}, \
                    actual segment count = {actual_segment_count}, \
                    total len = {expected_total_len:?}"
				);
			}
		}
	}
}