corevm-guest 0.1.28

#![allow(clippy::missing_safety_doc)]

use crate::{c, flags::RecvTransaction, GuestId, RecvTransactionOutcome};
use core::mem::size_of_val;

/// Opaque transaction failure.
#[derive(Debug)]
pub struct TransactionFailed;

impl core::fmt::Display for TransactionFailed {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
		f.write_str("Transaction failed")
	}
}

trait RetToResult {
	/// Convert host-call return value to result.
	fn into_result(self) -> Result<(), TransactionFailed>;
}

impl RetToResult for u64 {
	fn into_result(self) -> Result<(), TransactionFailed> {
		if self == 0 {
			Ok(())
		} else {
			Err(TransactionFailed)
		}
	}
}

/// Initiates a transaction with `target` providing the `payload`.
///
/// Blocks until the target guest calls [`recv_transaction`].
///
/// Returns `Ok(())` if both `initiate_transaction` and `recv_transaction` call succeed. Returns
/// `Err(TransactionFailed)` otherwise.
pub fn initiate_transaction(target: GuestId, payload: &[u8]) -> Result<(), TransactionFailed> {
	let address = payload.as_ptr() as u64;
	let len = size_of_val(payload) as u64;
	unsafe { c::initiate_transaction(u64::from(target), address, len) }.into_result()
}

/// Receives a transaction from the transaction initiator.
///
/// Returns the payload and the guest ID of the initiator or `None` if there is no pending
/// transaction.
///
/// This is a non-blocking variant of [`recv_transaction`].
#[cfg(feature = "alloc")]
pub fn try_recv_transaction() -> Option<(alloc::vec::Vec<u8>, GuestId)> {
	let ret = unsafe { c::recv_transaction(0, 0, RecvTransaction::NON_BLOCKING.bits()) };
	let (len, source) = RecvTransactionOutcome::from_reg(ret).into_inner()?;
	let mut buf = alloc::vec::Vec::with_capacity(len as usize);
	unsafe {
		c::recv_transaction(
			buf.as_mut_ptr() as u64,
			u64::from(len),
			RecvTransaction::NON_BLOCKING.bits(),
		)
	};
	unsafe { buf.set_len(len as usize) };
	Some((buf, source))
}

/// Receives a transaction from the transaction initiator.
///
/// Returns the payload and the guest ID of the initiator.
///
/// Blocks until some transaction against this guest becomes pending.
/// For a non-blocking variant see [`try_recv_transaction`].
#[cfg(feature = "alloc")]
pub fn recv_transaction() -> (alloc::vec::Vec<u8>, GuestId) {
	let ret = unsafe { c::recv_transaction(0, 0, 0) };
	let (len, source) = RecvTransactionOutcome::from_reg(ret)
		.into_inner()
		.expect("The call is blocking");
	let mut buf = alloc::vec::Vec::with_capacity(len as usize);
	unsafe { c::recv_transaction(buf.as_mut_ptr() as u64, u64::from(len), 0) };
	unsafe { buf.set_len(len as usize) };
	(buf, source)
}

/// Receives a transaction from the transaction initiator.
///
/// This is a non-blocking variant of [`recv_transaction_into`].
///
/// # Return value
///
/// - Returns `None` if there is no pending transaction.
/// - Returns `Some(Err(n))` where `n` is the payload size if the buffer is too small.
/// - On success returns `Some(Ok((n, guest_id)))` where `n` is the payload size and `guest_id` is
///   the guest ID of the initiator.
#[inline]
pub fn try_recv_transaction_into(buf: &mut [u8]) -> Option<Result<(usize, GuestId), usize>> {
	let ret = unsafe {
		c::recv_transaction(
			buf.as_mut_ptr() as u64,
			buf.len() as u64,
			RecvTransaction::NON_BLOCKING.bits(),
		)
	};
	let (len, source) = RecvTransactionOutcome::from_reg(ret).into_inner()?;
	if len > buf.len() as u32 {
		return Some(Err(len as usize));
	}
	Some(Ok((len as usize, source)))
}

/// Receives a transaction from the transaction initiator.
///
/// Blocks until some transaction against this guest becomes pending.
/// For a non-blocking variant see [`try_recv_transaction_into`].
///
/// # Return value
///
/// - Returns `Err(n)` where `n` is the payload size if the buffer is too small.
/// - On success returns `Ok((n, guest_id))` where `n` is the payload size and `guest_id` is the
///   guest ID of the initiator.
#[inline]
pub fn recv_transaction_into(buf: &mut [u8]) -> Result<(usize, GuestId), usize> {
	let ret = unsafe { c::recv_transaction(buf.as_mut_ptr() as u64, buf.len() as u64, 0) };
	let (len, source) = RecvTransactionOutcome::from_reg(ret)
		.into_inner()
		.expect("The call is blocking");
	if len > buf.len() as u32 {
		return Err(len as usize);
	}
	Ok((len as usize, source))
}

