corevm-host 0.1.28

use alloc::{borrow::Cow, vec::Vec};
use codec::{Compact, CompactLen, ConstEncodedLen, CountedInput, Decode, Encode, MaxEncodedLen};
use jam_types::SEGMENT_LEN;

/// Output stream identifier.
#[derive(Encode, Decode, MaxEncodedLen, Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum OutputStream {
	Stdout = 0,
	Stderr = 1,
	Video = 2,
	Audio = 3,
}

impl OutputStream {
	/// Total no. of output streams.
	pub const COUNT: usize = 4;

	/// All output streams.
	pub const ALL: [OutputStream; Self::COUNT] = {
		use OutputStream::*;
		[Stdout, Stderr, Video, Audio]
	};
}

impl ConstEncodedLen for OutputStream {}

/// A chunk of console output.
#[derive(Debug)]
pub struct ConsoleChunk {
	/// Offset from the start of the slot in milliseconds.
	pub time_offset: u64,
	pub buf: Vec<u8>,
}

/// The size of the chunk encoded as a part of [`ConsoleChunks`] (with delta encoding).
fn delta_encoded_size(time_offset: u64, len: usize, prev_time_offset: u64) -> usize {
	Compact::compact_len(&time_offset.wrapping_sub(prev_time_offset)) +
		Compact::compact_len(&(len as u64)) +
		len
}

/// Several chunks of console output.
#[derive(Debug)]
pub struct ConsoleChunks {
	chunks: Vec<ConsoleChunk>,
	/// Encoded size without the chunks vector length.
	///
	/// This field is computed automatically.
	encoded_size: usize,
}

impl ConsoleChunks {
	pub const fn new() -> Self {
		Self { chunks: Vec::new(), encoded_size: 0 }
	}

	pub fn clear(&mut self) {
		self.chunks.clear();
		self.encoded_size = 0;
	}

	pub fn into_inner(self) -> Vec<ConsoleChunk> {
		self.chunks
	}

	/// Append new console buffer.
	///
	/// If the supplied time offset matches the one of the last chunk, then this chunk is
	/// extended; otherwise a new chunk with the matching time offset is appended.
	pub fn append(&mut self, time_offset: u64, buf: Cow<'_, [u8]>) {
		match self.chunks.last_mut() {
			Some(chunk) if chunk.time_offset == time_offset => {
				self.encoded_size -= Compact::compact_len(&(chunk.buf.len() as u64));
				chunk.buf.extend_from_slice(buf.as_ref());
				self.encoded_size += Compact::compact_len(&(chunk.buf.len() as u64));
				self.encoded_size += buf.len();
			},
			_ => self.push(ConsoleChunk { time_offset, buf: buf.into_owned() }),
		}
	}

	/// Pre-allocate `len` bytes in the last chunk.
	///
	/// Appends a new chunk if there are no chunks or if the time offset of the last chunk doesn't
	/// match the supplied time offset.
	///
	/// Panics on length overflow.
	#[must_use]
	pub fn pre_allocate(&mut self, time_offset: u64, len: usize) -> &mut [u8] {
		match self.chunks.last_mut() {
			Some(chunk) if chunk.time_offset == time_offset => {},
			_ => self.push(ConsoleChunk { time_offset, buf: Vec::new() }),
		}
		let buf = &mut self.chunks.last_mut().expect("Initialized above").buf;
		self.encoded_size -= Compact::compact_len(&(buf.len() as u64));
		let offset = buf.len();
		buf.resize(offset + len, 0_u8);
		self.encoded_size += Compact::compact_len(&(buf.len() as u64));
		self.encoded_size += len;
		&mut buf[offset..]
	}

	fn push(&mut self, chunk: ConsoleChunk) {
		let prev_time_offset = self
			.chunks
			.last()
			.map(|last_chunk| {
				debug_assert!(
					last_chunk.time_offset < chunk.time_offset,
					"Console chunks are not monotonic: last chunk = {}, new chunk = {}",
					last_chunk.time_offset,
					chunk.time_offset
				);
				last_chunk.time_offset
			})
			.unwrap_or(0);
		self.encoded_size +=
			delta_encoded_size(chunk.time_offset, chunk.buf.len(), prev_time_offset);
		self.chunks.push(chunk);
	}

