use std::num::NonZeroU64;
use haematite_query::Refusal;
use haematite_query::handle::Error;
use haematite_query::page::{Allowance, Origin, RecordSource, Resume};
use serde::{Deserialize, Serialize};
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum Rows {
#[default]
Values,
Origins,
}
#[derive(Debug, Serialize)]
pub struct Named {
pub origin: Origin,
}
#[derive(Debug)]
pub struct Page {
pub rows: Vec<Named>,
pub next: Option<Resume>,
pub records_read: u64,
counts: Vec<u64>,
}
fn protocol(error: impl std::fmt::Display) -> Error {
Error::new(Refusal::QueryProtocol, error.to_string())
}
const fn after(origin: Origin, records_read: u64, passed_counts: Vec<u64>) -> Resume {
Resume {
records_read,
pending: origin,
discard_outputs: 0,
passed: true,
passed_counts,
}
}
pub fn page(
source: &mut impl RecordSource,
allowance: &Allowance,
resume: Option<&Resume>,
limits: &[NonZeroU64],
) -> Result<Page, Error> {
let read_before = resume.map_or(0, |resume| resume.records_read);
let counts = resume.map_or_else(
|| vec![0; limits.len()],
|resume| resume.passed_counts.clone(),
);
if counts.len() != limits.len()
|| counts
.iter()
.zip(limits)
.any(|(count, limit)| *count > limit.get())
{
return Err(protocol(
"origins continuation counters do not match the limit stages",
));
}
if resume.is_some_and(|resume| !resume.passed || resume.discard_outputs != 0) {
return Err(protocol(
"origins continuation must start after a passed place",
));
}
let mut page = Page {
rows: Vec::new(),
next: None,
records_read: read_before,
counts,
};
let mut bytes = 2_u64;
loop {
if page
.counts
.iter()
.zip(limits)
.any(|(count, limit)| *count == limit.get())
{
return Ok(page);
}
let given = u64::try_from(page.rows.len()).map_err(protocol)?;
if given == allowance.rows.get() {
return Ok(page.cut());
}
let Some(record) = source.next_record()? else {
page.next = source
.passed()
.map(|origin| after(origin, page.records_read, page.counts.clone()));
return Ok(page);
};
let row = Named {
origin: record.origin,
};
let size =
u64::try_from(serde_json::to_vec(&row).map_err(protocol)?.len()).map_err(protocol)?;
if size.saturating_add(2) > allowance.bytes.get() {
return Err(Error::new(
Refusal::QueryValueTooLarge,
format!(
"the origin of {}/{} needs more than {} bytes",
row.origin.collection, row.origin.id, allowance.bytes
),
));
}
let separator = u64::from(given > 0);
let grown = bytes.saturating_add(separator).saturating_add(size);
if grown > allowance.bytes.get() {
return Ok(page.cut());
}
bytes = grown;
page.records_read = page
.records_read
.checked_add(1)
.ok_or_else(|| protocol("records read overflow"))?;
for count in &mut page.counts {
*count = count
.checked_add(1)
.ok_or_else(|| protocol("limit count overflow"))?;
}
page.rows.push(row);
}
}
impl Page {
fn cut(mut self) -> Self {
self.next = self
.rows
.last()
.map(|row| after(row.origin.clone(), self.records_read, self.counts.clone()));
self
}
}