use crate::caps::{Capability, CapabilitySet};
use crate::error::{Error, Result};
use crate::record::{Header, Tag};
use crate::wire::{Reader, Writer};
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct Hello {
pub abi_major: u16,
pub abi_minor: u16,
pub window_bytes: u64,
pub max_batch_rows: u64,
pub offered: CapabilitySet,
pub source_bytes: u64,
}
impl Hello {
pub const VERSION: u16 = 1;
pub fn encode(&self, w: &mut Writer<'_>) -> Result<()> {
w.record(Tag::HELLO, Self::VERSION, |w| {
w.u16(self.abi_major)?;
w.u16(self.abi_minor)?;
w.u32(0)?;
w.u64(self.window_bytes)?;
w.u64(self.max_batch_rows)?;
w.var_bytes(self.offered.as_bytes())?;
w.u64(self.source_bytes)
})
}
pub fn decode(version: u16, p: &mut Reader<'_>) -> Result<Self> {
expect_version(Tag::HELLO, version)?;
let abi_major = p.u16()?;
let abi_minor = p.u16()?;
p.skip(4)?;
let window_bytes = p.u64()?;
let max_batch_rows = p.u64()?;
let offered = p.capability_set()?;
let source_bytes = p.opt_u64()?.unwrap_or(0);
Ok(Self {
abi_major,
abi_minor,
window_bytes,
max_batch_rows,
offered,
source_bytes,
})
}
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct HelloAck<'a> {
pub abi_major: u16,
pub abi_minor: u16,
pub required: CapabilitySet,
pub optional: CapabilitySet,
pub decoder_id: &'a str,
}
impl<'a> HelloAck<'a> {
pub const VERSION: u16 = 1;
pub fn encode(&self, w: &mut Writer<'_>) -> Result<()> {
w.record(Tag::HELLO_ACK, Self::VERSION, |w| {
w.u16(self.abi_major)?;
w.u16(self.abi_minor)?;
w.u32(0)?;
w.var_bytes(self.required.as_bytes())?;
w.var_bytes(self.optional.as_bytes())?;
w.var_str(self.decoder_id)
})
}
pub fn decode(version: u16, p: &mut Reader<'a>) -> Result<Self> {
expect_version(Tag::HELLO_ACK, version)?;
let abi_major = p.u16()?;
let abi_minor = p.u16()?;
p.skip(4)?;
let required = p.capability_set()?;
let optional = p.capability_set()?;
let decoder_id = p.var_str()?;
Ok(Self {
abi_major,
abi_minor,
required,
optional,
decoder_id,
})
}
}
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
pub struct RefusalReason(pub u16);
impl RefusalReason {
pub const MISSING_CAPABILITY: Self = Self(1);
pub const ABI_TOO_NEW: Self = Self(2);
pub const ABI_TOO_OLD: Self = Self(3);
pub const UNSUPPORTED_RECORD: Self = Self(4);
pub const MALFORMED: Self = Self(5);
pub const RESOURCE_LIMIT: Self = Self(6);
pub const POLICY: Self = Self(7);
#[must_use]
pub const fn name(self) -> Option<&'static str> {
match self {
Self::MISSING_CAPABILITY => Some("missing capability"),
Self::ABI_TOO_NEW => Some("ABI version too new"),
Self::ABI_TOO_OLD => Some("ABI version too old"),
Self::UNSUPPORTED_RECORD => Some("unsupported record"),
Self::MALFORMED => Some("malformed record"),
Self::RESOURCE_LIMIT => Some("resource limit"),
Self::POLICY => Some("refused by policy"),
_ => None,
}
}
}
impl core::fmt::Display for RefusalReason {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self.name() {
Some(name) => f.write_str(name),
None => write!(f, "refusal reason {}", self.0),
}
}
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct Refusal<'a> {
pub reason: RefusalReason,
pub capability: Capability,
pub detail: &'a str,
}
impl<'a> Refusal<'a> {
pub const VERSION: u16 = 1;
#[must_use]
pub const fn new(reason: RefusalReason, detail: &'a str) -> Self {
Self {
reason,
capability: Capability(u16::MAX),
detail,
}
}
pub fn encode(&self, w: &mut Writer<'_>) -> Result<()> {
w.record(Tag::REFUSAL, Self::VERSION, |w| {
w.u16(self.reason.0)?;
w.u16(self.capability.0)?;
w.u32(0)?;
w.var_str(self.detail)
})
}
pub fn decode(version: u16, p: &mut Reader<'a>) -> Result<Self> {
expect_version(Tag::REFUSAL, version)?;
let reason = RefusalReason(p.u16()?);
let capability = Capability(p.u16()?);
p.skip(4)?;
let detail = p.var_str()?;
Ok(Self {
reason,
capability,
detail,
})
}
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct Projection<'a> {
raw: &'a [u8],
}
impl<'a> Projection<'a> {
pub const ALL: Self = Self { raw: &[] };
pub fn from_bytes(raw: &'a [u8]) -> Result<Self> {
if !raw.len().is_multiple_of(4) {
return Err(Error::Malformed(
"a projection must be a whole number of four byte column indices",
));
}
Ok(Self { raw })
}
#[must_use]
pub const fn len(self) -> usize {
self.raw.len() / 4
}
#[must_use]
pub const fn is_empty(self) -> bool {
self.raw.is_empty()
}
pub fn iter(self) -> impl Iterator<Item = u32> + 'a {
self.raw
.as_chunks::<4>()
.0
.iter()
.map(|c| u32::from_le_bytes(*c))
}
#[must_use]
pub const fn as_bytes(self) -> &'a [u8] {
self.raw
}
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct ScanRequest<'a> {
pub row_start: u64,
pub row_count: u64,
pub flags: u64,
pub projection: Projection<'a>,
pub filter: &'a [u8],
}
impl<'a> ScanRequest<'a> {
pub const VERSION: u16 = 1;
#[must_use]
pub const fn everything() -> Self {
Self {
row_start: 0,
row_count: u64::MAX,
flags: 0,
projection: Projection::ALL,
filter: &[],
}
}
pub fn encode(&self, w: &mut Writer<'_>) -> Result<()> {
w.record(Tag::SCAN_REQUEST, Self::VERSION, |w| {
w.u64(self.row_start)?;
w.u64(self.row_count)?;
w.u64(self.flags)?;
w.var_bytes(self.projection.as_bytes())?;
w.var_bytes(self.filter)
})
}
pub fn decode(version: u16, p: &mut Reader<'a>) -> Result<Self> {
expect_version(Tag::SCAN_REQUEST, version)?;
let row_start = p.u64()?;
let row_count = p.u64()?;
let flags = p.u64()?;
let projection = Projection::from_bytes(p.var_bytes()?)?;
let filter = p.var_bytes()?;
Ok(Self {
row_start,
row_count,
flags,
projection,
filter,
})
}
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct RangeRequest {
pub offset: u64,
pub len: u64,
}
impl RangeRequest {
pub const VERSION: u16 = 1;
pub fn encode(&self, w: &mut Writer<'_>) -> Result<()> {
w.record(Tag::RANGE_REQUEST, Self::VERSION, |w| {
w.u64(self.offset)?;
w.u64(self.len)
})
}
pub fn decode(version: u16, p: &mut Reader<'_>) -> Result<Self> {
expect_version(Tag::RANGE_REQUEST, version)?;
let offset = p.u64()?;
let len = p.u64()?;
Ok(Self { offset, len })
}
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct Node {
pub length: u64,
pub null_count: u64,
}
impl Node {
pub const SIZE: usize = 16;
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct BufferRef {
pub offset: u64,
pub len: u64,
}
impl BufferRef {
pub const SIZE: usize = 16;
#[must_use]
pub const fn end(&self) -> Option<u64> {
self.offset.checked_add(self.len)
}
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct Nodes<'a> {
raw: &'a [u8],
}
impl<'a> Nodes<'a> {
pub const EMPTY: Self = Self { raw: &[] };
pub fn from_bytes(raw: &'a [u8]) -> Result<Self> {
if !raw.len().is_multiple_of(Node::SIZE) {
return Err(Error::Malformed(
"a node list must be a whole number of sixteen byte nodes",
));
}
Ok(Self { raw })
}
#[must_use]
pub const fn len(self) -> usize {
self.raw.len() / Node::SIZE
}
#[must_use]
pub const fn is_empty(self) -> bool {
self.raw.is_empty()
}
pub fn iter(self) -> impl Iterator<Item = Node> + 'a {
self.raw.as_chunks::<{ Node::SIZE }>().0.iter().map(|c| {
let (length, null_count) = c.split_at(8);
Node {
length: u64::from_le_bytes(length.try_into().unwrap_or([0; 8])),
null_count: u64::from_le_bytes(null_count.try_into().unwrap_or([0; 8])),
}
})
}
#[must_use]
pub const fn as_bytes(self) -> &'a [u8] {
self.raw
}
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct Buffers<'a> {
raw: &'a [u8],
}
impl<'a> Buffers<'a> {
pub const EMPTY: Self = Self { raw: &[] };
pub fn from_bytes(raw: &'a [u8]) -> Result<Self> {
if !raw.len().is_multiple_of(BufferRef::SIZE) {
return Err(Error::Malformed(
"a buffer list must be a whole number of sixteen byte references",
));
}
Ok(Self { raw })
}
#[must_use]
pub const fn len(self) -> usize {
self.raw.len() / BufferRef::SIZE
}
#[must_use]
pub const fn is_empty(self) -> bool {
self.raw.is_empty()
}
pub fn iter(self) -> impl Iterator<Item = BufferRef> + 'a {
self.raw
.as_chunks::<{ BufferRef::SIZE }>()
.0
.iter()
.map(|c| {
let (offset, len) = c.split_at(8);
BufferRef {
offset: u64::from_le_bytes(offset.try_into().unwrap_or([0; 8])),
len: u64::from_le_bytes(len.try_into().unwrap_or([0; 8])),
}
})
}
#[must_use]
pub const fn as_bytes(self) -> &'a [u8] {
self.raw
}
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct Batch<'a> {
pub rows: u64,
pub flags: u64,
pub nodes: Nodes<'a>,
pub buffers: Buffers<'a>,
}
impl<'a> Batch<'a> {
pub const VERSION: u16 = 1;
#[must_use]
pub const fn empty() -> Self {
Self {
rows: 0,
flags: 0,
nodes: Nodes::EMPTY,
buffers: Buffers::EMPTY,
}
}
pub fn encode(&self, w: &mut Writer<'_>) -> Result<()> {
w.record(Tag::BATCH, Self::VERSION, |w| {
w.u64(self.rows)?;
w.u64(self.flags)?;
w.var_bytes(self.nodes.as_bytes())?;
w.var_bytes(self.buffers.as_bytes())
})
}
pub fn decode(version: u16, p: &mut Reader<'a>) -> Result<Self> {
expect_version(Tag::BATCH, version)?;
let rows = p.u64()?;
let flags = p.u64()?;
let nodes = Nodes::from_bytes(p.var_bytes()?)?;
let buffers = Buffers::from_bytes(p.var_bytes()?)?;
Ok(Self {
rows,
flags,
nodes,
buffers,
})
}
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
#[non_exhaustive]
pub enum Message<'a> {
Hello(Hello),
HelloAck(HelloAck<'a>),
Refusal(Refusal<'a>),
ScanRequest(ScanRequest<'a>),
RangeRequest(RangeRequest),
Batch(Batch<'a>),
Unknown(
Header,
),
}
impl<'a> Reader<'a> {
pub fn message(&mut self) -> Result<Message<'a>> {
let (header, mut p) = self.record()?;
let v = header.version;
Ok(match header.tag {
Tag::HELLO => Message::Hello(Hello::decode(v, &mut p)?),
Tag::HELLO_ACK => Message::HelloAck(HelloAck::decode(v, &mut p)?),
Tag::REFUSAL => Message::Refusal(Refusal::decode(v, &mut p)?),
Tag::SCAN_REQUEST => Message::ScanRequest(ScanRequest::decode(v, &mut p)?),
Tag::RANGE_REQUEST => Message::RangeRequest(RangeRequest::decode(v, &mut p)?),
Tag::BATCH => Message::Batch(Batch::decode(v, &mut p)?),
_ => Message::Unknown(header),
})
}
}
fn expect_version(tag: Tag, version: u16) -> Result<()> {
let known = match tag {
Tag::HELLO => Hello::VERSION,
Tag::HELLO_ACK => HelloAck::VERSION,
Tag::REFUSAL => Refusal::VERSION,
Tag::SCAN_REQUEST => ScanRequest::VERSION,
Tag::RANGE_REQUEST => RangeRequest::VERSION,
Tag::BATCH => Batch::VERSION,
_ => return Err(Error::Malformed("no version is defined for this tag")),
};
if version == known {
Ok(())
} else {
Err(Error::UnsupportedVersion { tag, version })
}
}