	/// Returns the encoded size after a chunk with the supplied offset and having `len` bytes is
	/// appended.
	///
	/// Returns `None` on overflow.
	pub fn encoded_size_after(&self, time_offset: u64, len: usize) -> Option<usize> {
		let mut encoded_size = self.encoded_size;
		let mut chunks_len = self.chunks.len();
		match self.chunks.last() {
			Some(chunk) if chunk.time_offset == time_offset => {
				encoded_size -= Compact::compact_len(&(chunk.buf.len() as u64));
				encoded_size += Compact::compact_len(&(chunk.buf.len().checked_add(len)? as u64));
			},
			last => {
				let prev_time_offset = last.map(|chunk| chunk.time_offset).unwrap_or(0);
				if time_offset < prev_time_offset {
					// Non-monotonic.
					return None;
				}
				encoded_size += Compact::compact_len(&time_offset.wrapping_sub(prev_time_offset));
				encoded_size += Compact::compact_len(&(len as u64));
				chunks_len += 1;
			},
		}
		// Plus encoded size of the chunks vector length.
		Some(encoded_size.checked_add(len)? + Compact::compact_len(&(chunks_len as u64)))
	}
}

impl Default for ConsoleChunks {
	fn default() -> Self {
		Self::new()
	}
}

impl core::ops::Deref for ConsoleChunks {
	type Target = [ConsoleChunk];

	fn deref(&self) -> &Self::Target {
		&self.chunks[..]
	}
}

impl<'a> FromIterator<(u64, Cow<'a, [u8]>)> for ConsoleChunks {
	fn from_iter<I: IntoIterator<Item = (u64, Cow<'a, [u8]>)>>(iter: I) -> Self {
		let mut chunks = Self::new();
		chunks.extend(iter);
		chunks
	}
}

impl<'a> Extend<(u64, Cow<'a, [u8]>)> for ConsoleChunks {
	fn extend<I: IntoIterator<Item = (u64, Cow<'a, [u8]>)>>(&mut self, chunks: I) {
		for (time_offset, buf) in chunks.into_iter() {
			self.append(time_offset, buf);
		}
	}
}

impl FromIterator<ConsoleChunk> for ConsoleChunks {
	fn from_iter<I: IntoIterator<Item = ConsoleChunk>>(iter: I) -> Self {
		let mut chunks = Self::new();
		chunks.extend(iter);
		chunks
	}
}

impl Extend<ConsoleChunk> for ConsoleChunks {
	fn extend<I: IntoIterator<Item = ConsoleChunk>>(&mut self, chunks: I) {
		for ConsoleChunk { time_offset, buf } in chunks.into_iter() {
			self.append(time_offset, buf.into());
		}
	}
}

impl<'a> FromIterator<&'a ConsoleChunk> for ConsoleChunks {
	fn from_iter<I: IntoIterator<Item = &'a ConsoleChunk>>(iter: I) -> Self {
		let mut chunks = Self::new();
		chunks.extend(iter);
		chunks
	}
}

impl<'a> Extend<&'a ConsoleChunk> for ConsoleChunks {
	fn extend<I: IntoIterator<Item = &'a ConsoleChunk>>(&mut self, chunks: I) {
		for ConsoleChunk { time_offset, buf } in chunks.into_iter() {
			self.append(*time_offset, buf.into());
		}
	}
}

impl Encode for ConsoleChunks {
	fn encode_to<T: codec::Output + ?Sized>(&self, output: &mut T) {
		// Use delta encoding to reduce the encoded size of the time offsets.
		let mut prev = 0;
		Compact(self.chunks.len() as u64).encode_to(output);
		for ConsoleChunk { time_offset, buf } in self.chunks.iter() {
			Compact(time_offset.wrapping_sub(prev)).encode_to(output);
			buf.encode_to(output);
			prev = *time_offset;
		}
	}

	fn encoded_size(&self) -> usize {
		// Plus encoded size of the chunks vector length.
		self.encoded_size + Compact::compact_len(&(self.chunks.len() as u64))
	}
}

impl Decode for ConsoleChunks {
	fn decode<I: codec::Input>(input: &mut I) -> Result<Self, codec::Error> {
		let mut input = CountedInput::new(input);
		let len = Compact::<u64>::decode(&mut input)?.0 as usize;
		let mut chunks = Vec::with_capacity(len);
		let mut prev = 0;
		for _ in 0..len {
			let time_offset = Compact::<u64>::decode(&mut input)?.0.wrapping_add(prev);
			let buf: Vec<u8> = Decode::decode(&mut input)?;
			chunks.push(ConsoleChunk { time_offset, buf });
			prev = time_offset;
		}
		// Minus encoded size of the chunks vector length.
		let encoded_size = input.count() as usize - Compact::compact_len(&(chunks.len() as u64));
		Ok(Self { chunks, encoded_size })
	}
}

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

