jam-std-common 0.1.26

use super::{
	extrinsic::WorkReport,
	state::{Service, Statistics},
};
use bytes::Bytes;
use codec::{Decode, DecodeAll, Encode};
use futures::{stream::BoxStream, Stream, StreamExt};
use jam_types::{
	CoreIndex, Hash, HeaderHash, MmrPeakHash, Segment, SegmentTreeRoot, ServiceId, Slot,
	StateRootHash, WorkPackage, WorkPackageHash, WorkReportHash,
};
use std::{borrow::Cow, future::Future};

#[derive(Debug, thiserror::Error)]
pub enum Error {
	#[error("{0}")]
	Other(String),
	#[error("The block with header hash {0} is not available")]
	BlockUnavailable(HeaderHash),
	#[error("The work-report with hash {0} is not available")]
	WorkReportUnavailable(WorkReportHash),
	#[error("A segment could not be recovered")]
	SegmentUnavailable,
}

impl From<codec::Error> for Error {
	fn from(err: codec::Error) -> Self {
		Self::Other(format!("Codec error: {err}"))
	}
}

pub type NodeResult<T> = Result<T, Error>;

#[derive(Clone, Copy, serde::Serialize, serde::Deserialize, Debug, PartialEq, Eq)]
pub struct BlockDesc {
	pub header_hash: HeaderHash,
	pub slot: Slot,
}

/// An update from a subscription tracking the "best" or finalized chain.
#[derive(serde::Serialize, serde::Deserialize, Debug)]
pub struct ChainSubUpdate<T> {
	/// Header hash of the block that triggered this update.
	pub header_hash: HeaderHash,
	/// Slot of the block with header hash `header_hash`.
	pub slot: Slot,
	/// The tracked value according to the posterior state of the block with header hash
	/// `header_hash`.
	pub value: T,
}

impl<T> ChainSubUpdate<T> {
	pub fn map<R>(self, f: impl FnOnce(T) -> R) -> ChainSubUpdate<R> {
		ChainSubUpdate { header_hash: self.header_hash, slot: self.slot, value: f(self.value) }
	}
}

impl<T> ChainSubUpdate<Option<T>> {
	pub fn map_some<R>(self, f: impl FnOnce(T) -> R) -> ChainSubUpdate<Option<R>> {
		self.map(|value| value.map(f))
	}
}

impl ChainSubUpdate<Bytes> {
	pub fn decode<T: Decode>(&self) -> Result<ChainSubUpdate<T>, codec::Error> {
		Ok(ChainSubUpdate {
			header_hash: self.header_hash,
			slot: self.slot,
			value: T::decode_all(&mut &self.value[..])?,
		})
	}
}

impl ChainSubUpdate<Option<Bytes>> {
	pub fn decode<T: Decode>(&self) -> Result<ChainSubUpdate<Option<T>>, codec::Error> {
		Ok(ChainSubUpdate {
			header_hash: self.header_hash,
			slot: self.slot,
			value: match &self.value {
				Some(value) => Some(T::decode_all(&mut &value[..])?),
				None => None,
			},
		})
	}
}

pub type ChainSub<'a, T> = BoxStream<'a, NodeResult<ChainSubUpdate<T>>>;

#[derive(serde::Serialize, serde::Deserialize, Clone, Debug)]
pub enum VersionedParameters {
	/// Parameters of the chain, version 1.
	V1(jam_types::ProtocolParameters),
}

