Skip to content

Commit

Permalink
feat: add binary reader to improve reading of the binary nbt
Browse files Browse the repository at this point in the history
  • Loading branch information
Quozul committed Aug 15, 2024
1 parent 3c99e4f commit fb72067
Show file tree
Hide file tree
Showing 11 changed files with 281 additions and 262 deletions.
116 changes: 116 additions & 0 deletions src/nbt/binary_reader.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,116 @@
use std::string::FromUtf8Error;

macro_rules! impl_read_number {
($fn_name:ident, $type:ty) => {
pub fn $fn_name(&mut self) -> $type {
let size = std::mem::size_of::<$type>();
let bytes = &self.raw[self.index..self.index + size];
let integer = <$type>::from_be_bytes(bytes.try_into().unwrap());
self.index += size;
integer
}
};
}

macro_rules! impl_read_array {
($fn_name:ident, $type:ty, $reader:ident) => {
pub fn $fn_name(&mut self) -> Vec<$type> {
let size = self.read_i32();
let mut values = Vec::new();

for _ in 0..size {
let next_tag = self.$reader();
values.push(next_tag);
}

values
}
};
}

pub struct BinaryReader<'a> {
raw: &'a [u8],
index: usize,
}

impl<'a> BinaryReader<'a> {
pub fn new(raw: &'a [u8]) -> Self {
Self { raw, index: 0 }
}

pub fn read_string(&mut self) -> Result<String, FromUtf8Error> {
let size = self.read_u16() as usize;
let bytes = &self.raw[self.index..self.index + size];
self.index += size;
String::from_utf8(Vec::from(bytes))
}

pub fn read_name(&mut self) -> Option<String> {
self.read_string().ok().filter(|s| !s.is_empty())
}

pub fn read_type(&mut self) -> u8 {
self.read_u8()
}

impl_read_number!(read_i8, i8);
impl_read_number!(read_u8, u8);
impl_read_number!(read_i16, i16);
impl_read_number!(read_u16, u16);
impl_read_number!(read_i32, i32);
impl_read_number!(read_i64, i64);
impl_read_number!(read_f32, f32);
impl_read_number!(read_f64, f64);
impl_read_array!(read_byte_array, i8, read_i8);
impl_read_array!(read_int_array, i32, read_i32);
impl_read_array!(read_long_array, i64, read_i64);
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn test_read_i8() {
let data = [0x7F];
let mut reader = BinaryReader::new(&data);
assert_eq!(reader.read_i8(), 127);
}

#[test]
fn test_read_i16() {
let data = [0x7F, 0xFF];
let mut reader = BinaryReader::new(&data);
assert_eq!(reader.read_i16(), 32767);
}

#[test]
fn test_read_u16() {
let data = [0x0F, 0xFF];
let mut reader = BinaryReader::new(&data);
assert_eq!(reader.read_u16(), 4095);
}

#[test]
fn test_read_i32() {
let data = [0x7F, 0xFF, 0xFF, 0xFF];
let mut reader = BinaryReader::new(&data);
assert_eq!(reader.read_i32(), 2147483647);
}

#[test]
fn test_read_f32() {
let data = [0x3F, 0x80, 0x00, 0x00];
let mut reader = BinaryReader::new(&data);
assert_eq!(reader.read_f32(), 1.0);
}

#[test]
fn test_read_string() {
let data = [0, 5, 72, 69, 76, 76, 79];
let mut reader = BinaryReader::new(&data);
let parsed = reader.read_string().unwrap();

assert_eq!(parsed, "HELLO");
}
}
3 changes: 2 additions & 1 deletion src/nbt/mod.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
pub mod binary_reader;
pub mod parse;
mod readers;
mod parsers;
pub mod tag;
mod writers;
128 changes: 9 additions & 119 deletions src/nbt/parse.rs
Original file line number Diff line number Diff line change
@@ -1,119 +1,10 @@
use crate::nbt::readers::{
read_f32, read_f64, read_i16, read_i32, read_i64, read_i8, read_name, read_string, read_type,
};
use crate::nbt::binary_reader::BinaryReader;
use crate::nbt::parsers::parse_with_type::parse_with_type;
use crate::nbt::tag::Tag;

pub fn parse_tag(raw: &[u8], index: &mut usize) -> Tag {
let tag_type = read_type(raw, index);
parse_of_type(raw, index, tag_type, false)
}

fn parse_of_type(raw: &[u8], index: &mut usize, tag_type: u8, skip_name: bool) -> Tag {
let name = if skip_name || tag_type == 0 {
None
} else {
read_name(raw, index)
};

match tag_type {
0 => Tag::End,
1 => {
let value = read_i8(raw, index);
Tag::Byte { name, value }
}
2 => {
let value = read_i16(raw, index);
Tag::Short { name, value }
}
3 => {
let value = read_i32(raw, index);
Tag::Int { name, value }
}
4 => {
let value = read_i64(raw, index);
Tag::Long { name, value }
}
5 => {
let value = read_f32(raw, index);
Tag::Float { name, value }
}
6 => {
let value = read_f64(raw, index);
Tag::Double { name, value }
}
7 => {
let value = parse_array(raw, index, read_i8);
Tag::ByteArray { name, value }
}
8 => {
let value = read_string(raw, index).unwrap_or_default();
Tag::String { name, value }
}
9 => {
let (tag_type, value) = parse_list_tag(raw, index);
Tag::List {
name,
value,
tag_type,
}
}
10 => {
let value = parse_compound_tag(raw, index);
Tag::Compound { name, value }
}
11 => {
let value = parse_array(raw, index, read_i32);
Tag::IntArray { name, value }
}
12 => {
let value = parse_array(raw, index, read_i64);
Tag::LongArray { name, value }
}
_ => panic!("Unsupported tag type {tag_type}"),
}
}