impl OutputBuffers {
	/// Create output buffers from the output streams stored in `segments`.
	///
	/// Each stream size is specified in `stream_len` array.
	///
	/// The data should start at the beginning of the first segment and end somewhere in the last
	/// segment. The total length equals the sum of lengths in `stream_len` array.
	pub fn from_segments(
		segments: &[impl AsRef<[u8; SEGMENT_LEN]>],
		stream_len: &[u32; OutputStream::COUNT],
	) -> Result<Self, codec::Error> {
		let total_len: usize = stream_len.iter().map(|len| *len as usize).sum();
		let mut input = SegmentedInput::new(segments, 0, total_len);
		let mut buffers = Self::default();
		for (src_len, stream) in stream_len.iter().zip(OutputStream::ALL) {
			if *src_len == 0 {
				continue;
			}
			let src_len = *src_len as usize;
			debug_assert!(
				src_len <= input.end - input.offset,
				"src len = {src_len}, remaining len = {}",
				input.end - input.offset
			);
			buffers.set_stream(stream, src_len, &mut input)?;
		}
		Ok(buffers)
	}

	/// Create output buffers from the specified output stream stored in `segments`.
	pub fn from_segments_one(
		segments: &[impl AsRef<[u8; SEGMENT_LEN]>],
		stream: OutputStream,
		stream_start: usize,
		stream_end: usize,
	) -> Result<Self, codec::Error> {
		let src_len = stream_end - stream_start;
		let mut input = SegmentedInput::new(segments, stream_start, stream_end);
		let mut buffers = Self::default();
		debug_assert!(
			src_len <= input.end - input.offset,
			"src len = {src_len}, remaining len = {}",
			input.end - input.offset
		);
		buffers.set_stream(stream, src_len, &mut input)?;
		Ok(buffers)
	}

	fn set_stream(
		&mut self,
		i: OutputStream,
		src_len: usize,
		input: &mut SegmentedInput<'_, impl AsRef<[u8; SEGMENT_LEN]>>,
	) -> Result<(), codec::Error> {
		use OutputStream::*;
		match i {
			Stdout | Stderr => {
				let old_offset = input.offset;
				let chunks = ConsoleChunks::decode(input)?;
				if input.offset - old_offset != src_len {
					return Err("Invalid console chunks size".into());
				}
				self.console_buffers[i as usize] = chunks;
			},
			Video | Audio => {
				let dst = &mut self.buffers[i as usize - 2];
				let dst_offset = dst.len();
				dst.resize(dst_offset + src_len, 0_u8);
				input.read(&mut dst[dst_offset..]);
			},
		}
		Ok(())
	}

	/// Get decoded stdout/stderr buffer.
	pub fn get_console_buf(&self, i: OutputStream) -> &[ConsoleChunk] {
		assert!(matches!(i, OutputStream::Stdout | OutputStream::Stderr));
		&self.console_buffers[i as usize][..]
	}

	/// Take decoded stdout/stderr buffer.
	pub fn take_console_buf(&mut self, i: OutputStream) -> Vec<ConsoleChunk> {
		assert!(matches!(i, OutputStream::Stdout | OutputStream::Stderr));
		core::mem::take(&mut self.console_buffers[i as usize]).into_inner()
	}

	/// Get encoded video/audio buffer.
	pub fn get_encoded_buf(&self, i: OutputStream) -> &[u8] {
		assert!(matches!(i, OutputStream::Video | OutputStream::Audio));
		&self.buffers[i as usize - 2]
	}

	/// Take encoded video/audio buffer.
	pub fn take_encoded_buf(&mut self, i: OutputStream) -> Vec<u8> {
		assert!(matches!(i, OutputStream::Video | OutputStream::Audio));
		core::mem::take(&mut self.buffers[i as usize - 2])
	}

	/// Clear all buffers.
	pub fn clear(&mut self) {
		for buf in self.console_buffers.iter_mut() {
			buf.clear();
		}
		for buf in self.buffers.iter_mut() {
			buf.clear();
		}
	}
}

/// A reader that reads directly from segments.
///
/// This avoids unnecessary copying potentially large amount of data.
#[derive(Debug)]
struct SegmentedInput<'a, S> {
	segments: &'a [S],
	end: usize,
	offset: usize,
}

impl<'a, S: AsRef<[u8; SEGMENT_LEN]>> SegmentedInput<'a, S> {
	fn new(segments: &'a [S], start: usize, end: usize) -> Self {
		Self { segments, end, offset: start }
	}