/// Status of a work-package following execution of a block.
///
/// Note that the status of a work-package may be different on different forks. As
/// `Node::subscribe_work_package_status(finalized: false)` follows the best block, it can jump
/// between forks, and thus potentially yield "impossible" status transitions.
#[derive(Clone, serde::Serialize, serde::Deserialize, Debug, PartialEq, Eq)]
pub enum WorkPackageStatus {
	/// The work-package has not yet been reported, but could be reported in a descendant block.
	Reportable {
		/// The number of blocks remaining until the work-package can no longer be reported.
		///
		/// 1 for example means that the next block is the last block in which the work-package can
		/// be reported.
		remaining_blocks: u16,
	},
	/// The work-package has been reported but is not yet available.
	Reported {
		/// The block in which the work-package was reported.
		reported_in: BlockDesc,
		/// The core on which the work-package was reported.
		core: CoreIndex,
		/// The hash of the work-report that was included on-chain.
		report_hash: WorkReportHash,
	},
	/// The work-package is ready, ie it is either available or has been audited.
	///
	/// A ready work-package is queued for accumulation once its prerequisites are met.
	/// Accumulation of a ready work-package is not guaranteed, in particular its prerequisites may
	/// never be met. Note that there is no "accumulated" status to indicate when accumulation has
	/// happened. To determine if/when a work-package is accumulated, you should monitor the
	/// service's state for the expected changes using eg `Node::subscribe_service_value`.
	Ready {
		/// The block in which the work-package was reported.
		reported_in: BlockDesc,
		/// The core on which the work-package was reported.
		core: CoreIndex,
		/// The hash of the work-report that was included on-chain.
		report_hash: WorkReportHash,
		/// The block in which the work-package became ready.
		ready_in: BlockDesc,
	},
	/// The work-package cannot become ready _on this fork_.
	///
	/// This could be because:
	///
	/// - Its anchor is on a different fork.
	/// - It was not reported in time.
	/// - It did not become available in time.
	///
	/// The `Cow` is a freeform message giving details.
	Failed(Cow<'static, str>),
}

#[derive(Clone, Copy, serde::Serialize, serde::Deserialize, Debug)]
pub struct SyncState {
	pub num_peers: u32,
	pub status: SyncStatus,
}

#[derive(Clone, Copy, serde::Serialize, serde::Deserialize, Debug, Eq, PartialEq)]
pub enum SyncStatus {
	InProgress,
	Completed,
}

/// Generic JAM node interface.
///
/// This trait is implemented directly by eg PolkaJam, and is also implemented for all types
/// implementing [`RpcClient`](crate::RpcClient). JAM tools and builders should ideally use this
/// trait internally, so that they can support both embedded nodes as well as RPC.
///
/// # Chain subscriptions
///
/// The `subscribe` methods which return a [`ChainSub`] all take a `finalized` argument. This
/// argument determines which chain to track: if `finalized` is true, the subscription tracks the
/// latest finalized block; if `finalized` is false, the subscription tracks the head of the "best"
/// chain.
///
/// Note that as the "best" chain may switch to a different fork at any time:
///
/// - Updates yielded by a subscription following the best chain are not guaranteed to ever be
///   included in the finalized chain.
/// - Subscriptions following the best chain may yield "impossible" update sequences. For example,
///   `subscribe_work_package_status(finalized: false)` may yield a `Reported` status followed by a
///   `Reportable` status, if the best chain switches from a fork where the package has been
///   reported to a fork where it has not.
///
/// If these behaviours are unacceptable, use subscriptions tracking the latest finalized block
/// instead. Such subscriptions are well-behaved, but may be significantly delayed compared to
/// best-chain subscriptions.
#[async_trait::async_trait]
pub trait Node: Send + Sync {
	/// Returns the chain parameters.
	async fn parameters(&self) -> NodeResult<VersionedParameters>;

	/// Returns the header hash and slot of the head of the "best" chain.
	async fn best_block(&self) -> NodeResult<BlockDesc>;

	/// Subscribe to updates of the head of the "best" chain, as returned by `best_block`.
	async fn subscribe_best_block(&self) -> NodeResult<BoxStream<NodeResult<BlockDesc>>>;

	/// Returns the header hash and slot of the latest finalized block.
	async fn finalized_block(&self) -> NodeResult<BlockDesc>;