fn parse_list_tag(raw: &[u8], index: &mut usize) -> (u8, Vec<Tag>) {
let mut values = Vec::new();

let tag_type = read_type(raw, index);
let list_length = read_i32(raw, index);
if list_length <= 0 && tag_type == 0 {
return (tag_type, values);
}

for _ in 0..list_length {
let next_tag = parse_of_type(raw, index, tag_type, true);
values.push(next_tag);
}

(tag_type, values)
}

fn parse_compound_tag(raw: &[u8], index: &mut usize) -> Vec<Tag> {
let mut values = Vec::new();

loop {
let next_tag = parse_tag(raw, index);
if next_tag == Tag::End {
break;
}
values.push(next_tag);
}

values
}

fn parse_array<T>(raw: &[u8], index: &mut usize, parser: fn(&[u8], &mut usize) -> T) -> Vec<T> {
let size = read_i32(raw, index);
let mut values = Vec::new();

for _ in 0..size {
let next_tag = parser(raw, index);
values.push(next_tag);
}

values
pub fn parse_tag(reader: &mut BinaryReader) -> Tag {
let tag_type = reader.read_type();
parse_with_type(reader, tag_type, false)
}

#[cfg(test)]
Expand All @@ -123,8 +14,8 @@ mod tests {
#[test]
fn test_hello_world() {
let data = include_bytes!("../../test_files/hello_world.nbt");
let mut index = 0_usize;
let result = parse_tag(data, &mut index);
let mut reader = BinaryReader::new(data);
let result = parse_tag(&mut reader);

assert_eq!(
result,
Expand All @@ -144,9 +35,8 @@ mod tests {
#[test]
fn test_bigtest() {
let data = include_bytes!("../../test_files/bigtest.nbt");

let mut index = 0_usize;
let result = parse_tag(data, &mut index);
let mut reader = BinaryReader::new(data);
let result = parse_tag(&mut reader);

// Build the ByteArray
let mut value = Vec::new();
Expand Down
3 changes: 3 additions & 0 deletions src/nbt/parsers/mod.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
pub mod parse_compound_tag;
pub mod parse_list_tag;
pub mod parse_with_type;
17 changes: 17 additions & 0 deletions src/nbt/parsers/parse_compound_tag.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
use crate::nbt::binary_reader::BinaryReader;
use crate::nbt::parse::parse_tag;
use crate::nbt::tag::Tag;

pub fn parse_compound_tag(reader: &mut BinaryReader) -> Vec<Tag> {
let mut values = Vec::new();

loop {
let next_tag = parse_tag(reader);
if next_tag == Tag::End {
break;
}
values.push(next_tag);
}

values
}
20 changes: 20 additions & 0 deletions src/nbt/parsers/parse_list_tag.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
use crate::nbt::binary_reader::BinaryReader;
use crate::nbt::parsers::parse_with_type::parse_with_type;
use crate::nbt::tag::Tag;

pub fn parse_list_tag(reader: &mut BinaryReader) -> (u8, Vec<Tag>) {
let mut values = Vec::new();

let tag_type = reader.read_type();
let list_length = reader.read_i32();
if list_length <= 0 && tag_type == 0 {
return (tag_type, values);
}

for _ in 0..list_length {
let next_tag = parse_with_type(reader, tag_type, true);
values.push(next_tag);
}

(tag_type, values)
}
69 changes: 69 additions & 0 deletions src/nbt/parsers/parse_with_type.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,69 @@
use crate::nbt::binary_reader::BinaryReader;
use crate::nbt::parsers::parse_compound_tag::parse_compound_tag;
use crate::nbt::parsers::parse_list_tag::parse_list_tag;
use crate::nbt::tag::Tag;

pub fn parse_with_type(reader: &mut BinaryReader, tag_type: u8, skip_name: bool) -> Tag {
let name = if skip_name || tag_type == 0 {
None
} else {
reader.read_name()
};

match tag_type {
0 => Tag::End,
1 => {
let value = reader.read_i8();
Tag::Byte { name, value }
}
2 => {
let value = reader.read_i16();
Tag::Short { name, value }
}
3 => {
let value = reader.read_i32();
Tag::Int { name, value }
}
4 => {
let value = reader.read_i64();
Tag::Long { name, value }
}
5 => {
let value = reader.read_f32();
Tag::Float { name, value }
}
6 => {
let value = reader.read_f64();
Tag::Double { name, value }
}
7 => {
let value = reader.read_byte_array();
Tag::ByteArray { name, value }
}
8 => {
let value = reader.read_string().unwrap_or_default();
Tag::String { name, value }
}
9 => {
let (tag_type, value) = parse_list_tag(reader);
Tag::List {
name,
value,
tag_type,
}
}
10 => {
let value = parse_compound_tag(reader);
Tag::Compound { name, value }
}
11 => {
let value = reader.read_int_array();
Tag::IntArray { name, value }
}
12 => {
let value = reader.read_long_array();
Tag::LongArray { name, value }
}
_ => panic!("Unsupported tag type {tag_type}"),
}
}
Loading

0 comments on commit fb72067

Please sign in to comment.