	fn read(&mut self, dest: &mut [u8]) {
		let dest_len = dest.len();
		let mut dest_offset = 0;
		while dest_offset != dest_len {
			let i = self.offset / SEGMENT_LEN;
			let src_offset = self.offset % SEGMENT_LEN;
			let n = (SEGMENT_LEN - src_offset).min(dest_len - dest_offset);
			dest[dest_offset..dest_offset + n]
				.copy_from_slice(&self.segments[i].as_ref()[src_offset..src_offset + n]);
			dest_offset += n;
			self.offset += n;
		}
	}
}

impl<S: AsRef<[u8; SEGMENT_LEN]>> codec::Input for SegmentedInput<'_, S> {
	fn remaining_len(&mut self) -> Result<Option<usize>, codec::Error> {
		Ok(Some(self.end - self.offset))
	}

	fn read(&mut self, dest: &mut [u8]) -> Result<(), codec::Error> {
		SegmentedInput::read(self, dest);
		Ok(())
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	use alloc::{vec, vec::Vec};
	use jam_types::SegmentBytes;
	use rand::Rng;

	#[test]
	fn segmented_input_works() {
		#[derive(Encode, Decode, Debug, PartialEq, Eq)]
		struct Dummy {
			x: u64,
			y: u32,
			z: Vec<Dummy>,
		}
		let mut rng = rand::rng();
		let num_values = rng.random_range(0..1000);
		let mut values: Vec<Dummy> = Vec::with_capacity(num_values);
		for _ in 0..num_values {
			let z_len = rng.random_range(0..10);
			let mut z = Vec::with_capacity(z_len);
			for _ in 0..z_len {
				z.push(Dummy { x: rng.random(), y: rng.random(), z: Vec::new() });
			}
			values.push(Dummy { x: rng.random(), y: rng.random(), z });
		}
		let mut encoded = values.encode();
		let encoded_len = encoded.len();
		while !encoded.len().is_multiple_of(SEGMENT_LEN) {
			encoded.push(0_u8);
		}
		let segments = encoded
			.chunks(SEGMENT_LEN)
			.map(|slice| slice.to_vec().try_into().unwrap())
			.collect::<Vec<SegmentBytes>>();
		let mut input = SegmentedInput::new(&segments, 0, encoded_len);
		let actual_values: Vec<Dummy> = Decode::decode(&mut input).unwrap();
		assert_eq!(values, actual_values);
	}

	#[test]
	fn console_chunks_encoded_size_works() {
		macro_rules! check {
			($chunks: expr) => {{
				let chunks = $chunks;
				assert_eq!(chunks.encode().len(), chunks.encoded_size(), "Chunks = {chunks:?}");
			}};
		}
		check!(ConsoleChunks::default());
		check!({
			let mut chunks = ConsoleChunks::new();
			chunks.push(ConsoleChunk { time_offset: 99999, buf: vec![0_u8; 123] });
			chunks
		});
		check!({
			let mut chunks = ConsoleChunks::new();
			chunks.push(ConsoleChunk { time_offset: 99999, buf: vec![0_u8; 0] });
			chunks
		});
		check!({
			let mut chunks = ConsoleChunks::new();
			chunks.push(ConsoleChunk { time_offset: 99999, buf: vec![0_u8; 123] });
			chunks.append(99999, vec![0_u8; 123].into());
			chunks
		});
		check!({
			let mut chunks = ConsoleChunks::new();
			chunks.push(ConsoleChunk { time_offset: 99999, buf: vec![0_u8; 123] });
			chunks.append(100_000, vec![0_u8; 123].into());
			chunks
		});
		check!({
			let mut chunks = ConsoleChunks::new();
			chunks.push(ConsoleChunk { time_offset: 99999, buf: vec![0_u8; 123] });
			let _ = chunks.pre_allocate(100_000, 123);
			chunks
		});
		check!({
			let mut chunks = ConsoleChunks::new();
			chunks.push(ConsoleChunk { time_offset: 99999, buf: vec![0_u8; 123] });
			let _ = chunks.pre_allocate(100_000, 123);
			let bytes = chunks.encode();
			ConsoleChunks::decode(&mut &bytes[..]).unwrap()
		});
	}

	#[test]
	fn console_chunks_encoded_size_after_works() {
		let mut chunks = ConsoleChunks::new();
		let expected_encoded_size = chunks.encoded_size_after(99999, 123).unwrap();
		chunks.push(ConsoleChunk { time_offset: 99999, buf: vec![0_u8; 123] });
		assert_eq!(expected_encoded_size, chunks.encode().len());
	}
}