	/// Subscribe to updates of the latest finalized block, as returned by `finalized_block`.
	async fn subscribe_finalized_block(&self) -> NodeResult<BoxStream<NodeResult<BlockDesc>>>;

	/// Returns the header hash and slot of the parent of the block with the given header hash.
	async fn parent(&self, header_hash: HeaderHash) -> NodeResult<BlockDesc>;

	/// Returns the posterior state root of the block with the given header hash.
	async fn state_root(&self, header_hash: HeaderHash) -> NodeResult<StateRootHash>;

	/// Returns the BEEFY root of the block with the given header hash.
	async fn beefy_root(&self, header_hash: HeaderHash) -> NodeResult<MmrPeakHash>;

	/// Returns the activity statistics stored in the posterior state of the block with the given
	/// header hash.
	///
	/// The statistics are encoded as per the GP.
	async fn encoded_statistics(&self, header_hash: HeaderHash) -> NodeResult<Bytes>;

	/// Subscribe to updates of the activity statistics stored in chain state.
	///
	/// The statistics are encoded as per the GP.
	async fn subscribe_encoded_statistics(&self, finalized: bool) -> NodeResult<ChainSub<Bytes>>;

	/// Returns the storage data for the service with the given ID.
	///
	/// The data are encoded as per the GP. `None` is returned if there is no service with the
	/// given ID.
	async fn encoded_service_data(
		&self,
		header_hash: HeaderHash,
		id: ServiceId,
	) -> NodeResult<Option<Bytes>>;

	/// Subscribe to updates of the storage data for the service with the given ID.
	///
	/// The `value` field of updates will contain service data encoded as per the GP. It will be
	/// `None` when there is no service with the given ID.
	async fn subscribe_encoded_service_data(
		&self,
		id: ServiceId,
		finalized: bool,
	) -> NodeResult<ChainSub<Option<Bytes>>>;

	/// Returns the value associated with the given service ID and key in the posterior state of
	/// the block with the given header hash.
	///
	/// `None` is returned if there is no value associated with the given service ID and key.
	///
	/// This method can be used to query arbitrary key-value pairs set by service accumulation
	/// logic.
	async fn service_value(
		&self,
		header_hash: HeaderHash,
		id: ServiceId,
		key: &[u8],
	) -> NodeResult<Option<Bytes>>;

	/// Subscribe to updates of the value associated with the given service ID and key.
	///
	/// The `value` field of updates will be `None` when there is no value associated with the
	/// given service ID and key.
	async fn subscribe_service_value(
		&self,
		id: ServiceId,
		key: &[u8],
		finalized: bool,
	) -> NodeResult<ChainSub<Option<Bytes>>>;

	/// Returns the preimage of the given hash, if it has been provided to the given service in the
	/// posterior state of the block with the given header hash.
	///
	/// `None` is returned if the preimage has not been provided to the given service.
	async fn service_preimage(
		&self,
		header_hash: HeaderHash,
		id: ServiceId,
		hash: Hash,
	) -> NodeResult<Option<Bytes>>;

	/// Subscribe to updates of the preimage associated with the given service ID and hash.
	///
	/// The `value` field of updates will be `None` if the preimage has not been provided to the
	/// service. Otherwise, it will be the preimage of the given hash.
	async fn subscribe_service_preimage(
		&self,
		id: ServiceId,
		hash: Hash,
		finalized: bool,
	) -> NodeResult<ChainSub<Option<Bytes>>>;

	/// Returns the preimage request associated with the given service ID and hash/length in the
	/// posterior state of the block with the given header hash.
	///
	/// `None` is returned if the preimage with the given hash/length has neither been requested by
	/// nor provided to the given service. An empty list is returned if the preimage has been
	/// requested, but not yet provided. A non-empty list means that the preimage has been
	/// provided; the meaning of the slots in the list is as follows:
	///
	/// - The first slot is the slot in which the preimage was provided.
	/// - The second slot, if present, is the slot in which the preimage was "forgotten".
	/// - The third slot, if present, is the slot in which the preimage was requested again.
	async fn service_request(
		&self,
		header_hash: HeaderHash,
		id: ServiceId,
		hash: Hash,
		len: u32,
	) -> NodeResult<Option<Vec<Slot>>>;

