use std::future::Future;
use std::marker::{Send, Sync};
use celestia_types::namespace_data::NamespaceData;
use celestia_types::nmt::Namespace;
use celestia_types::sample::{RawSample, Sample, SampleId};
use celestia_types::{ExtendedDataSquare, RawShare, Share, ShareProof};
use jsonrpsee::core::client::{ClientT, Error};
use jsonrpsee::proc_macros::rpc;
use serde::{Deserialize, Serialize};
use crate::custom_client_error;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "PascalCase")]
pub struct GetRangeResponse {
pub shares: Vec<Share>,
pub proof: ShareProof,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "UPPERCASE")]
pub enum RowSide {
Left,
Right,
Both,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct GetRowResponse {
pub shares: Vec<Share>,
pub side: RowSide,
}
#[derive(Debug, Copy, Clone, PartialEq, Serialize, Deserialize)]
pub struct SampleCoordinates {
pub row: u16,
#[serde(rename = "col")]
pub column: u16,
}
impl SampleCoordinates {
pub fn new(row: u16, column: u16) -> Self {
SampleCoordinates { row, column }
}
}
impl From<(u16, u16)> for SampleCoordinates {
fn from(value: (u16, u16)) -> Self {
SampleCoordinates::new(value.0, value.1)
}
}
mod rpc {
use super::*;
use celestia_types::eds::RawExtendedDataSquare;
use jsonrpsee::core::RpcResult;
#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
pub struct RawGetRowResponse {
pub(crate) shares: Vec<RawShare>,
pub(crate) side: RowSide,
}
#[rpc(client, server, namespace = "share", namespace_separator = ".")]
trait Share {
#[method(name = "GetEDS")]
async fn share_get_eds(&self, height: u64) -> RpcResult<RawExtendedDataSquare>;
#[method(name = "GetRange")]
async fn share_get_range(
&self,
height: u64,
start: u64,
end: u64,
) -> RpcResult<GetRangeResponse>;
#[method(name = "GetSamples")]
async fn share_get_samples(
&self,
height: u64,
indices: Vec<SampleCoordinates>, ) -> RpcResult<Vec<RawSample>>;
#[method(name = "GetRow")]
async fn share_get_row(&self, height: u64, row: u16) -> RpcResult<RawGetRowResponse>;
#[method(name = "GetShare")]
async fn share_get_share(&self, height: u64, row: u16, col: u16) -> RpcResult<RawShare>;
#[method(name = "GetNamespaceData")]
async fn share_get_namespace_data(
&self,
height: u64,
namespace: Namespace,
) -> RpcResult<NamespaceData>;
#[method(name = "SharesAvailable")]
async fn share_shares_available(&self, height: u64) -> RpcResult<()>;
}
}
pub trait ShareClient: ClientT {
fn share_get_eds<'a, 'fut>(
&'a self,
height: u64,
) -> impl Future<Output = Result<ExtendedDataSquare, Error>> + Send + 'fut
where
'a: 'fut,
Self: Sized + Sync + 'fut,
{
async move {
let raw_eds = rpc::ShareClient::share_get_eds(self, height).await?;
ExtendedDataSquare::from_raw(raw_eds).map_err(custom_client_error)
}
}
fn share_get_range<'a, 'fut>(
&'a self,
height: u64,
start: u64,
end: u64,
) -> impl Future<Output = Result<GetRangeResponse, Error>> + Send + 'fut
where
'a: 'fut,
Self: Sized + Sync + 'fut,
{
rpc::ShareClient::share_get_range(self, height, start, end)
}
fn share_get_samples<'a, 'fut, I, C>(
&'a self,
height: u64,
coordinates: I,
) -> impl Future<Output = Result<Vec<Sample>, Error>> + Send + 'fut
where
'a: 'fut,
Self: Sized + Sync + 'fut,
I: IntoIterator<Item = C>,
C: Into<SampleCoordinates>,
{
let coordinates = coordinates
.into_iter()
.map(|c| c.into())
.collect::<Vec<_>>();
async move {
let raw_samples =
rpc::ShareClient::share_get_samples(self, height, coordinates.clone()).await?;
let mut samples = Vec::with_capacity(raw_samples.len());
for (coords, raw_sample) in coordinates.iter().zip(raw_samples.into_iter()) {
let sample_id = SampleId::new(coords.row, coords.column, height)
.map_err(custom_client_error)?;
let sample =
Sample::from_raw(sample_id, raw_sample).map_err(custom_client_error)?;
samples.push(sample);
}
Ok(samples)
}
}
fn share_get_row<'a, 'fut>(
&'a self,
height: u64,
square_width: u16,
row: u16,
) -> impl Future<Output = Result<GetRowResponse, Error>> + Send + 'fut
where
'a: 'fut,
Self: Sized + Sync + 'fut,
{
async move {
let resp = rpc::ShareClient::share_get_row(self, height, row).await?;
let expected_len = match resp.side {
RowSide::Left | RowSide::Right => square_width / 2,
RowSide::Both => square_width,
};
if resp.shares.len() != usize::from(expected_len) {
return Err(Error::Custom(format!(
"GetRowResponse::shares should have length of {} but received {}",
expected_len,
resp.shares.len(),
)));
}
let shares = resp
.shares
.into_iter()
.enumerate()
.map(|(relative_col, raw_share)| {
let relative_col = u16::try_from(relative_col)
.expect("square_width and expected_len already validated this");
let col = match resp.side {
RowSide::Left | RowSide::Both => relative_col,
RowSide::Right => (square_width / 2) + relative_col,
};
let share = if is_ods_square(row, col, square_width) {
Share::from_raw(&raw_share.data)
} else {
Share::parity(&raw_share.data)
}
.map_err(custom_client_error)?;
Ok(share)
})
.collect::<Result<Vec<_>, Error>>()?;
Ok(GetRowResponse {
shares,
side: resp.side,
})
}
}
fn share_get_share<'a, 'fut>(
&'a self,
height: u64,
square_width: u16,
row: u16,
col: u16,
) -> impl Future<Output = Result<Share, Error>> + Send + 'fut
where
'a: 'fut,
Self: Sized + Sync + 'fut,
{
async move {
let share = rpc::ShareClient::share_get_share(self, height, row, col).await?;
let share = if is_ods_square(row, col, square_width) {
Share::from_raw(&share.data)
} else {
Share::parity(&share.data)
}
.map_err(custom_client_error)?;
Ok(share)
}
}
fn share_get_namespace_data<'a, 'fut>(
&'a self,
height: u64,
namespace: Namespace,
) -> impl Future<Output = Result<NamespaceData, Error>> + Send + 'fut
where
'a: 'fut,
Self: Sized + Sync + 'fut,
{
rpc::ShareClient::share_get_namespace_data(self, height, namespace)
}
fn share_shares_available<'a, 'fut>(
&'a self,
height: u64,
) -> impl Future<Output = Result<(), Error>> + Send + 'fut
where
'a: 'fut,
Self: Sized + Sync + 'fut,
{
rpc::ShareClient::share_shares_available(self, height)
}
}
impl<T> ShareClient for T where T: ClientT {}
fn is_ods_square(row: u16, column: u16, square_width: u16) -> bool {
let ods_width = square_width / 2;
row < ods_width && column < ods_width
}
pub trait ShareServer: rpc::ShareServer {}
impl<T> ShareServer for T where T: rpc::ShareServer {}
pub use rpc::ShareServer as ShareRpcServer;