From 336c7c9dd6a80578f7e97e03b444e9557b7d8252 Mon Sep 17 00:00:00 2001 From: Adam Ierymenko Date: Fri, 29 Sep 2023 18:14:16 -0400 Subject: [PATCH] A few additions to buf, errors, and restore varint from attic. --- src/buf.rs | 10 ++++- src/error.rs | 17 +++++++++ src/lib.rs | 1 + src/varint.rs | 102 ++++++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 128 insertions(+), 2 deletions(-) create mode 100644 src/varint.rs diff --git a/src/buf.rs b/src/buf.rs index ad41b4b..7a1be5a 100644 --- a/src/buf.rs +++ b/src/buf.rs @@ -46,6 +46,11 @@ impl Buf { } } + #[inline(always)] + pub fn iter(&self) -> impl Iterator { + self.as_slice().iter() + } + #[inline(always)] pub fn len(&self) -> usize { unsafe { *self.0.as_ptr() as usize } @@ -90,11 +95,12 @@ impl Buf { #[inline] #[must_use] pub fn append(&mut self, buf: &[u8]) -> bool { - let new_len = self.len() + buf.len(); + let old_len = self.len(); + let new_len = old_len + buf.len(); if new_len <= self.capacity() { unsafe { *self.0.as_ptr() = new_len as u32; - copy_nonoverlapping(buf.as_ptr(), self.0.as_ptr().cast::().add(8), buf.len()) + copy_nonoverlapping(buf.as_ptr(), self.0.as_ptr().cast::().add(8 + old_len), buf.len()) }; true } else { diff --git a/src/error.rs b/src/error.rs index 723099f..64f8d29 100644 --- a/src/error.rs +++ b/src/error.rs @@ -9,6 +9,23 @@ use std::error::Error; use std::fmt::{Debug, Display}; +pub struct InvalidStateError; + +impl Display for InvalidStateError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("InvalidStateError") + } +} + +impl Debug for InvalidStateError { + #[inline(always)] + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + ::fmt(self, f) + } +} + +impl Error for InvalidStateError {} + pub struct InvalidFormatError; impl Display for InvalidFormatError { diff --git a/src/lib.rs b/src/lib.rs index 7e3fa8b..572d722 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -23,6 +23,7 @@ pub mod ringbuffer; pub mod str; pub mod sync; pub mod tofrombytes; +pub mod varint; /// Initial value that should be used for monotonic tick time variables. pub const NEVER_HAPPENED_TICKS: i64 = i64::MIN / 2; diff --git a/src/varint.rs b/src/varint.rs new file mode 100644 index 0000000..fa51417 --- /dev/null +++ b/src/varint.rs @@ -0,0 +1,102 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +use std::io::{Read, Write}; + +pub const VARINT_MAX_SIZE_BYTES: usize = 10; + +/// Encode an integer as a varint. +/// +/// WARNING: if the supplied byte slice does not have at least 10 bytes available this may panic. +/// This is checked in debug mode by an assertion. +#[inline] +pub fn encode(b: &mut [u8], mut v: u64) -> usize { + debug_assert!(b.len() >= VARINT_MAX_SIZE_BYTES); + let mut i = 0; + loop { + if v > 0x7f { + b[i] = (v as u8) & 0x7f; + i += 1; + v = v.wrapping_shr(7); + } else { + b[i] = (v as u8) | 0x80; + i += 1; + break; + } + } + i +} + +/// Write a variable length integer, which can consume up to 10 bytes. +#[inline] +pub fn write(w: &mut W, v: u64) -> std::io::Result<()> { + let mut b = [0_u8; VARINT_MAX_SIZE_BYTES]; + let i = encode(&mut b, v); + w.write_all(&b[0..i]) +} + +/// Dencode up to 10 bytes as a varint. +/// +/// if the supplied byte slice does not contain a valid varint encoding this will return None. +/// if the supplied byte slice is shorter than expected this will return None. +#[inline] +pub fn decode(b: &[u8]) -> Option<(u64, usize)> { + let mut v = 0_u64; + let mut pos = 0; + let mut i = 0_usize; + while i < b.len() && i < VARINT_MAX_SIZE_BYTES { + let b = b[i]; + i += 1; + if b <= 0x7f { + v |= (b as u64).wrapping_shl(pos); + pos += 7; + } else { + v |= ((b & 0x7f) as u64).wrapping_shl(pos); + return Some((v, i)); + } + } + None +} + +/// Read a variable length integer, returning the value and the number of bytes written. +#[inline] +pub fn read(r: &mut R) -> std::io::Result<(u64, usize)> { + let mut v = 0_u64; + let mut buf = [0_u8; 1]; + let mut pos = 0; + let mut i = 0_usize; + loop { + r.read_exact(&mut buf)?; + let b = buf[0]; + i += 1; + if b <= 0x7f { + v |= (b as u64).wrapping_shl(pos); + pos += 7; + } else { + v |= ((b & 0x7f) as u64).wrapping_shl(pos); + return Ok((v, i)); + } + } +} + +#[cfg(test)] +mod tests { + use crate::varint::*; + + #[test] + fn varint() { + let mut t: Vec = Vec::new(); + for i in 0..131072 { + t.clear(); + let ii = (u64::MAX / 131072) * i; + assert!(write(&mut t, ii).is_ok()); + let mut t2 = t.as_slice(); + assert_eq!(read(&mut t2).unwrap().0, ii); + } + } +}