use yo_common::{Code, Error, Result};
use yo_kv::{Cursor, KeyCursor};
use crate::reply::Out;
pub(super) trait Resume: Copy {
fn from_raw(raw: u64) -> Self;
fn raw(self) -> u64;
}
impl Resume for Cursor {
fn from_raw(raw: u64) -> Cursor {
Cursor::from_raw(raw)
}
fn raw(self) -> u64 {
Cursor::raw(self)
}
}
impl Resume for KeyCursor {
fn from_raw(raw: u64) -> KeyCursor {
KeyCursor::from_raw(raw)
}
fn raw(self) -> u64 {
KeyCursor::raw(self)
}
}
pub(super) const BAD_CURSOR: &str = "invalid cursor";
pub(super) const COUNT: usize = 10;
pub(super) fn parse_cursor<C: Resume>(arg: &[u8]) -> Result<C> {
let bad = || Error::new(Code::Invalid, BAD_CURSOR);
let digits = arg.trim_ascii_start();
let digits = digits.strip_prefix(b"+").unwrap_or(digits);
if digits.is_empty() || !digits.iter().all(u8::is_ascii_digit) {
return Err(bad());
}
let mut raw: u64 = 0;
for b in digits {
raw = raw
.checked_mul(10)
.and_then(|v| v.checked_add(u64::from(b - b'0')))
.ok_or_else(bad)?;
}
Ok(C::from_raw(raw))
}
pub(super) fn reply<C: Resume>(
out: &mut Out,
walk: impl FnOnce(&mut Out) -> Result<(C, usize)>,
) -> Result<()> {
out.array(2);
let at = out.len();
let (next, n) = walk(out)?;
out.close_array(at, n);
let body = out.len() - at;
out.bulk_u64(next.raw());
let cursor = out.len() - at - body;
out.hoist(at, cursor);
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_cursor_is_read_the_way_strtoull_reads_one() {
for ok in [
&b"0"[..],
b"+0",
b" 0",
b" +0",
b"007",
b"18446744073709551615",
] {
assert!(
parse_cursor::<Cursor>(ok).is_ok(),
"{:?}",
core::str::from_utf8(ok)
);
}
for bad in [
&b""[..],
b"-1",
b"0 ",
b"0abc",
b"abc",
b"0x10",
b"+ 0",
b"18446744073709551616",
] {
let got = parse_cursor::<Cursor>(bad).expect_err("this is not a cursor");
assert_eq!(got.message(), BAD_CURSOR, "{:?}", core::str::from_utf8(bad));
}
}
#[test]
fn the_cursor_ends_up_in_front_of_the_elements() {
let mut out = Out::new(crate::proto::Proto::Resp2);
reply(&mut out, |out| {
out.bulk(b"a");
out.bulk(b"b");
Ok((Cursor::from_raw(7), 2))
})
.expect("a walk that cannot fail");
assert_eq!(
core::str::from_utf8(out.as_slice()).expect("ascii"),
"*2\r\n$1\r\n7\r\n*2\r\n$1\r\na\r\n$1\r\nb\r\n"
);
}
}