	/// Subscribe to updates of the preimage request associated with the given service ID and
	/// hash/length.
	///
	/// The `value` field of updates will either be `None` or a list of slots, with the same
	/// semantics as `service_request` return values.
	async fn subscribe_service_request(
		&self,
		id: ServiceId,
		hash: Hash,
		len: u32,
		finalized: bool,
	) -> NodeResult<ChainSub<Option<Vec<Slot>>>>;

	/// Returns the work-report with the given hash.
	async fn work_report(&self, hash: WorkReportHash) -> NodeResult<WorkReport>;

	/// Submit a work-package to the guarantors currently assigned to the given core.
	async fn submit_encoded_work_package(
		&self,
		core: CoreIndex,
		package: Bytes,
		extrinsics: &[Bytes],
	) -> NodeResult<()>;

	/// Submit a work-package bundle to the guarantors currently assigned to the given core.
	async fn submit_encoded_work_package_bundle(
		&self,
		core: CoreIndex,
		bundle: Bytes,
	) -> NodeResult<()>;

	/// Returns the status of the given work-package following execution of the block with the
	/// given header hash.
	///
	/// If `anchor` does not match the anchor of the work-package then an error or an incorrect
	/// status may be returned. An error may be returned if `anchor` is too old.
	async fn work_package_status(
		&self,
		header_hash: HeaderHash,
		hash: WorkPackageHash,
		anchor: HeaderHash,
	) -> NodeResult<WorkPackageStatus>;

	/// Subscribe to status updates for the given work-package.
	///
	/// If `anchor` does not match the anchor of the work-package then the subscription may fail or
	/// yield incorrect statuses. The subscription may fail if `anchor` is too old.
	async fn subscribe_work_package_status(
		&self,
		hash: WorkPackageHash,
		anchor: HeaderHash,
		finalized: bool,
	) -> NodeResult<ChainSub<WorkPackageStatus>>;

	/// Submit a preimage which is being requested by the given service.
	///
	/// Note that this method does not wait for the preimage to be distributed or integrated
	/// on-chain; it returns immediately.
	async fn submit_preimage(&self, requester: ServiceId, preimage: Bytes) -> NodeResult<()>;

	/// Returns a list of all services currently known to be on JAM.
	///
	/// This is a best-effort list and may not reflect the true state. Nodes could e.g. reasonably
	/// hide services which are not recently active from this list.
	async fn list_services(&self, header_hash: HeaderHash) -> NodeResult<Vec<ServiceId>>;

	/// Fetches a list of segments from the DA layer, exported by the work-package with the given
	/// hash.
	async fn fetch_work_package_segments(
		&self,
		wp_hash: WorkPackageHash,
		indices: Vec<u16>,
	) -> NodeResult<Vec<Segment>>;

	/// Fetches a list of segments from the DA layer, exported by a work-package with the given
	/// segment root hash.
	async fn fetch_segments(
		&self,
		segment_root: SegmentTreeRoot,
		indices: Vec<u16>,
	) -> NodeResult<Vec<Segment>>;

	/// Returns the sync status of the node.
	async fn sync_state(&self) -> NodeResult<SyncState>;

	/// Subscribe to changes in sync status.
	async fn subscribe_sync_status(&self) -> NodeResult<BoxStream<NodeResult<SyncStatus>>>;
}

/// An extension trait for `Node`s which provides various convenience methods. In particular, it
/// provides wrapper methods which encode/decode data provided to/returned from the underlying
/// `Node` methods.
pub trait NodeExt: Node {
	/// Returns the activity statistics stored in the posterior state of the block with the given
	/// header hash.
	fn statistics(
		&self,
		header_hash: HeaderHash,
	) -> impl Future<Output = NodeResult<Statistics>> + Send {
		async move {
			let statistics = self.encoded_statistics(header_hash).await?;
			Ok(Statistics::decode_all(&mut &statistics[..])?)
		}
	}