/// Commits the transaction that was started by either [`initiate_transaction`] or
/// [`recv_transaction`] or any of non-blocking variants of these calls.
///
/// Blocks until the other guest either calls [`commit_transaction`], or [`rollback_transaction`],
/// or panics.
///
/// Panics if there is no active transaction.
#[inline]
pub fn commit_transaction() -> Result<(), TransactionFailed> {
	unsafe { c::commit_transaction() }.into_result()
}

/// Cancels the transaction that was started by either [`initiate_transaction`] or
/// [`recv_transaction`] or any of non-blocking variants of these calls.
#[inline]
pub fn rollback_transaction() {
	unsafe { c::rollback_transaction() }
}

/// Ends the transaction that was started by either [`initiate_transaction`] or
/// [`recv_transaction`] or any of non-blocking variants of these calls.
#[inline]
pub fn end_transaction() {
	unsafe { c::end_transaction() }
}

#[derive(Debug, PartialEq, Eq)]
enum TransactionState {
	/// Transaction was initiated or received.
	Initial,
	/// Transaction was committed or rolled back.
	CommittedOrRolledBack,
	/// Transaction ended.
	Ended,
}

/// Transaction wrapper that automatically rolls back and ends unfinished transaction once it goes
/// out of scope.
///
/// # Example
///
/// Initiator:
/// ```rust,ignore
/// let mut tx = Transaction::initiate(123, &[]);
/// // Do the changes here...
/// if tx.commit().is_err() {
///     // Rollback the changes here...
/// }
/// tx.end();
/// ```
///
/// Processor:
/// ```rust,ignore
/// let (mut tx, payload, initiator_id) = Transaction::recv();
/// // Do the changes here...
/// if tx.commit().is_err() {
///     // Rollback the changes here...
/// }
/// tx.end();
/// ```
pub struct Transaction {
	state: TransactionState,
}

impl Transaction {
	/// Initiates a transaction with `target` providing the `payload`.
	///
	/// Wraps [`initiate_transaction`].
	#[inline]
	pub fn initiate(target: GuestId, payload: &[u8]) -> Result<Self, TransactionFailed> {
		initiate_transaction(target, payload)?;
		Ok(Self { state: TransactionState::Initial })
	}

	/// Receives a transaction from the transaction initiator.
	///
	/// Returns new [`Transaction`] instance, the payload and the guest ID of the initiator.
	///
	/// Wraps [`recv_transaction`].
	#[cfg(feature = "alloc")]
	#[inline]
	pub fn recv() -> Result<(Self, alloc::vec::Vec<u8>, GuestId), TransactionFailed> {
		let (payload, source) = recv_transaction();
		let tx = Self { state: TransactionState::Initial };
		Ok((tx, payload, source))
	}

	/// Receives a transaction from the transaction initiator.
	///
	/// Copies the payload to `payload`.
	///
	/// Returns new [`Transaction`] instance, the size of the payload and the guest ID of the
	/// initiator.
	///
	/// Wraps [`recv_transaction_into`].
	#[inline]
	pub fn recv_into(payload: &mut [u8]) -> Result<(Self, usize, GuestId), usize> {
		let (payload_len, source) = recv_transaction_into(payload)?;
		let tx = Self { state: TransactionState::Initial };
		Ok((tx, payload_len, source))
	}

	/// Commits the transaction.
	///
	/// Panics if the transaction has already been committed or rolled back.
	///
	/// Wraps [`commit_transaction`].
	#[inline]
	pub fn commit(&mut self) -> Result<(), TransactionFailed> {
		assert_eq!(TransactionState::Initial, self.state);
		commit_transaction()?;
		self.state = TransactionState::CommittedOrRolledBack;
		Ok(())
	}

	/// Rolls back the transaction.
	///
	/// Panics if the transaction has already been committed or rolled back.
	///
	/// Wraps [`rollback_transaction`].
	#[inline]
	pub fn rollback(&mut self) {
		assert_eq!(TransactionState::Initial, self.state);
		rollback_transaction();
		self.state = TransactionState::CommittedOrRolledBack;
	}

	/// Ends the transaction.
	///
	/// Panics if the transaction hasn't been committed or rolled back.
	///
	/// Wraps [`end_transaction`].
	#[inline]
	pub fn end(mut self) {
		self.inner_end();
	}

	fn inner_end(&mut self) {
		assert_eq!(TransactionState::CommittedOrRolledBack, self.state);
		end_transaction();
		self.state = TransactionState::Ended;
	}
}

impl Drop for Transaction {
	fn drop(&mut self) {
		if self.state == TransactionState::Initial {
			self.rollback();
		}
		if self.state == TransactionState::CommittedOrRolledBack {
			self.inner_end();
		}
	}
}