scylla_cql/frame/
types.rs

1//! CQL binary protocol in-wire types.
2
3use super::frame_errors::LowLevelDeserializationError;
4use super::TryFromPrimitiveError;
5use byteorder::{BigEndian, ReadBytesExt};
6use bytes::Bytes;
7#[cfg(test)]
8use bytes::BytesMut;
9use bytes::{Buf, BufMut};
10use std::collections::HashMap;
11use std::convert::TryFrom;
12use std::convert::TryInto;
13use std::net::IpAddr;
14use std::net::SocketAddr;
15use std::str;
16use thiserror::Error;
17use uuid::Uuid;
18
19/// A setting that defines a successful write or read by the number of cluster replicas
20/// that acknowledge the write or respond to the read request, respectively.
21/// See [ScyllaDB docs](https://docs.scylladb.com/manual/stable/cql/consistency.html)
22/// for more detailed description and guidelines.
23#[derive(Debug, Copy, Clone, Default, PartialEq, Eq, PartialOrd, Ord)]
24#[cfg_attr(feature = "serde", derive(serde::Deserialize))]
25#[cfg_attr(feature = "serde", serde(rename_all = "SCREAMING_SNAKE_CASE"))]
26#[repr(u16)]
27pub enum Consistency {
28    /// **Write-only**. Closest replica, as determined by the
29    /// [Snitch](https://docs.scylladb.com/manual/stable/reference/glossary.html#term-Snitch),
30    /// must respond. If all replica nodes are down, write succeeds after a hinted handoff.
31    /// Provides low latency, guarantees writes never fail.
32    Any = 0x0000,
33    /// The closest replica as determined by the Snitch must respond. Consistency requirements
34    /// are not too strict.
35    One = 0x0001,
36    /// The closest two replicas as determined by the Snitch must respond.
37    Two = 0x0002,
38    /// The closest three replicas as determined by the Snitch must respond.
39    Three = 0x0003,
40    /// A simple majority of all replicas across all datacenters must respond.
41    /// This CL allows for some level of failure.
42    Quorum = 0x0004,
43    /// _All_ replicas in the cluster must respond. May cause performance issues.
44    All = 0x0005,
45    /// Same as [QUORUM](Consistency::Quorum), but confined to the same datacenter as the coordinator.
46    #[default]
47    LocalQuorum = 0x0006,
48    /// **Write-only**. A simple majority in each datacenter must respond.
49    EachQuorum = 0x0007,
50    /// Same as [ONE](Consistency::One), but confined to the local datacenter.
51    LocalOne = 0x000A,
52
53    // Apparently, Consistency can be set to Serial or LocalSerial in SELECT statements
54    // to make them use Paxos.
55    /// **Read-only**. Returns results with the most recent data. Including uncommitted in-flight LWTs.
56    Serial = 0x0008,
57    /// **Read-only**. Same as [SERIAL](Consistency::Serial), but confined to a local datacenter.
58    LocalSerial = 0x0009,
59}
60
61impl TryFrom<u16> for Consistency {
62    type Error = TryFromPrimitiveError<u16>;
63
64    fn try_from(value: u16) -> Result<Self, Self::Error> {
65        match value {
66            0x0000 => Ok(Consistency::Any),
67            0x0001 => Ok(Consistency::One),
68            0x0002 => Ok(Consistency::Two),
69            0x0003 => Ok(Consistency::Three),
70            0x0004 => Ok(Consistency::Quorum),
71            0x0005 => Ok(Consistency::All),
72            0x0006 => Ok(Consistency::LocalQuorum),
73            0x0007 => Ok(Consistency::EachQuorum),
74            0x000A => Ok(Consistency::LocalOne),
75            0x0008 => Ok(Consistency::Serial),
76            0x0009 => Ok(Consistency::LocalSerial),
77            _ => Err(TryFromPrimitiveError {
78                enum_name: "Consistency",
79                primitive: value,
80            }),
81        }
82    }
83}
84
85/// [Consistency] for Lightweight Transactions (LWTs).
86///
87/// [SerialConsistency] sets the consistency level for the serial phase of conditional updates.
88/// This option will be ignored for anything else that a conditional update/insert.
89#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
90#[cfg_attr(feature = "serde", derive(serde::Deserialize))]
91#[cfg_attr(feature = "serde", serde(rename_all = "SCREAMING_SNAKE_CASE"))]
92#[repr(i16)]
93pub enum SerialConsistency {
94    /// Guarantees linearizable semantics across the whole cluster.
95    Serial = 0x0008,
96    /// Guarantees linearizable semantics in a local datacenter.
97    LocalSerial = 0x0009,
98}
99
100impl TryFrom<i16> for SerialConsistency {
101    type Error = TryFromPrimitiveError<i16>;
102
103    fn try_from(value: i16) -> Result<Self, Self::Error> {
104        match value {
105            0x0008 => Ok(Self::Serial),
106            0x0009 => Ok(Self::LocalSerial),
107            _ => Err(TryFromPrimitiveError {
108                enum_name: "SerialConsistency",
109                primitive: value,
110            }),
111        }
112    }
113}
114
115impl Consistency {
116    /// Checks if the consistency is a serial consistency.
117    pub fn is_serial(&self) -> bool {
118        matches!(self, Consistency::Serial | Consistency::LocalSerial)
119    }
120}
121
122/// Error returned when a serial consistency what expected, yet got another kind of consistency.
123#[derive(Debug, Error)]
124#[error("Expected Consistency Serial or LocalSerial, got: {0}")]
125pub struct NonSerialConsistencyError(Consistency);
126
127impl TryFrom<Consistency> for SerialConsistency {
128    type Error = NonSerialConsistencyError;
129
130    fn try_from(c: Consistency) -> Result<Self, Self::Error> {
131        match c {
132            Consistency::Any
133            | Consistency::One
134            | Consistency::Two
135            | Consistency::Three
136            | Consistency::Quorum
137            | Consistency::All
138            | Consistency::LocalQuorum
139            | Consistency::EachQuorum
140            | Consistency::LocalOne => Err(NonSerialConsistencyError(c)),
141            Consistency::Serial => Ok(SerialConsistency::Serial),
142            Consistency::LocalSerial => Ok(SerialConsistency::LocalSerial),
143        }
144    }
145}
146
147impl std::fmt::Display for Consistency {
148    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
149        write!(f, "{self:?}")
150    }
151}
152
153impl std::fmt::Display for SerialConsistency {
154    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
155        write!(f, "{self:?}")
156    }
157}
158
159#[derive(Copy, Clone, Debug, PartialEq, Eq)]
160pub enum RawValue<'a> {
161    Null,
162    Unset,
163    Value(&'a [u8]),
164}
165
166impl<'a> RawValue<'a> {
167    #[inline]
168    pub fn as_value(&self) -> Option<&'a [u8]> {
169        match self {
170            RawValue::Value(v) => Some(v),
171            RawValue::Null | RawValue::Unset => None,
172        }
173    }
174}
175
176pub(crate) fn read_raw_bytes<'a>(
177    count: usize,
178    buf: &mut &'a [u8],
179) -> Result<&'a [u8], LowLevelDeserializationError> {
180    if buf.len() < count {
181        return Err(LowLevelDeserializationError::TooFewBytesReceived {
182            expected: count,
183            received: buf.len(),
184        });
185    }
186    let (ret, rest) = buf.split_at(count);
187    *buf = rest;
188    Ok(ret)
189}
190
191pub fn read_int(buf: &mut &[u8]) -> Result<i32, std::io::Error> {
192    let v = buf.read_i32::<BigEndian>()?;
193    Ok(v)
194}
195
196pub fn write_int(v: i32, buf: &mut impl BufMut) {
197    buf.put_i32(v);
198}
199
200pub fn read_int_length(buf: &mut &[u8]) -> Result<usize, LowLevelDeserializationError> {
201    let v = read_int(buf)?;
202    let v: usize = v.try_into()?;
203
204    Ok(v)
205}
206
207pub(crate) fn write_int_length(
208    v: usize,
209    buf: &mut impl BufMut,
210) -> Result<(), std::num::TryFromIntError> {
211    let v: i32 = v.try_into()?;
212
213    write_int(v, buf);
214    Ok(())
215}
216
217#[test]
218fn type_int() {
219    let vals = [i32::MIN, -1, 0, 1, i32::MAX];
220    for val in vals.iter() {
221        let mut buf = Vec::new();
222        write_int(*val, &mut buf);
223        assert_eq!(read_int(&mut &buf[..]).unwrap(), *val);
224    }
225}
226
227pub fn read_long(buf: &mut &[u8]) -> Result<i64, std::io::Error> {
228    let v = buf.read_i64::<BigEndian>()?;
229    Ok(v)
230}
231
232pub fn write_long(v: i64, buf: &mut impl BufMut) {
233    buf.put_i64(v);
234}
235
236#[test]
237fn type_long() {
238    let vals = [i64::MIN, -1, 0, 1, i64::MAX];
239    for val in vals.iter() {
240        let mut buf = Vec::new();
241        write_long(*val, &mut buf);
242        assert_eq!(read_long(&mut &buf[..]).unwrap(), *val);
243    }
244}
245
246pub fn read_short(buf: &mut &[u8]) -> Result<u16, std::io::Error> {
247    let v = buf.read_u16::<BigEndian>()?;
248    Ok(v)
249}
250
251pub fn write_short(v: u16, buf: &mut impl BufMut) {
252    buf.put_u16(v);
253}
254
255pub(crate) fn read_short_length(buf: &mut &[u8]) -> Result<usize, std::io::Error> {
256    let v = read_short(buf)?;
257    let v: usize = v.into();
258    Ok(v)
259}
260
261pub(crate) fn write_short_length(
262    v: usize,
263    buf: &mut impl BufMut,
264) -> Result<(), std::num::TryFromIntError> {
265    let v: u16 = v.try_into()?;
266    write_short(v, buf);
267    Ok(())
268}
269
270#[test]
271fn type_short() {
272    let vals: [u16; 3] = [0, 1, u16::MAX];
273    for val in vals.iter() {
274        let mut buf = Vec::new();
275        write_short(*val, &mut buf);
276        assert_eq!(read_short(&mut &buf[..]).unwrap(), *val);
277    }
278}
279
280// https://github.com/apache/cassandra/blob/trunk/doc/native_protocol_v4.spec#L208
281pub fn read_bytes_opt<'a>(
282    buf: &mut &'a [u8],
283) -> Result<Option<&'a [u8]>, LowLevelDeserializationError> {
284    let len = read_int(buf)?;
285    if len < 0 {
286        return Ok(None);
287    }
288    let len = len as usize;
289    let v = Some(read_raw_bytes(len, buf)?);
290    Ok(v)
291}
292
293// Same as read_bytes, but we assume the value won't be `null`
294pub fn read_bytes<'a>(buf: &mut &'a [u8]) -> Result<&'a [u8], LowLevelDeserializationError> {
295    let len = read_int_length(buf)?;
296    let v = read_raw_bytes(len, buf)?;
297    Ok(v)
298}
299
300pub fn read_value<'a>(buf: &mut &'a [u8]) -> Result<RawValue<'a>, LowLevelDeserializationError> {
301    let len = read_int(buf)?;
302    match len {
303        -2 => Ok(RawValue::Unset),
304        -1 => Ok(RawValue::Null),
305        len if len >= 0 => {
306            let v = read_raw_bytes(len as usize, buf)?;
307            Ok(RawValue::Value(v))
308        }
309        len => Err(LowLevelDeserializationError::InvalidValueLength(len)),
310    }
311}
312
313pub fn read_short_bytes<'a>(buf: &mut &'a [u8]) -> Result<&'a [u8], LowLevelDeserializationError> {
314    let len = read_short_length(buf)?;
315    let v = read_raw_bytes(len, buf)?;
316    Ok(v)
317}
318
319pub fn write_bytes(v: &[u8], buf: &mut impl BufMut) -> Result<(), std::num::TryFromIntError> {
320    write_int_length(v.len(), buf)?;
321    buf.put_slice(v);
322    Ok(())
323}
324
325pub fn write_bytes_opt(
326    v: Option<impl AsRef<[u8]>>,
327    buf: &mut impl BufMut,
328) -> Result<(), std::num::TryFromIntError> {
329    match v {
330        Some(bytes) => {
331            write_int_length(bytes.as_ref().len(), buf)?;
332            buf.put_slice(bytes.as_ref());
333        }
334        None => write_int(-1, buf),
335    }
336
337    Ok(())
338}
339
340pub fn write_short_bytes(v: &[u8], buf: &mut impl BufMut) -> Result<(), std::num::TryFromIntError> {
341    write_short_length(v.len(), buf)?;
342    buf.put_slice(v);
343    Ok(())
344}
345
346pub fn read_bytes_map(
347    buf: &mut &[u8],
348) -> Result<HashMap<String, Bytes>, LowLevelDeserializationError> {
349    let len = read_short_length(buf)?;
350    let mut v = HashMap::with_capacity(len);
351    for _ in 0..len {
352        let key = read_string(buf)?.to_owned();
353        let val = Bytes::copy_from_slice(read_bytes(buf)?);
354        v.insert(key, val);
355    }
356    Ok(v)
357}
358
359pub fn write_bytes_map<B>(
360    v: &HashMap<String, B>,
361    buf: &mut impl BufMut,
362) -> Result<(), std::num::TryFromIntError>
363where
364    B: AsRef<[u8]>,
365{
366    let len = v.len();
367    write_short_length(len, buf)?;
368    for (key, val) in v.iter() {
369        write_string(key, buf)?;
370        write_bytes(val.as_ref(), buf)?;
371    }
372    Ok(())
373}
374
375#[test]
376fn type_bytes_map() {
377    let mut val = HashMap::new();
378    val.insert("".to_owned(), Bytes::new());
379    val.insert("EXTENSION1".to_owned(), Bytes::from_static(&[1, 2, 3]));
380    val.insert("EXTENSION2".to_owned(), Bytes::from_static(&[4, 5, 6]));
381    let mut buf = BytesMut::new();
382    write_bytes_map(&val, &mut buf).unwrap();
383    assert_eq!(read_bytes_map(&mut &*buf).unwrap(), val);
384}
385
386pub fn read_string<'a>(buf: &mut &'a [u8]) -> Result<&'a str, LowLevelDeserializationError> {
387    let len = read_short_length(buf)?;
388    let raw = read_raw_bytes(len, buf)?;
389    let v = str::from_utf8(raw)?;
390    Ok(v)
391}
392
393pub fn write_string(v: &str, buf: &mut impl BufMut) -> Result<(), std::num::TryFromIntError> {
394    let raw = v.as_bytes();
395    write_short_length(v.len(), buf)?;
396    buf.put_slice(raw);
397    Ok(())
398}
399
400#[test]
401fn type_string() {
402    let vals = [String::from(""), String::from("hello, world!")];
403    for val in vals.iter() {
404        let mut buf = Vec::new();
405        write_string(val, &mut buf).unwrap();
406        assert_eq!(read_string(&mut &buf[..]).unwrap(), *val);
407    }
408}
409
410pub fn read_long_string<'a>(buf: &mut &'a [u8]) -> Result<&'a str, LowLevelDeserializationError> {
411    let len = read_int_length(buf)?;
412    let raw = read_raw_bytes(len, buf)?;
413    let v = str::from_utf8(raw)?;
414    Ok(v)
415}
416
417pub fn write_long_string(v: &str, buf: &mut impl BufMut) -> Result<(), std::num::TryFromIntError> {
418    let raw = v.as_bytes();
419    let len = raw.len();
420    write_int_length(len, buf)?;
421    buf.put_slice(raw);
422    Ok(())
423}
424
425#[test]
426fn type_long_string() {
427    let vals = [String::from(""), String::from("hello, world!")];
428    for val in vals.iter() {
429        let mut buf = Vec::new();
430        write_long_string(val, &mut buf).unwrap();
431        assert_eq!(read_long_string(&mut &buf[..]).unwrap(), *val);
432    }
433}
434
435pub fn read_string_map(
436    buf: &mut &[u8],
437) -> Result<HashMap<String, String>, LowLevelDeserializationError> {
438    let len = read_short_length(buf)?;
439    let mut v = HashMap::with_capacity(len);
440    for _ in 0..len {
441        let key = read_string(buf)?.to_owned();
442        let val = read_string(buf)?.to_owned();
443        v.insert(key, val);
444    }
445    Ok(v)
446}
447
448pub fn write_string_map(
449    v: &HashMap<impl AsRef<str>, impl AsRef<str>>,
450    buf: &mut impl BufMut,
451) -> Result<(), std::num::TryFromIntError> {
452    let len = v.len();
453    write_short_length(len, buf)?;
454    for (key, val) in v.iter() {
455        write_string(key.as_ref(), buf)?;
456        write_string(val.as_ref(), buf)?;
457    }
458    Ok(())
459}
460
461#[test]
462fn type_string_map() {
463    let mut val = HashMap::new();
464    val.insert(String::from(""), String::from(""));
465    val.insert(String::from("CQL_VERSION"), String::from("3.0.0"));
466    val.insert(String::from("THROW_ON_OVERLOAD"), String::from(""));
467    let mut buf = Vec::new();
468    write_string_map(&val, &mut buf).unwrap();
469    assert_eq!(read_string_map(&mut &buf[..]).unwrap(), val);
470}
471
472pub fn read_string_list(buf: &mut &[u8]) -> Result<Vec<String>, LowLevelDeserializationError> {
473    let len = read_short_length(buf)?;
474    let mut v = Vec::with_capacity(len);
475    for _ in 0..len {
476        v.push(read_string(buf)?.to_owned());
477    }
478    Ok(v)
479}
480
481pub fn write_string_list(
482    v: &[String],
483    buf: &mut impl BufMut,
484) -> Result<(), std::num::TryFromIntError> {
485    let len = v.len();
486    write_short_length(len, buf)?;
487    for v in v.iter() {
488        write_string(v, buf)?;
489    }
490    Ok(())
491}
492
493#[test]
494fn type_string_list() {
495    let val = vec![
496        "".to_owned(),
497        "CQL_VERSION".to_owned(),
498        "THROW_ON_OVERLOAD".to_owned(),
499    ];
500
501    let mut buf = Vec::new();
502    write_string_list(&val, &mut buf).unwrap();
503    assert_eq!(read_string_list(&mut &buf[..]).unwrap(), val);
504}
505
506pub fn read_string_multimap(
507    buf: &mut &[u8],
508) -> Result<HashMap<String, Vec<String>>, LowLevelDeserializationError> {
509    let len = read_short_length(buf)?;
510    let mut v = HashMap::with_capacity(len);
511    for _ in 0..len {
512        let key = read_string(buf)?.to_owned();
513        let val = read_string_list(buf)?;
514        v.insert(key, val);
515    }
516    Ok(v)
517}
518
519pub fn write_string_multimap(
520    v: &HashMap<String, Vec<String>>,
521    buf: &mut impl BufMut,
522) -> Result<(), std::num::TryFromIntError> {
523    let len = v.len();
524    write_short_length(len, buf)?;
525    for (key, val) in v.iter() {
526        write_string(key, buf)?;
527        write_string_list(val, buf)?;
528    }
529    Ok(())
530}
531
532#[test]
533fn type_string_multimap() {
534    let mut val = HashMap::new();
535    val.insert(String::from(""), vec![String::from("")]);
536    val.insert(
537        String::from("versions"),
538        vec![String::from("3.0.0"), String::from("4.2.0")],
539    );
540    val.insert(String::from("empty"), vec![]);
541    let mut buf = Vec::new();
542    write_string_multimap(&val, &mut buf).unwrap();
543    assert_eq!(read_string_multimap(&mut &buf[..]).unwrap(), val);
544}
545
546pub fn read_uuid(buf: &mut &[u8]) -> Result<Uuid, LowLevelDeserializationError> {
547    let raw = read_raw_bytes(16, buf)?;
548
549    // It's safe to unwrap here because the conversion only fails
550    // if the argument slice's length does not match, which
551    // `read_raw_bytes` prevents.
552    let raw_array: &[u8; 16] = raw.try_into().unwrap();
553
554    Ok(Uuid::from_bytes(*raw_array))
555}
556
557pub fn write_uuid(uuid: &Uuid, buf: &mut impl BufMut) {
558    buf.put_slice(&uuid.as_bytes()[..]);
559}
560
561#[test]
562fn type_uuid() {
563    let u = Uuid::parse_str("f3b4958c-52a1-11e7-802a-010203040506").unwrap();
564    let mut buf = Vec::new();
565    write_uuid(&u, &mut buf);
566    let u2 = read_uuid(&mut &*buf).unwrap();
567    assert_eq!(u, u2);
568}
569
570pub fn read_consistency(buf: &mut &[u8]) -> Result<Consistency, LowLevelDeserializationError> {
571    let raw = read_short(buf)?;
572    Consistency::try_from(raw).map_err(LowLevelDeserializationError::UnknownConsistency)
573}
574
575pub fn write_consistency(c: Consistency, buf: &mut impl BufMut) {
576    write_short(c as u16, buf);
577}
578
579pub fn write_serial_consistency(c: SerialConsistency, buf: &mut impl BufMut) {
580    write_short(c as u16, buf);
581}
582
583#[test]
584fn type_consistency() {
585    let c = Consistency::Quorum;
586    let mut buf = Vec::new();
587    write_consistency(c, &mut buf);
588    let c2 = read_consistency(&mut &*buf).unwrap();
589    assert_eq!(c, c2);
590
591    let c: i16 = 0x1234;
592    buf.clear();
593    buf.put_i16(c);
594    let c_result = read_consistency(&mut &*buf);
595    assert!(c_result.is_err());
596
597    // Check that the error message contains information about the invalid value
598    let err_str = format!("{}", c_result.unwrap_err());
599    assert!(err_str.contains(&format!("{c}")));
600}
601
602pub fn read_inet(buf: &mut &[u8]) -> Result<SocketAddr, LowLevelDeserializationError> {
603    let len = buf.read_u8()?;
604    let ip_addr = match len {
605        4 => {
606            let ret = IpAddr::from(<[u8; 4]>::try_from(&buf[0..4])?);
607            buf.advance(4);
608            ret
609        }
610        16 => {
611            let ret = IpAddr::from(<[u8; 16]>::try_from(&buf[0..16])?);
612            buf.advance(16);
613            ret
614        }
615        v => return Err(LowLevelDeserializationError::InvalidInetLength(v)),
616    };
617    let port = read_int(buf)?;
618
619    Ok(SocketAddr::new(ip_addr, port as u16))
620}
621
622pub fn write_inet(addr: SocketAddr, buf: &mut impl BufMut) {
623    match addr.ip() {
624        IpAddr::V4(v4) => {
625            buf.put_u8(4);
626            buf.put_slice(&v4.octets());
627        }
628        IpAddr::V6(v6) => {
629            buf.put_u8(16);
630            buf.put_slice(&v6.octets());
631        }
632    }
633
634    write_int(addr.port() as i32, buf)
635}
636
637#[test]
638fn type_inet() {
639    use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
640
641    let iv4 = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 1234);
642    let iv6 = SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)), 2345);
643    let mut buf = Vec::new();
644
645    write_inet(iv4, &mut buf);
646    let read_iv4 = read_inet(&mut &*buf).unwrap();
647    assert_eq!(iv4, read_iv4);
648    buf.clear();
649
650    write_inet(iv6, &mut buf);
651    let read_iv6 = read_inet(&mut &*buf).unwrap();
652    assert_eq!(iv6, read_iv6);
653}
654
655fn zig_zag_encode(v: i64) -> u64 {
656    ((v >> 63) ^ (v << 1)) as u64
657}
658
659fn zig_zag_decode(v: u64) -> i64 {
660    ((v >> 1) as i64) ^ -((v & 1) as i64)
661}
662
663pub(crate) fn unsigned_vint_encode(v: u64, buf: &mut Vec<u8>) {
664    let mut v = v;
665    let mut number_of_bytes = (639 - 9 * v.leading_zeros()) >> 6;
666    if number_of_bytes <= 1 {
667        return buf.put_u8(v as u8);
668    }
669
670    if number_of_bytes != 9 {
671        let extra_bytes = number_of_bytes - 1;
672        let length_bits = !(0xff >> extra_bytes);
673        v |= (length_bits as u64) << (8 * extra_bytes);
674    } else {
675        buf.put_u8(0xff);
676        number_of_bytes -= 1;
677    }
678    buf.put_uint(v, number_of_bytes as usize)
679}
680
681pub(crate) fn unsigned_vint_decode(buf: &mut &[u8]) -> Result<u64, std::io::Error> {
682    let first_byte = buf.read_u8()?;
683    let extra_bytes = first_byte.leading_ones() as usize;
684
685    let mut v = if extra_bytes != 8 {
686        let first_byte_bits = first_byte & (0xffu8 >> extra_bytes);
687        (first_byte_bits as u64) << (8 * extra_bytes)
688    } else {
689        0
690    };
691
692    if extra_bytes != 0 {
693        v += buf.read_uint::<BigEndian>(extra_bytes)?;
694    }
695
696    Ok(v)
697}
698
699pub(crate) fn vint_encode(v: i64, buf: &mut Vec<u8>) {
700    unsigned_vint_encode(zig_zag_encode(v), buf)
701}
702
703pub(crate) fn vint_decode(buf: &mut &[u8]) -> Result<i64, std::io::Error> {
704    unsigned_vint_decode(buf).map(zig_zag_decode)
705}
706
707#[test]
708fn zig_zag_encode_test() {
709    assert_eq!(zig_zag_encode(0), 0);
710    assert_eq!(zig_zag_encode(-1), 1);
711    assert_eq!(zig_zag_encode(1), 2);
712    assert_eq!(zig_zag_encode(-2), 3);
713    assert_eq!(zig_zag_encode(2), 4);
714    assert_eq!(zig_zag_encode(-3), 5);
715    assert_eq!(zig_zag_encode(3), 6);
716}
717
718#[test]
719fn zig_zag_decode_test() {
720    assert_eq!(zig_zag_decode(0), 0);
721    assert_eq!(zig_zag_decode(1), -1);
722    assert_eq!(zig_zag_decode(2), 1);
723    assert_eq!(zig_zag_decode(3), -2);
724    assert_eq!(zig_zag_decode(4), 2);
725    assert_eq!(zig_zag_decode(5), -3);
726    assert_eq!(zig_zag_decode(6), 3);
727}
728
729#[test]
730fn unsigned_vint_encode_and_decode_test() {
731    let unsigned_vint_encoding = vec![
732        (0, vec![0]),
733        (1, vec![1]),
734        (2, vec![2]),
735        ((1 << 2) - 1, vec![3]),
736        (1 << 2, vec![4]),
737        ((1 << 2) + 1, vec![5]),
738        ((1 << 3) - 1, vec![7]),
739        (1 << 3, vec![8]),
740        ((1 << 3) + 1, vec![9]),
741        ((1 << 4) - 1, vec![15]),
742        (1 << 4, vec![16]),
743        ((1 << 4) + 1, vec![17]),
744        ((1 << 5) - 1, vec![31]),
745        (1 << 5, vec![32]),
746        ((1 << 5) + 1, vec![33]),
747        ((1 << 6) - 1, vec![63]),
748        (1 << 6, vec![64]),
749        ((1 << 6) + 1, vec![65]),
750        ((1 << 7) - 1, vec![127]),
751        (1 << 7, vec![128, 128]),
752        ((1 << 7) + 1, vec![128, 129]),
753        ((1 << 8) - 1, vec![128, 255]),
754        (1 << 8, vec![129, 0]),
755        ((1 << 8) + 1, vec![129, 1]),
756        ((1 << 9) - 1, vec![129, 255]),
757        (1 << 9, vec![130, 0]),
758        ((1 << 9) + 1, vec![130, 1]),
759        ((1 << 10) - 1, vec![131, 255]),
760        (1 << 10, vec![132, 0]),
761        ((1 << 10) + 1, vec![132, 1]),
762        ((1 << 11) - 1, vec![135, 255]),
763        (1 << 11, vec![136, 0]),
764        ((1 << 11) + 1, vec![136, 1]),
765        ((1 << 12) - 1, vec![143, 255]),
766        (1 << 12, vec![144, 0]),
767        ((1 << 12) + 1, vec![144, 1]),
768        ((1 << 13) - 1, vec![159, 255]),
769        (1 << 13, vec![160, 0]),
770        ((1 << 13) + 1, vec![160, 1]),
771        ((1 << 14) - 1, vec![191, 255]),
772        (1 << 14, vec![192, 64, 0]),
773        ((1 << 14) + 1, vec![192, 64, 1]),
774        ((1 << 15) - 1, vec![192, 127, 255]),
775        (1 << 15, vec![192, 128, 0]),
776        ((1 << 15) + 1, vec![192, 128, 1]),
777        ((1 << 16) - 1, vec![192, 255, 255]),
778        (1 << 16, vec![193, 0, 0]),
779        ((1 << 16) + 1, vec![193, 0, 1]),
780        ((1 << 17) - 1, vec![193, 255, 255]),
781        (1 << 17, vec![194, 0, 0]),
782        ((1 << 17) + 1, vec![194, 0, 1]),
783        ((1 << 18) - 1, vec![195, 255, 255]),
784        (1 << 18, vec![196, 0, 0]),
785        ((1 << 18) + 1, vec![196, 0, 1]),
786        ((1 << 19) - 1, vec![199, 255, 255]),
787        (1 << 19, vec![200, 0, 0]),
788        ((1 << 19) + 1, vec![200, 0, 1]),
789        ((1 << 20) - 1, vec![207, 255, 255]),
790        (1 << 20, vec![208, 0, 0]),
791        ((1 << 20) + 1, vec![208, 0, 1]),
792        ((1 << 21) - 1, vec![223, 255, 255]),
793        (1 << 21, vec![224, 32, 0, 0]),
794        ((1 << 21) + 1, vec![224, 32, 0, 1]),
795        ((1 << 22) - 1, vec![224, 63, 255, 255]),
796        (1 << 22, vec![224, 64, 0, 0]),
797        ((1 << 22) + 1, vec![224, 64, 0, 1]),
798        ((1 << 23) - 1, vec![224, 127, 255, 255]),
799        (1 << 23, vec![224, 128, 0, 0]),
800        ((1 << 23) + 1, vec![224, 128, 0, 1]),
801        ((1 << 24) - 1, vec![224, 255, 255, 255]),
802        (1 << 24, vec![225, 0, 0, 0]),
803        ((1 << 24) + 1, vec![225, 0, 0, 1]),
804        ((1 << 25) - 1, vec![225, 255, 255, 255]),
805        (1 << 25, vec![226, 0, 0, 0]),
806        ((1 << 25) + 1, vec![226, 0, 0, 1]),
807        ((1 << 26) - 1, vec![227, 255, 255, 255]),
808        (1 << 26, vec![228, 0, 0, 0]),
809        ((1 << 26) + 1, vec![228, 0, 0, 1]),
810        ((1 << 27) - 1, vec![231, 255, 255, 255]),
811        (1 << 27, vec![232, 0, 0, 0]),
812        ((1 << 27) + 1, vec![232, 0, 0, 1]),
813        ((1 << 28) - 1, vec![239, 255, 255, 255]),
814        (1 << 28, vec![240, 16, 0, 0, 0]),
815        ((1 << 28) + 1, vec![240, 16, 0, 0, 1]),
816        ((1 << 29) - 1, vec![240, 31, 255, 255, 255]),
817        (1 << 29, vec![240, 32, 0, 0, 0]),
818        ((1 << 29) + 1, vec![240, 32, 0, 0, 1]),
819        ((1 << 30) - 1, vec![240, 63, 255, 255, 255]),
820        (1 << 30, vec![240, 64, 0, 0, 0]),
821        ((1 << 30) + 1, vec![240, 64, 0, 0, 1]),
822        ((1 << 31) - 1, vec![240, 127, 255, 255, 255]),
823        (1 << 31, vec![240, 128, 0, 0, 0]),
824        ((1 << 31) + 1, vec![240, 128, 0, 0, 1]),
825        ((1 << 32) - 1, vec![240, 255, 255, 255, 255]),
826        (1 << 32, vec![241, 0, 0, 0, 0]),
827        ((1 << 32) + 1, vec![241, 0, 0, 0, 1]),
828        ((1 << 33) - 1, vec![241, 255, 255, 255, 255]),
829        (1 << 33, vec![242, 0, 0, 0, 0]),
830        ((1 << 33) + 1, vec![242, 0, 0, 0, 1]),
831        ((1 << 34) - 1, vec![243, 255, 255, 255, 255]),
832        (1 << 34, vec![244, 0, 0, 0, 0]),
833        ((1 << 34) + 1, vec![244, 0, 0, 0, 1]),
834        ((1 << 35) - 1, vec![247, 255, 255, 255, 255]),
835        (1 << 35, vec![248, 8, 0, 0, 0, 0]),
836        ((1 << 35) + 1, vec![248, 8, 0, 0, 0, 1]),
837        ((1 << 36) - 1, vec![248, 15, 255, 255, 255, 255]),
838        (1 << 36, vec![248, 16, 0, 0, 0, 0]),
839        ((1 << 36) + 1, vec![248, 16, 0, 0, 0, 1]),
840        ((1 << 37) - 1, vec![248, 31, 255, 255, 255, 255]),
841        (1 << 37, vec![248, 32, 0, 0, 0, 0]),
842        ((1 << 37) + 1, vec![248, 32, 0, 0, 0, 1]),
843        ((1 << 38) - 1, vec![248, 63, 255, 255, 255, 255]),
844        (1 << 38, vec![248, 64, 0, 0, 0, 0]),
845        ((1 << 38) + 1, vec![248, 64, 0, 0, 0, 1]),
846        ((1 << 39) - 1, vec![248, 127, 255, 255, 255, 255]),
847        (1 << 39, vec![248, 128, 0, 0, 0, 0]),
848        ((1 << 39) + 1, vec![248, 128, 0, 0, 0, 1]),
849        ((1 << 40) - 1, vec![248, 255, 255, 255, 255, 255]),
850        (1 << 40, vec![249, 0, 0, 0, 0, 0]),
851        ((1 << 40) + 1, vec![249, 0, 0, 0, 0, 1]),
852        ((1 << 41) - 1, vec![249, 255, 255, 255, 255, 255]),
853        (1 << 41, vec![250, 0, 0, 0, 0, 0]),
854        ((1 << 41) + 1, vec![250, 0, 0, 0, 0, 1]),
855        ((1 << 42) - 1, vec![251, 255, 255, 255, 255, 255]),
856        (1 << 42, vec![252, 4, 0, 0, 0, 0, 0]),
857        ((1 << 42) + 1, vec![252, 4, 0, 0, 0, 0, 1]),
858        ((1 << 43) - 1, vec![252, 7, 255, 255, 255, 255, 255]),
859        (1 << 43, vec![252, 8, 0, 0, 0, 0, 0]),
860        ((1 << 43) + 1, vec![252, 8, 0, 0, 0, 0, 1]),
861        ((1 << 44) - 1, vec![252, 15, 255, 255, 255, 255, 255]),
862        (1 << 44, vec![252, 16, 0, 0, 0, 0, 0]),
863        ((1 << 44) + 1, vec![252, 16, 0, 0, 0, 0, 1]),
864        ((1 << 45) - 1, vec![252, 31, 255, 255, 255, 255, 255]),
865        (1 << 45, vec![252, 32, 0, 0, 0, 0, 0]),
866        ((1 << 45) + 1, vec![252, 32, 0, 0, 0, 0, 1]),
867        ((1 << 46) - 1, vec![252, 63, 255, 255, 255, 255, 255]),
868        (1 << 46, vec![252, 64, 0, 0, 0, 0, 0]),
869        ((1 << 46) + 1, vec![252, 64, 0, 0, 0, 0, 1]),
870        ((1 << 47) - 1, vec![252, 127, 255, 255, 255, 255, 255]),
871        (1 << 47, vec![252, 128, 0, 0, 0, 0, 0]),
872        ((1 << 47) + 1, vec![252, 128, 0, 0, 0, 0, 1]),
873        ((1 << 48) - 1, vec![252, 255, 255, 255, 255, 255, 255]),
874        (1 << 48, vec![253, 0, 0, 0, 0, 0, 0]),
875        ((1 << 48) + 1, vec![253, 0, 0, 0, 0, 0, 1]),
876        ((1 << 49) - 1, vec![253, 255, 255, 255, 255, 255, 255]),
877        (1 << 49, vec![254, 2, 0, 0, 0, 0, 0, 0]),
878        ((1 << 49) + 1, vec![254, 2, 0, 0, 0, 0, 0, 1]),
879        ((1 << 50) - 1, vec![254, 3, 255, 255, 255, 255, 255, 255]),
880        (1 << 50, vec![254, 4, 0, 0, 0, 0, 0, 0]),
881        ((1 << 50) + 1, vec![254, 4, 0, 0, 0, 0, 0, 1]),
882        ((1 << 51) - 1, vec![254, 7, 255, 255, 255, 255, 255, 255]),
883        (1 << 51, vec![254, 8, 0, 0, 0, 0, 0, 0]),
884        ((1 << 51) + 1, vec![254, 8, 0, 0, 0, 0, 0, 1]),
885        ((1 << 52) - 1, vec![254, 15, 255, 255, 255, 255, 255, 255]),
886        (1 << 52, vec![254, 16, 0, 0, 0, 0, 0, 0]),
887        ((1 << 52) + 1, vec![254, 16, 0, 0, 0, 0, 0, 1]),
888        ((1 << 53) - 1, vec![254, 31, 255, 255, 255, 255, 255, 255]),
889        (1 << 53, vec![254, 32, 0, 0, 0, 0, 0, 0]),
890        ((1 << 53) + 1, vec![254, 32, 0, 0, 0, 0, 0, 1]),
891        ((1 << 54) - 1, vec![254, 63, 255, 255, 255, 255, 255, 255]),
892        (1 << 54, vec![254, 64, 0, 0, 0, 0, 0, 0]),
893        ((1 << 54) + 1, vec![254, 64, 0, 0, 0, 0, 0, 1]),
894        ((1 << 55) - 1, vec![254, 127, 255, 255, 255, 255, 255, 255]),
895        (1 << 55, vec![254, 128, 0, 0, 0, 0, 0, 0]),
896        ((1 << 55) + 1, vec![254, 128, 0, 0, 0, 0, 0, 1]),
897        ((1 << 56) - 1, vec![254, 255, 255, 255, 255, 255, 255, 255]),
898        (1 << 56, vec![255, 1, 0, 0, 0, 0, 0, 0, 0]),
899        ((1 << 56) + 1, vec![255, 1, 0, 0, 0, 0, 0, 0, 1]),
900        (
901            (1 << 57) - 1,
902            vec![255, 1, 255, 255, 255, 255, 255, 255, 255],
903        ),
904        (1 << 57, vec![255, 2, 0, 0, 0, 0, 0, 0, 0]),
905        ((1 << 57) + 1, vec![255, 2, 0, 0, 0, 0, 0, 0, 1]),
906        (
907            (1 << 58) - 1,
908            vec![255, 3, 255, 255, 255, 255, 255, 255, 255],
909        ),
910        (1 << 58, vec![255, 4, 0, 0, 0, 0, 0, 0, 0]),
911        ((1 << 58) + 1, vec![255, 4, 0, 0, 0, 0, 0, 0, 1]),
912        (
913            (1 << 59) - 1,
914            vec![255, 7, 255, 255, 255, 255, 255, 255, 255],
915        ),
916        (1 << 59, vec![255, 8, 0, 0, 0, 0, 0, 0, 0]),
917        ((1 << 59) + 1, vec![255, 8, 0, 0, 0, 0, 0, 0, 1]),
918        (
919            (1 << 60) - 1,
920            vec![255, 15, 255, 255, 255, 255, 255, 255, 255],
921        ),
922        (1 << 60, vec![255, 16, 0, 0, 0, 0, 0, 0, 0]),
923        ((1 << 60) + 1, vec![255, 16, 0, 0, 0, 0, 0, 0, 1]),
924        (
925            (1 << 61) - 1,
926            vec![255, 31, 255, 255, 255, 255, 255, 255, 255],
927        ),
928        (1 << 61, vec![255, 32, 0, 0, 0, 0, 0, 0, 0]),
929        ((1 << 61) + 1, vec![255, 32, 0, 0, 0, 0, 0, 0, 1]),
930        (
931            (1 << 62) - 1,
932            vec![255, 63, 255, 255, 255, 255, 255, 255, 255],
933        ),
934        (1 << 62, vec![255, 64, 0, 0, 0, 0, 0, 0, 0]),
935        ((1 << 62) + 1, vec![255, 64, 0, 0, 0, 0, 0, 0, 1]),
936        (
937            (1 << 63) - 1,
938            vec![255, 127, 255, 255, 255, 255, 255, 255, 255],
939        ),
940        (1 << 63, vec![255, 128, 0, 0, 0, 0, 0, 0, 0]),
941        ((1 << 63) + 1, vec![255, 128, 0, 0, 0, 0, 0, 0, 1]),
942        (u64::MAX, vec![255, 255, 255, 255, 255, 255, 255, 255, 255]),
943    ];
944
945    let mut buf = Vec::new();
946
947    for (v, result) in unsigned_vint_encoding.into_iter() {
948        unsigned_vint_encode(v, &mut buf);
949        assert_eq!(buf, result);
950        let decoded_v = unsigned_vint_decode(&mut buf.as_slice()).unwrap();
951        assert_eq!(v, decoded_v);
952        buf.clear();
953    }
954}
955
956#[test]
957fn vint_encode_and_decode_test() {
958    let mut buf: Vec<u8> = Vec::with_capacity(128);
959
960    let mut check = |n: i64| {
961        vint_encode(n, &mut buf);
962        assert_eq!(vint_decode(&mut buf.as_slice()).unwrap(), n);
963        buf.clear();
964    };
965
966    for i in 0..63 {
967        check((1 << i) - 1);
968        check(1 - (1 << i));
969        check(1 << i);
970        check(-(1 << i));
971        check((1 << i) + 1);
972        check(-1 - (1 << i));
973    }
974    check(i64::MAX);
975    check(-i64::MAX);
976    check(i64::MIN)
977}