From c6d703f1673a9b2481a0559aa817afc8b0b3dd9e Mon Sep 17 00:00:00 2001 From: OMGeeky Date: Thu, 14 Nov 2024 18:01:31 +0100 Subject: [PATCH] fix var_int and add var_long --- src/types.rs | 1 + src/types/var_int.rs | 166 +++++++++++++++++++++++++++++------------- src/types/var_long.rs | 138 +++++++++++++++++++++++++++++++++++ 3 files changed, 256 insertions(+), 49 deletions(-) create mode 100644 src/types/var_long.rs diff --git a/src/types.rs b/src/types.rs index 5d95333..58101fb 100644 --- a/src/types.rs +++ b/src/types.rs @@ -23,3 +23,4 @@ pub trait McRustRepr { } pub mod string; pub mod var_int; +pub mod var_long; diff --git a/src/types/var_int.rs b/src/types/var_int.rs index 76093a3..efc544a 100644 --- a/src/types/var_int.rs +++ b/src/types/var_int.rs @@ -1,5 +1,6 @@ use crate::types::{McRead, McRustRepr, McWrite}; +use num_traits::ToPrimitive; use std::ops::Deref; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; @@ -22,63 +23,16 @@ impl McWrite for VarInt { where Self: Sized, { - let mut value = self.0 as u32; - loop { - if (value & Self::SEGMENT_BITS as u32) == 0 { - let _ = stream.write(&[value.to_le_bytes()[0]]).await?; - return Ok(1); - } - - let x = value & Self::SEGMENT_BITS as u32 | Self::CONTINUE_BIT as u32; - let x = x.to_le_bytes()[0]; - let _ = stream.write(&[x]).await?; - value >>= 7; - } + write_varint(stream, self.0).await } } impl McRead for VarInt { type Error = String; async fn read_stream(b: &mut T) -> Result { - let mut value = 0i32; - let mut position = 0; - // println!("CONTINUE bit: {:0>32b}", Self::CONTINUE_BIT); - // println!("SEGMENT bit: {:0>32b}", Self::SEGMENT_BITS); - - loop { - let mut current_byte = 0u8; - b.read_exact(std::slice::from_mut(&mut current_byte)) - .await - .map_err(|x| x.to_string())?; - // println!( - // "b: {:0>32b}\nm: {:0>32b}\nr: {:0>32b}\n>: {:0>32b} ({position})\nv: {:0>32b}\nr2:{:0>32b}", - // current_byte, - // Self::SEGMENT_BITS, - // current_byte & Self::SEGMENT_BITS, - // ((current_byte & Self::SEGMENT_BITS) as i32) << position, - // value, - // value | (((current_byte & Self::SEGMENT_BITS) as i32) << position) - // ); - value |= ((current_byte & Self::SEGMENT_BITS) as i32) << position; - // println!("{x}:\n{current_byte:>32b}:\n{value:>32b}:\n{value}"); - - if (current_byte & Self::CONTINUE_BIT) == 0 { - break; - } - - position += 7; - - if position >= 32 { - return Err("VarInt is too big".to_string()); - } - } - + let value = read_stream(b, 32).await? as i32; Ok(Self(value)) } } -impl VarInt { - const SEGMENT_BITS: u8 = 0x7F; - const CONTINUE_BIT: u8 = 0x80; -} impl McRustRepr for VarInt { type RustRepresentation = i32; @@ -94,3 +48,117 @@ impl McRustRepr for VarInt { &self.0 } } + +pub(crate) async fn read_stream( + b: &mut T, + max_size: usize, +) -> Result { + let mut value = 0u64; + let mut position = 0; + + loop { + let current_byte = b.read_u8().await.map_err(|x| x.to_string())?; + value |= ((current_byte & SEGMENT_BITS_U8) as u64) << position; + if (current_byte & CONTINUE_BIT_U8) == 0 { + break; + } + + position += 7; + + if position >= max_size { + return Err("Number is too big".to_string()); + } + } + + Ok(value) +} + +async fn write_varint( + writer: &mut W, + mut value: i32, +) -> std::io::Result { + let mut written = 0; + loop { + written += 1; + if (value & !SEGMENT_BITS) == 0 { + writer.write_u8(value as u8).await.unwrap(); + break; + } + + writer + .write_u8(((value & SEGMENT_BITS) | CONTINUE_BIT) as u8) + .await + .unwrap(); + + value = (value as u32 >> 7) as i32; // Simulate Java's >>> + } + Ok(written) +} + +pub(crate) fn get_number_bytes_size(mut value: u64) -> usize { + let mut size = 1; + + while value >= CONTINUE_BIT_U8 as u64 { + value >>= 7; + size += 1; + } + + size +} + +impl VarInt { + pub(crate) fn get_size(&self) -> usize { + let value = self.0 as u64; // Convert to u64 for bitwise operations + get_number_bytes_size(value) + } +} + +const SEGMENT_BITS: i32 = 0x7F; +const CONTINUE_BIT: i32 = 0x80; +const SEGMENT_BITS_U8: u8 = SEGMENT_BITS as u8; +const CONTINUE_BIT_U8: u8 = CONTINUE_BIT as u8; + +#[cfg(test)] +mod tests { + use crate::types::var_int::VarInt; + use crate::types::{McRead, McWrite}; + use tokio::io::{AsyncWriteExt, BufReader, BufWriter}; + + #[tokio::test] + async fn test_varint_read_stream() { + let test_cases = [ + (vec![0x00], 0), + (vec![0x01], 1), + (vec![0x02], 2), + (vec![0x7f], 127), + (vec![0x80, 0x01], 128), + (vec![0xff, 0x01], 255), + (vec![0xdd, 0xc7, 0x01], 25565), + (vec![0xff, 0xff, 0x7f], 2097151), + (vec![0xff, 0xff, 0xff, 0xff, 0x07], 2147483647), + (vec![0xff, 0xff, 0xff, 0xff, 0x0f], -1), + (vec![0x80, 0x80, 0x80, 0x80, 0x08], -2147483648), + ]; + + for (input_bytes, expected_value) in test_cases { + let mut reader = BufReader::new(input_bytes.as_slice()); + let varint = VarInt::read_stream(&mut reader).await.unwrap(); + assert_eq!( + varint.0, expected_value, + "Decoding mismatch for bytes: {:?}", + input_bytes + ); + let mut reversed_bytes = Vec::new(); + + let mut writer = BufWriter::new(&mut reversed_bytes); + varint.write_stream(&mut writer).await.unwrap(); + writer.flush().await.unwrap(); + + assert_eq!( + input_bytes, reversed_bytes, + "Back and forth did not work for {}", + *varint + ) + } + } +} diff --git a/src/types/var_long.rs b/src/types/var_long.rs new file mode 100644 index 0000000..2bce607 --- /dev/null +++ b/src/types/var_long.rs @@ -0,0 +1,138 @@ +use crate::types::{McRead, McRustRepr, McWrite}; + +use std::ops::Deref; +use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt}; + +#[derive(Debug, Copy, Clone)] +pub struct VarLong(pub i64); +impl Deref for VarLong { + type Target = i64; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl VarLong { + pub(crate) fn get_size(&self) -> usize { + let value = self.0 as u64; // Convert to u64 for bitwise operations + crate::types::var_int::get_number_bytes_size(value) + } +} +impl McWrite for VarLong { + type Error = std::io::Error; + + async fn write_stream( + &self, + stream: &mut T, + ) -> Result + where + Self: Sized, + { + let value = self.0 as u64; + write_varlong(stream, self.0).await + } +} +impl McRead for VarLong { + type Error = String; + async fn read_stream(b: &mut T) -> Result { + let value = crate::types::var_int::read_stream(b, 64).await? as i64; + Ok(Self(value)) + } +} +impl McRustRepr for VarLong { + type RustRepresentation = i64; + + fn into_rs(self) -> Self::RustRepresentation { + self.0 + } + + fn to_rs(&self) -> Self::RustRepresentation { + self.0 + } + + fn as_rs(&self) -> &Self::RustRepresentation { + &self.0 + } +} + +async fn write_varlong( + writer: &mut W, + mut value: i64, +) -> std::io::Result { + let mut written = 0; + loop { + written += 1; + if (value & !SEGMENT_BITS) == 0 { + writer.write_u8(value as u8).await.unwrap(); + break; + } + + writer + .write_u8(((value & SEGMENT_BITS) | CONTINUE_BIT) as u8) + .await + .unwrap(); + + value = (value as u64 >> 7) as i64; // Simulate Java's >>> + } + Ok(written) +} + +const SEGMENT_BITS: i64 = 0x7F; +const CONTINUE_BIT: i64 = 0x80; + +#[cfg(test)] +mod tests { + use crate::types::var_long::VarLong; + use crate::types::{McRead, McWrite}; + use tokio::io::{AsyncWriteExt, BufReader, BufWriter}; + use tokio_util::bytes::BufMut; + + #[tokio::test] + async fn test_varlong_read_stream() { + let test_cases = [ + (vec![0x00], 0), + (vec![0x01], 1), + (vec![0x02], 2), + (vec![0x7f], 127), + (vec![0x80, 0x01], 128), + (vec![0xff, 0x01], 255), + (vec![0xff, 0xff, 0xff, 0xff, 0x07], 2147483647), + ( + vec![0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x7f], + 9223372036854775807, + ), + ( + vec![0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x01], + -1, + ), + ( + vec![0x80, 0x80, 0x80, 0x80, 0xf8, 0xff, 0xff, 0xff, 0xff, 0x01], + -2147483648, + ), + ( + vec![0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x01], + -9223372036854775808, + ), + ]; + + for (input_bytes, expected_value) in test_cases { + let mut reader = BufReader::new(input_bytes.as_slice()); + let varint = VarLong::read_stream(&mut reader).await.unwrap(); + assert_eq!( + *varint, expected_value, + "Decoding mismatch for bytes: {:?}", + input_bytes + ); + let mut reversed_bytes = Vec::new(); + let mut writer = BufWriter::new(&mut reversed_bytes); + varint.write_stream(&mut writer).await.unwrap(); + writer.flush().await.unwrap(); + assert_eq!( + input_bytes, reversed_bytes, + "Back and forth did not work for {}", + *varint + ) + } + } +}