use crate::proto::Proto;
use yo_common::num::{
DIGITS_MAX, i64_len, push_double, push_fixed4, push_human, push_i64, push_u64, u64_digits,
u64_len,
};
#[derive(Debug, Clone)]
pub struct Out {
buf: Vec<u8>,
proto: Proto,
}
impl Out {
pub fn new(proto: Proto) -> Out {
Out {
buf: Vec::new(),
proto,
}
}
pub fn with_capacity(proto: Proto, cap: usize) -> Out {
Out {
buf: Vec::with_capacity(cap),
proto,
}
}
#[inline]
pub const fn proto(&self) -> Proto {
self.proto
}
#[inline]
pub const fn set_proto(&mut self, proto: Proto) {
self.proto = proto;
}
#[inline]
pub fn as_slice(&self) -> &[u8] {
&self.buf
}
#[inline]
pub fn len(&self) -> usize {
self.buf.len()
}
#[inline]
pub fn capacity(&self) -> usize {
self.buf.capacity()
}
#[inline]
pub fn is_empty(&self) -> bool {
self.buf.is_empty()
}
#[inline]
pub fn clear(&mut self) {
self.buf.clear();
}
pub fn consume(&mut self, n: usize) {
assert!(n <= self.buf.len(), "consumed past the end of the reply");
self.buf.drain(..n);
}
#[inline]
pub fn truncate(&mut self, len: usize) {
self.buf.truncate(len);
}
#[inline]
pub fn reserve(&mut self, n: usize) {
self.buf.reserve(n);
}
pub fn into_inner(self) -> Vec<u8> {
self.buf
}
#[inline]
pub fn raw(&mut self, bytes: &[u8]) {
self.buf.extend_from_slice(bytes);
}
#[inline]
pub fn simple(&mut self, s: &[u8]) {
debug_assert!(
!s.contains(&b'\r') && !s.contains(&b'\n'),
"a simple string cannot carry a line ending, use a bulk string"
);
self.buf.reserve(s.len() + 3);
self.buf.push(b'+');
self.buf.extend_from_slice(s);
self.crlf();
}
#[inline]
pub fn ok(&mut self) {
self.buf.extend_from_slice(b"+OK\r\n");
}
#[inline]
pub fn error(&mut self, msg: &[u8]) {
debug_assert!(
!msg.contains(&b'\r') && !msg.contains(&b'\n'),
"an error line cannot carry a line ending"
);
self.buf.reserve(msg.len() + 3);
self.buf.push(b'-');
self.buf.extend_from_slice(msg);
self.crlf();
}
pub fn error_line(&mut self, prefix: &[u8], msg: &[u8]) {
self.buf.reserve(prefix.len() + msg.len() + 4);
self.buf.push(b'-');
self.buf.extend_from_slice(prefix);
for &b in msg {
self.buf
.push(if b == b'\r' || b == b'\n' { b' ' } else { b });
}
self.crlf();
}
pub fn blob_error(&mut self, msg: &[u8]) {
if self.proto.is_resp3() {
self.blob(b'!', msg);
} else {
self.buf.reserve(msg.len() + 3);
self.buf.push(b'-');
for &b in msg {
self.buf
.push(if b == b'\r' || b == b'\n' { b' ' } else { b });
}
self.crlf();
}
}
#[inline]
pub fn int(&mut self, n: i64) {
self.buf.reserve(i64_len(n) + 3);
self.buf.push(b':');
push_i64(&mut self.buf, n);
self.crlf();
}
#[inline]
pub fn uint(&mut self, n: u64) {
self.buf.reserve(u64_len(n) + 3);
self.buf.push(b':');
push_u64(&mut self.buf, n);
self.crlf();
}
#[inline]
pub fn bulk(&mut self, s: &[u8]) {
self.blob(b'$', s);
}
pub fn bulk_int(&mut self, n: i64) {
let digits = i64_len(n);
self.buf.reserve(digits + 16);
self.buf.push(b'$');
push_u64(&mut self.buf, digits as u64);
self.crlf();
push_i64(&mut self.buf, n);
self.crlf();
}
pub fn bulk_u64(&mut self, n: u64) {
let mut digits = [0u8; DIGITS_MAX];
self.bulk(u64_digits(&mut digits, n));
}
pub fn bulk_double(&mut self, d: f64) {
self.bulk_written(|buf| push_double(buf, d));
}
pub fn human_double(&mut self, d: f64) {
self.bulk_written(|buf| push_human(buf, d));
}
pub fn distance(&mut self, d: f64) {
self.bulk_written(|buf| push_fixed4(buf, d));
}
fn bulk_written(&mut self, f: impl FnOnce(&mut Vec<u8>)) {
self.buf.reserve(48);
let start = self.buf.len();
f(&mut self.buf);
let digits = self.buf.len() - start;
self.buf.push(b'$');
push_u64(&mut self.buf, digits as u64);
self.crlf();
let header = self.buf.len() - start - digits;
self.buf[start..].rotate_right(header);
self.crlf();
}
pub fn verbatim(&mut self, format: &[u8; 3], text: &[u8]) {
if !self.proto.is_resp3() {
self.bulk(text);
return;
}
let len = text.len() + 4;
self.buf.reserve(len + 16);
self.buf.push(b'=');
push_u64(&mut self.buf, len as u64);
self.crlf();
self.buf.extend_from_slice(format);
self.buf.push(b':');
self.buf.extend_from_slice(text);
self.crlf();
}
pub fn big_number(&mut self, digits: &[u8]) {
if self.proto.is_resp3() {
self.buf.reserve(digits.len() + 3);
self.buf.push(b'(');
self.buf.extend_from_slice(digits);
self.crlf();
} else {
self.bulk(digits);
}
}
#[inline]
pub fn nil(&mut self) {
self.buf.extend_from_slice(if self.proto.is_resp3() {
b"_\r\n"
} else {
b"$-1\r\n"
});
}
#[inline]
pub fn nil_array(&mut self) {
self.buf.extend_from_slice(if self.proto.is_resp3() {
b"_\r\n"
} else {
b"*-1\r\n"
});
}
pub fn double(&mut self, d: f64) {
if self.proto.is_resp3() {
self.buf.reserve(32);
self.buf.push(b',');
push_double(&mut self.buf, d);
self.crlf();
return;
}
self.bulk_double(d);
}
#[inline]
pub fn bool(&mut self, b: bool) {
self.buf
.extend_from_slice(match (self.proto.is_resp3(), b) {
(true, true) => b"#t\r\n",
(true, false) => b"#f\r\n",
(false, true) => b":1\r\n",
(false, false) => b":0\r\n",
});
}
#[inline]
pub fn array(&mut self, n: usize) {
self.header(b'*', n);
}
pub fn hoist(&mut self, start: usize, tail: usize) {
assert!(
start + tail <= self.buf.len(),
"hoisted more than was written"
);
self.buf[start..].rotate_right(tail);
}
pub fn close_array(&mut self, start: usize, n: usize) {
self.close(b'*', start, n);
}
pub fn close_set(&mut self, start: usize, n: usize) {
self.close(if self.proto.is_resp3() { b'~' } else { b'*' }, start, n);
}
pub fn close_map(&mut self, start: usize, n: usize) {
self.close(if self.proto.is_resp3() { b'%' } else { b'*' }, start, n);
}
fn close(&mut self, tag: u8, start: usize, n: usize) {
let body = self.buf.len() - start;
self.buf.push(tag);
push_u64(&mut self.buf, n as u64);
self.crlf();
let header = self.buf.len() - start - body;
self.hoist(start, header);
}
#[inline]
pub fn map(&mut self, n: usize) {
if self.proto.is_resp3() {
self.header(b'%', n);
} else {
self.header(b'*', n * 2);
}
}
#[inline]
pub fn set(&mut self, n: usize) {
self.header(if self.proto.is_resp3() { b'~' } else { b'*' }, n);
}
#[inline]
pub fn push(&mut self, n: usize) {
self.header(if self.proto.is_resp3() { b'>' } else { b'*' }, n);
}
#[inline]
pub fn attribute(&mut self, n: usize) {
debug_assert!(
self.proto.is_resp3(),
"RESP2 has no attributes, check the protocol first"
);
self.header(b'|', n);
}
#[inline]
pub const fn bulk_len(value_len: usize) -> usize {
1 + digits_of(value_len as u64) + 2 + value_len + 2
}
#[inline]
pub const fn int_len(n: i64) -> usize {
1 + i64_len(n) + 2
}
#[inline]
pub const fn header_len(n: usize) -> usize {
1 + digits_of(n as u64) + 2
}
#[inline]
fn header(&mut self, kind: u8, n: usize) {
self.buf.reserve(24);
self.buf.push(kind);
push_u64(&mut self.buf, n as u64);
self.crlf();
}
#[inline]
fn blob(&mut self, kind: u8, s: &[u8]) {
self.buf.reserve(Out::bulk_len(s.len()));
self.buf.push(kind);
push_u64(&mut self.buf, s.len() as u64);
self.crlf();
self.buf.extend_from_slice(s);
self.crlf();
}
#[inline]
fn crlf(&mut self) {
self.buf.extend_from_slice(b"\r\n");
}
}
const fn digits_of(n: u64) -> usize {
let mut d = 1;
let mut v = n;
while v >= 10 {
v /= 10;
d += 1;
}
d
}
#[cfg(test)]
mod tests {
use super::*;
fn both(f: impl Fn(&mut Out)) -> (String, String) {
let mut two = Out::new(Proto::Resp2);
let mut three = Out::new(Proto::Resp3);
f(&mut two);
f(&mut three);
(
String::from_utf8(two.into_inner()).unwrap(),
String::from_utf8(three.into_inner()).unwrap(),
)
}
fn one(proto: Proto, f: impl Fn(&mut Out)) -> String {
let mut out = Out::new(proto);
f(&mut out);
String::from_utf8(out.into_inner()).unwrap()
}
#[test]
fn the_three_types_both_protocols_share_are_written_the_same_way() {
let (two, three) = both(|o| {
o.simple(b"PONG");
o.error(b"WRONGTYPE Operation against a key holding the wrong kind of value");
o.int(-42);
o.ok();
});
assert_eq!(two, three);
assert_eq!(
two,
"+PONG\r\n-WRONGTYPE Operation against a key holding the wrong kind of value\r\n:-42\r\n+OK\r\n"
);
}
#[test]
fn a_bulk_string_carries_its_length_and_anything_in_it() {
let (two, three) = both(|o| {
o.bulk(b"hello");
o.bulk(b"");
o.bulk(b"a\r\nb");
});
assert_eq!(two, three);
assert_eq!(two, "$5\r\nhello\r\n$0\r\n\r\n$4\r\na\r\nb\r\n");
}
#[test]
fn a_number_as_a_string_gets_the_right_length() {
assert_eq!(one(Proto::Resp2, |o| o.bulk_int(0)), "$1\r\n0\r\n");
assert_eq!(one(Proto::Resp2, |o| o.bulk_int(-1234)), "$5\r\n-1234\r\n");
assert_eq!(
one(Proto::Resp2, |o| o.bulk_int(i64::MIN)),
"$20\r\n-9223372036854775808\r\n"
);
}
#[test]
fn an_array_can_be_headed_after_its_elements_are_written() {
let (two, three) = both(|o| {
let start = o.len();
o.bulk(b"a");
o.bulk(b"bb");
o.close_array(start, 2);
});
assert_eq!(two, three);
assert_eq!(two, "*2\r\n$1\r\na\r\n$2\r\nbb\r\n");
assert_eq!(
one(Proto::Resp2, |o| {
o.int(1);
let start = o.len();
o.bulk(b"x");
o.close_array(start, 1);
}),
":1\r\n*1\r\n$1\r\nx\r\n"
);
assert_eq!(
one(Proto::Resp2, |o| {
let start = o.len();
o.close_array(start, 0);
}),
"*0\r\n"
);
let long = one(Proto::Resp2, |o| {
let start = o.len();
for _ in 0..100 {
o.int(7);
}
o.close_array(start, 100);
});
assert!(long.starts_with("*100\r\n:7\r\n"));
assert!(long.ends_with(":7\r\n"));
assert_eq!(long.len(), "*100\r\n".len() + 100 * ":7\r\n".len());
}
#[test]
fn resp2_has_two_nulls_and_resp3_has_one() {
let (two, three) = both(|o| {
o.nil();
o.nil_array();
});
assert_eq!(two, "$-1\r\n*-1\r\n");
assert_eq!(three, "_\r\n_\r\n");
}
#[test]
fn a_map_becomes_a_flat_array_on_resp2() {
let (two, three) = both(|o| {
o.map(2);
o.bulk(b"a");
o.bulk(b"1");
o.bulk(b"b");
o.bulk(b"2");
});
assert_eq!(two, "*4\r\n$1\r\na\r\n$1\r\n1\r\n$1\r\nb\r\n$1\r\n2\r\n");
assert_eq!(three, "%2\r\n$1\r\na\r\n$1\r\n1\r\n$1\r\nb\r\n$1\r\n2\r\n");
}
#[test]
fn a_set_and_a_push_become_arrays_on_resp2() {
let (two, three) = both(|o| {
o.set(1);
o.bulk(b"x");
o.push(2);
o.bulk(b"message");
o.bulk(b"ch");
});
assert_eq!(two, "*1\r\n$1\r\nx\r\n*2\r\n$7\r\nmessage\r\n$2\r\nch\r\n");
assert_eq!(
three,
"~1\r\n$1\r\nx\r\n>2\r\n$7\r\nmessage\r\n$2\r\nch\r\n"
);
}
#[test]
fn a_boolean_is_an_integer_on_resp2() {
let (two, three) = both(|o| {
o.bool(true);
o.bool(false);
});
assert_eq!(two, ":1\r\n:0\r\n");
assert_eq!(three, "#t\r\n#f\r\n");
}
#[test]
fn a_double_is_a_bulk_string_on_resp2() {
let (two, three) = both(|o| {
o.double(1.5);
o.double(3.0);
o.double(f64::INFINITY);
});
assert_eq!(two, "$3\r\n1.5\r\n$1\r\n3\r\n$3\r\ninf\r\n");
assert_eq!(three, ",1.5\r\n,3\r\n,inf\r\n");
}
#[test]
fn a_verbatim_string_loses_its_format_on_resp2() {
let (two, three) = both(|o| o.verbatim(b"txt", b"Some string"));
assert_eq!(two, "$11\r\nSome string\r\n");
assert_eq!(three, "=15\r\ntxt:Some string\r\n");
}
#[test]
fn a_big_number_is_a_bulk_string_on_resp2() {
let n = b"3492890328409238509324850943850943825024385";
let (two, three) = both(|o| o.big_number(n));
assert_eq!(
two,
format!("${}\r\n{}\r\n", n.len(), str::from_utf8(n).unwrap())
);
assert_eq!(three, format!("({}\r\n", str::from_utf8(n).unwrap()));
}
#[test]
fn a_blob_error_keeps_its_newlines_on_resp3_and_loses_them_on_resp2() {
let (two, three) = both(|o| o.blob_error(b"SYNTAX bad\nline two"));
assert_eq!(two, "-SYNTAX bad line two\r\n");
assert_eq!(three, "!19\r\nSYNTAX bad\nline two\r\n");
}
#[test]
fn the_predicted_sizes_are_the_sizes_actually_written() {
for len in [0usize, 1, 9, 10, 99, 100, 1000, 65536] {
let value = vec![b'x'; len];
let written = one(Proto::Resp2, |o| o.bulk(&value));
assert_eq!(Out::bulk_len(len), written.len(), "bulk of {len}");
}
for n in [0i64, 7, -7, 100, i64::MAX, i64::MIN] {
let written = one(Proto::Resp2, |o| o.int(n));
assert_eq!(Out::int_len(n), written.len(), "int {n}");
}
for n in [0usize, 5, 1234] {
let written = one(Proto::Resp2, |o| o.array(n));
assert_eq!(Out::header_len(n), written.len(), "array header {n}");
}
}
#[test]
fn hello_switches_the_protocol_for_everything_after_it() {
let mut out = Out::new(Proto::Resp2);
out.nil();
out.set_proto(Proto::Resp3);
out.nil();
assert_eq!(out.as_slice(), b"$-1\r\n_\r\n");
}
#[test]
fn a_partial_write_leaves_the_rest_behind() {
let mut out = Out::new(Proto::Resp2);
out.ok();
out.ok();
out.consume(5);
assert_eq!(out.as_slice(), b"+OK\r\n");
out.clear();
assert!(out.is_empty());
}
}