corevm-engine 0.1.28

//! These structs help count how many input chunks were read from the extrinsics.

use alloc::collections::VecDeque;
use bytes::{Buf as _, Bytes};

/// This macro counts the current chunk as read if it hasn't already been counted.
///
/// You should invoke this macro each time a chunk is peeked or extracted from the queue.
/// Peeking is counted because its outcome depends on the total number of chunks which might be
/// different in the builder and in the JAM node. See [`ChunkedBytesInput`] for further explanation.
///
/// This could have been a method if Rust borrow checker allowed us.
macro_rules! count_chunk {
	($self: ident) => {
		if !$self.chunk_counted {
			$self.num_chunks_read += 1;
			$self.chunk_counted = true;
		}
	};
}

/// Input stream that uses a collection of [`Bytes`] as the underlying data source and counts how
/// many chunks were read.
///
/// The purpose of this struct is to reliably count how many chunks were read by the guest. The
/// counting is done in the simulator which receives all input chunks regardless of the input size,
/// and the resulting number is then used by the builder to truncate the input chunks so that they
/// don't overflow the max. input size of the work package.
///
/// # Warning ⚠️
///
/// We shouldn't expose any fields that depend on the number of chunks because this number can be
/// different when the guest is executed by the builder and by the JAM node!
pub struct ChunkedBytesInput {
	chunks: VecDeque<Bytes>,
	position: usize,
	num_chunks_read: usize,
	chunk_counted: bool,
}

impl ChunkedBytesInput {
	pub fn new(chunks: VecDeque<Bytes>) -> Self {
		Self { chunks, position: 0, num_chunks_read: 0, chunk_counted: false }
	}

	/// Returns the number of chunks that were read/accessed.
	pub fn num_chunks_read(&self) -> usize {
		// This is safe to expose: it doesn't depend on the total number of chunks.
		self.num_chunks_read
	}

	/// Returns the number of bytes read so far.
	pub fn position(&self) -> usize {
		// This is safe to expose: it doesn't depend on the total number of chunks.
		self.position
	}

	/// Peeks the data available in the current chunk while skipping empty chunks.
	///
	/// Returns `None` if there are no more chunks or if all the remaining chunks are empty.
	///
	/// - Skips empty chunks.
	/// - Returns the data by reference via [`Bytes`].
	///
	/// # Note
	///
	/// This method might increase [`num_chunks_read`](Self::num_chunks_read) even if it returns
	/// `None`. This is because empty chunks are skipped.
	pub fn peek_within_chunk_skipping_empty_chunks(&mut self) -> Option<Bytes> {
		while let Some(chunk) = self.chunks.front_mut() {
			count_chunk!(self);
			if chunk.is_empty() {
				self.chunks.pop_front();
				self.chunk_counted = false;
				continue;
			}
			return Some(chunk.clone());
		}
		None
	}

	/// Reads up to `length` bytes of the remaining data from the current chunk while skipping
	/// empty chunks.
	///
	/// Returns `None` if there are no more chunks or if all the remaining chunks are empty.
	///
	/// This method is similar to [`std::io::Read::read`], i.e. it might not return all the
	/// remaining bytes in the stream.
	///
	/// - Skips empty chunks.
	/// - Returns the data by reference via [`Bytes`].
	///
	/// # Note
	///
	/// This method might increase [`num_chunks_read`](Self::num_chunks_read) even if it returns
	/// `None`. This is because empty chunks are skipped.
	pub fn read_within_chunk_skipping_empty_chunks(&mut self, length: usize) -> Option<Bytes> {
		while let Some(chunk) = self.chunks.front_mut() {
			count_chunk!(self);
			if chunk.is_empty() {
				self.chunks.pop_front();
				self.chunk_counted = false;
				continue;
			}
			let n = length.min(chunk.len());
			self.position += n;
			return Some(chunk.split_to(n));
		}
		None
	}
}

impl codec::Input for ChunkedBytesInput {
	fn remaining_len(&mut self) -> Result<Option<usize>, codec::Error> {
		// NOTE Can't use total length of all chunks here because the number of chunks can be
		// different when the guest is executed by CoreVM builder and by JAM node.
		Ok(None)
	}

	fn read(&mut self, mut buf: &mut [u8]) -> Result<(), codec::Error> {
		if buf.is_empty() {
			return Ok(());
		}
		while let Some(chunk) = self.chunks.front_mut() {
			count_chunk!(self);
			if chunk.is_empty() {
				self.chunks.pop_front();
				self.chunk_counted = false;
				continue;
			}
			let n = buf.len().min(chunk.len());
			let (part, rest) = buf.split_at_mut(n);
			part.copy_from_slice(&chunk[..n]);
			buf = rest;
			self.position += n;
			chunk.advance(n);
			if buf.is_empty() {
				return Ok(());
			}
		}
		if !buf.is_empty() {
			return Err("Buffer overflow".into());
		}
		Ok(())
	}
}

/// Chunked queue of bytes that tracks how many chunks were read.
///
/// See [`ChunkedBytesInput`] for the explanation of why this is needed.
///
/// # ⚠️ Warning ⚠️
///
/// We shouldn't expose any fields that depend on the number of chunks because this number can be
/// different when the guest is executed by the builder and by the JAM node!
pub struct ChunkedVecDequeBytes {
	chunks: VecDeque<VecDeque<Bytes>>,
	num_chunks_read: usize,
	chunk_counted: bool,
}

impl ChunkedVecDequeBytes {
	pub fn new(chunks: VecDeque<VecDeque<Bytes>>) -> Self {
		Self { chunks, num_chunks_read: 0, chunk_counted: false }
	}

	/// Returns the number of chunks that were read/accessed.
	pub fn num_chunks_read(&self) -> usize {
		// This is safe to expose: it doesn't depend on the total number of chunks.
		self.num_chunks_read
	}

	/// Peek the front chunk while skipping empty chunks.
	pub fn peek_front_skipping_empty_chunks(&mut self) -> Option<Bytes> {
		self.pop_empty_chunks();
		match self.chunks.front_mut() {
			Some(chunk) => {
				count_chunk!(self);
				chunk.front().cloned()
			},
			None => None,
		}
	}

	/// Pop the front chunk while skipping empty chunks.
	pub fn pop_front_skipping_empty_chunks(&mut self) -> Option<Bytes> {
		self.pop_empty_chunks();
		match self.chunks.front_mut() {
			Some(chunk) => {
				count_chunk!(self);
				chunk.pop_front()
			},
			None => None,
		}
	}

	fn pop_empty_chunks(&mut self) {
		while let Some(chunk) = self.chunks.front() {
			count_chunk!(self);
			if !chunk.is_empty() {
				break;
			}
			self.chunks.pop_front();
			self.chunk_counted = false;
		}
	}
}