	/// Subscribe to updates of the activity statistics stored in chain state.
	fn subscribe_statistics(
		&self,
		finalized: bool,
	) -> impl Future<
		Output = NodeResult<impl Stream<Item = NodeResult<ChainSubUpdate<Statistics>>> + Send>,
	> + Send {
		async move {
			let sub = self.subscribe_encoded_statistics(finalized).await?;
			Ok(sub.map(|res| res.and_then(|update| Ok(update.decode()?))))
		}
	}

	/// Returns the storage data for the service with the given ID.
	///
	/// `None` is returned if there is no service with the given ID.
	fn service_data(
		&self,
		header_hash: HeaderHash,
		id: ServiceId,
	) -> impl Future<Output = NodeResult<Option<Service>>> + Send {
		async move {
			let Some(service) = self.encoded_service_data(header_hash, id).await? else {
				return Ok(None)
			};
			Ok(Some(Service::decode_all(&mut &service[..])?))
		}
	}

	/// Subscribe to updates of the storage data for the service with the given ID.
	///
	/// The `value` field of updates will be `None` when there is no service with the given ID.
	fn subscribe_service_data(
		&self,
		id: ServiceId,
		finalized: bool,
	) -> impl Future<
		Output = NodeResult<impl Stream<Item = NodeResult<ChainSubUpdate<Option<Service>>>> + Send>,
	> + Send {
		async move {
			let sub = self.subscribe_encoded_service_data(id, finalized).await?;
			Ok(sub.map(|res| res.and_then(|update| Ok(update.decode()?))))
		}
	}

	/// Returns the value associated with the given service ID and key in the posterior state of
	/// the block with the given header hash.
	///
	/// `None` is returned if there is no value associated with the given service ID and key.
	///
	/// The key is encoded prior to lookup in the state, and the returned value is decoded as type
	/// `V`.
	fn typed_service_value<V: Decode>(
		&self,
		header_hash: HeaderHash,
		id: ServiceId,
		key: &(impl Encode + Sync + ?Sized),
	) -> impl Future<Output = NodeResult<Option<V>>> + Send {
		async move {
			let key = key.encode();
			let Some(value) = self.service_value(header_hash, id, &key).await? else {
				return Ok(None)
			};
			Ok(Some(V::decode_all(&mut &value[..])?))
		}
	}

	/// Subscribe to updates of the value associated with the given service ID and key.
	///
	/// The `value` field of updates will be `None` when there is no value associated with the
	/// given service ID and key.
	///
	/// The key is encoded prior to lookup in the state, and values are decoded as type `V`.
	fn subscribe_typed_service_value<V: Decode>(
		&self,
		id: ServiceId,
		key: &(impl Encode + Sync + ?Sized),
		finalized: bool,
	) -> impl Future<
		Output = NodeResult<impl Stream<Item = NodeResult<ChainSubUpdate<Option<V>>>> + Send>,
	> + Send {
		async move {
			let key = key.encode();
			let sub = self.subscribe_service_value(id, &key, finalized).await?;
			Ok(sub.map(|res| res.and_then(|update| Ok(update.decode()?))))
		}
	}

	/// Submit a work-package to the guarantors currently assigned to the given core.
	fn submit_work_package(
		&self,
		core: CoreIndex,
		package: &WorkPackage,
		extrinsics: &[Bytes],
	) -> impl Future<Output = NodeResult<()>> + Send {
		async move {
			self.submit_encoded_work_package(core, package.encode().into(), extrinsics)
				.await
		}
	}
}

impl<T: Node + ?Sized> NodeExt for T {}