mat/simdnbt
1
0
Fork 0
mirror of https://github.com/azalea-rs/simdnbt.git synced 2025-08-02 07:26:04 +00:00
Code Activity Actions

make functions that could cause UB private, and add new functions for reading at simdnbt::borrow and owned

Browse source
This commit is contained in:
mat 2024-05-14 10:16:32 +00:00
parent c83d29855a
commit d7b54cd3f9
16 changed files with 204 additions and 60 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
simdnbt/README.md
View file

@ -18,10 +18,9 @@ The difference is that the "borrow" variant requires you to keep a reference to
```rust,no_run
use std::borrow::Cow;
use std::io::Cursor;
use simdnbt::borrow::Nbt;
fn example(item_bytes: &[u8]) {
let nbt = Nbt::read(&mut Cursor::new(item_bytes))
let nbt = simdnbt::borrow::read(&mut Cursor::new(item_bytes))
.unwrap()
.unwrap();
let skyblock_id: Cow<str> = nbt

8
simdnbt/benches/compare.rs
View file

@ -27,13 +27,13 @@ fn bench_read_file(filename: &str, c: &mut Criterion) {
group.bench_function("simdnbt_borrow_parse", |b| {
b.iter(|| {
black_box(simdnbt::borrow::Nbt::read(&mut input_stream).unwrap());
black_box(simdnbt::borrow::read(&mut input_stream).unwrap());
input_stream.set_position(0);
})
});
group.bench_function("simdnbt_owned_parse", |b| {
b.iter(|| {
black_box(simdnbt::owned::Nbt::read(&mut input_stream).unwrap());
black_box(simdnbt::owned::read(&mut input_stream).unwrap());
input_stream.set_position(0);
})
});
@ -84,7 +84,7 @@ fn bench_read_file(filename: &str, c: &mut Criterion) {
})
});
let nbt = simdnbt::borrow::Nbt::read(&mut Cursor::new(&input))
let nbt = simdnbt::borrow::read(&mut Cursor::new(&input))
.unwrap()
.unwrap();
group.bench_function("simdnbt_borrow_write", |b| {
@ -95,7 +95,7 @@ fn bench_read_file(filename: &str, c: &mut Criterion) {
})
});
let nbt = simdnbt::owned::Nbt::read(&mut Cursor::new(&input))
let nbt = simdnbt::owned::read(&mut Cursor::new(&input))
.unwrap()
.unwrap();
group.bench_function("simdnbt_owned_write", |b| {

4
simdnbt/benches/compare_realworld.rs
View file

@ -33,7 +33,7 @@ pub fn bench_read_file(filename: &str, c: &mut Criterion) {
graphite_items_from_nbt(graphite_binary::nbt::decode::read(&mut &input[..]).unwrap())
.unwrap();
let simdnbt_nbt = simdnbt_items_from_nbt(
simdnbt::borrow::Nbt::read(&mut Cursor::new(input))
simdnbt::borrow::read(&mut Cursor::new(input))
.unwrap()
.unwrap(),
)
@ -63,7 +63,7 @@ pub fn bench_read_file(filename: &str, c: &mut Criterion) {
group.bench_function("simdnbt_parse", |b| {
b.iter(|| {
let input = black_box(input);
let nbt = black_box(simdnbt::borrow::Nbt::read(&mut Cursor::new(input)));
let nbt = black_box(simdnbt::borrow::read(&mut Cursor::new(input)));
let nbt = nbt.unwrap().unwrap();
black_box(simdnbt_items_from_nbt(nbt));
})

6
simdnbt/benches/nbt_borrow.rs
View file

@ -27,14 +27,12 @@ fn bench_file(filename: &str, c: &mut Criterion) {
group.bench_function("Decode", |b| {
b.iter(|| {
black_box(simdnbt::borrow::Nbt::read(&mut input_stream).unwrap());
black_box(simdnbt::borrow::read(&mut input_stream).unwrap());
input_stream.set_position(0);
})
});
let nbt = simdnbt::borrow::Nbt::read(&mut input_stream)
.unwrap()
.unwrap();
let nbt = simdnbt::borrow::read(&mut input_stream).unwrap().unwrap();
group.bench_function("Get", |b| {
b.iter(|| {
let level = nbt.compound("abilities").unwrap();

4
simdnbt/benches/nbt_owned.rs
View file

@ -27,12 +27,12 @@ fn bench_file(filename: &str, c: &mut Criterion) {
group.bench_function("Decode", |b| {
b.iter(|| {
black_box(simdnbt::owned::Nbt::read(&mut decoded_src_stream).unwrap());
black_box(simdnbt::owned::read(&mut decoded_src_stream).unwrap());
decoded_src_stream.set_position(0);
})
});
let nbt = simdnbt::owned::Nbt::read(&mut decoded_src_stream)
let nbt = simdnbt::owned::read(&mut decoded_src_stream)
.unwrap()
.unwrap();
group.bench_function("Get", |b| {

4
simdnbt/examples/hypixel.rs
View file

@ -68,7 +68,7 @@ fn main() {
let input = black_box(include_bytes!("../tests/realworld.nbt"));
for _ in 0..1 {
let nbt = simdnbt::borrow::Nbt::read(&mut Cursor::new(input));
let nbt = simdnbt::borrow::read(&mut Cursor::new(input));
let nbt = black_box(nbt.unwrap().unwrap());
let data = Base::from_nbt(&nbt).unwrap();
@ -76,7 +76,7 @@ fn main() {
// roundtrip
let mut new_nbt_bytes = Vec::new();
data.clone().to_nbt().write(&mut new_nbt_bytes);
let new_nbt = simdnbt::borrow::Nbt::read(&mut Cursor::new(&new_nbt_bytes[..]))
let new_nbt = simdnbt::borrow::read(&mut Cursor::new(&new_nbt_bytes[..]))
.unwrap()
.unwrap();
let new_data = Base::from_nbt(&new_nbt).unwrap();

4
simdnbt/examples/hypixel_no_derive.rs
View file

@ -1,6 +1,6 @@
use std::{collections::HashMap, hint::black_box, io::Cursor};
use simdnbt::borrow::{BaseNbt, Nbt};
use simdnbt::borrow::BaseNbt;
#[derive(Clone, PartialEq, Debug)]
pub struct Item {
@ -103,7 +103,7 @@ fn main() {
let input = black_box(include_bytes!("../tests/realworld.nbt"));
for _ in 0..1 {
let nbt = Nbt::read(&mut Cursor::new(input));
let nbt = simdnbt::borrow::read(&mut Cursor::new(input));
let nbt = black_box(nbt.unwrap().unwrap());
black_box(items_from_nbt(nbt));
}

2
simdnbt/examples/read_and_write.rs
View file

@ -10,7 +10,7 @@ fn main() {
}
let input = input.as_slice();
let nbt = simdnbt::owned::Nbt::read(&mut Cursor::new(input))
let nbt = simdnbt::owned::read(&mut Cursor::new(input))
.unwrap()
.unwrap();

11
simdnbt/src/borrow/compound.rs
View file

@ -19,11 +19,18 @@ pub struct NbtCompound<'a> {
}
impl<'a> NbtCompound<'a> {
pub fn read(data: &mut Cursor<&'a [u8]>, alloc: &TagAllocator<'a>) -> Result<Self, Error> {
/// # Safety
/// The given TagAllocator must be valid for the lifetime of all the tags in this NBT.
pub(crate) unsafe fn read(
data: &mut Cursor<&'a [u8]>,
alloc: &TagAllocator<'a>,
) -> Result<Self, Error> {
Self::read_with_depth(data, alloc, 0, 0)
}
pub fn read_with_depth(
/// # Safety
/// The given TagAllocator must be valid for the lifetime of all the tags in this NBT.
pub(crate) unsafe fn read_with_depth(
data: &mut Cursor<&'a [u8]>,
alloc: &TagAllocator<'a>,
compound_depth: usize,

14
simdnbt/src/borrow/list.rs
View file

@ -35,7 +35,9 @@ pub enum NbtList<'a> {
LongArray(&'a [RawList<'a, i64>]) = LONG_ARRAY_ID,
}
impl<'a> NbtList<'a> {
pub fn read(
/// # Safety
/// The given TagAllocator must be valid for the lifetime of all the tags in this NBT.
pub(crate) unsafe fn read(
data: &mut Cursor<&'a [u8]>,
alloc: &TagAllocator<'a>,
compound_depth: usize,
@ -105,12 +107,10 @@ impl<'a> NbtList<'a> {
let length = read_u32(data)?;
let mut tags = alloc.get().unnamed_compound.start(list_depth);
for _ in 0..length {
let tag = match NbtCompound::read_with_depth(
data,
alloc,
compound_depth + 1,
list_depth,
) {
let tag_res = unsafe {
NbtCompound::read_with_depth(data, alloc, compound_depth + 1, list_depth)
};
let tag = match tag_res {
Ok(tag) => tag,
Err(e) => {
alloc.get().unnamed_compound.finish(tags, list_depth);

132
simdnbt/src/borrow/mod.rs
View file

@ -21,6 +21,51 @@ use crate::{
use self::tag_alloc::TagAllocator;
pub use self::{compound::NbtCompound, list::NbtList};
/// Read a normal root NBT compound. This is either empty or has a name and compound tag.
///
/// Returns `Ok(Nbt::None)` if there is no data.
pub fn read<'a>(data: &mut Cursor<&'a [u8]>) -> Result<Nbt<'a>, Error> {
Nbt::read(data)
}
/// Read a root NBT compound, but without reading the name. This is used in Minecraft when reading
/// NBT over the network.
///
/// This is similar to [`read_tag`], but returns an [`Nbt`] instead (guaranteeing it'll be either
/// empty or a compound).
pub fn read_unnamed<'a>(data: &mut Cursor<&'a [u8]>) -> Result<Nbt<'a>, Error> {
Nbt::read_unnamed(data)
}
/// Read a compound tag. This may have any number of items.
pub fn read_compound<'a>(data: &mut Cursor<&'a [u8]>) -> Result<BaseNbtCompound<'a>, Error> {
let tag_alloc = TagAllocator::new();
let tag = unsafe { NbtCompound::read(data, &tag_alloc) }?;
Ok(BaseNbtCompound {
tag,
_tag_alloc: tag_alloc,
})
}
/// Read an NBT tag, without reading its name. This may be any type of tag except for an end tag. If you need to be able to
/// handle end tags, use [`read_optional_tag`].
pub fn read_tag<'a>(data: &mut Cursor<&'a [u8]>) -> Result<BaseNbtTag<'a>, Error> {
let tag_alloc = TagAllocator::new();
let tag = unsafe { NbtTag::read(data, &tag_alloc) }?;
Ok(BaseNbtTag {
tag,
_tag_alloc: tag_alloc,
})
}
/// Read any NBT tag, without reading its name. This may be any type of tag, including an end tag.
///
/// Returns `Ok(None)` if there is no data.
pub fn read_optional_tag<'a>(data: &mut Cursor<&'a [u8]>) -> Result<Option<BaseNbtTag<'a>>, Error> {
let tag_alloc = TagAllocator::new();
let tag = unsafe { NbtTag::read_optional(data, &tag_alloc) }?;
Ok(tag.map(|tag| BaseNbtTag {
tag,
_tag_alloc: tag_alloc,
}))
}
/// A complete NBT container. This contains a name and a compound tag.
#[derive(Debug)]
pub struct BaseNbt<'a> {
@ -30,6 +75,20 @@ pub struct BaseNbt<'a> {
_tag_alloc: TagAllocator<'a>,
}
/// A nameless NBT container. This only contains a compound tag. This contains a `TagAllocator`,
/// so it can exist independently from a [`BaseNbt`].
pub struct BaseNbtCompound<'a> {
tag: NbtCompound<'a>,
_tag_alloc: TagAllocator<'a>,
}
/// A nameless NBT tag.
pub struct BaseNbtTag<'a> {
tag: NbtTag<'a>,
_tag_alloc: TagAllocator<'a>,
}
/// Either a complete NBT container, or nothing.
#[derive(Debug, PartialEq, Default)]
pub enum Nbt<'a> {
Some(BaseNbt<'a>),
@ -38,8 +97,8 @@ pub enum Nbt<'a> {
}
impl<'a> Nbt<'a> {
/// Reads NBT from the given data. Returns `Ok(None)` if there is no data.
pub fn read(data: &mut Cursor<&'a [u8]>) -> Result<Nbt<'a>, Error> {
/// Reads NBT from the given data. Returns `Ok(Nbt::None)` if there is no data.
fn read(data: &mut Cursor<&'a [u8]>) -> Result<Nbt<'a>, Error> {
let root_type = data.read_u8().map_err(|_| Error::UnexpectedEof)?;
if root_type == END_ID {
return Ok(Nbt::None);
@ -50,7 +109,7 @@ impl<'a> Nbt<'a> {
let tag_alloc = TagAllocator::new();
let name = read_string(data)?;
let tag = NbtCompound::read_with_depth(data, &tag_alloc, 0, 0)?;
let tag = unsafe { NbtCompound::read(data, &tag_alloc) }?;
Ok(Nbt::Some(BaseNbt {
name,
@ -59,6 +118,25 @@ impl<'a> Nbt<'a> {
}))
}
fn read_unnamed(data: &mut Cursor<&'a [u8]>) -> Result<Nbt<'a>, Error> {
let root_type = data.read_u8().map_err(|_| Error::UnexpectedEof)?;
if root_type == END_ID {
return Ok(Nbt::None);
}
if root_type != COMPOUND_ID {
return Err(Error::InvalidRootType(root_type));
}
let tag_alloc = TagAllocator::new();
let tag = unsafe { NbtCompound::read(data, &tag_alloc) }?;
Ok(Nbt::Some(BaseNbt {
name: Mutf8Str::from_slice(&[]),
tag,
_tag_alloc: tag_alloc,
}))
}
pub fn write(&self, data: &mut Vec<u8>) {
match self {
Nbt::Some(nbt) => nbt.write(data),
@ -109,6 +187,20 @@ impl<'a> Deref for BaseNbt<'a> {
&self.tag
}
}
impl<'a> Deref for BaseNbtCompound<'a> {
type Target = NbtCompound<'a>;
fn deref(&self) -> &Self::Target {
&self.tag
}
}
impl<'a> Deref for BaseNbtTag<'a> {
type Target = NbtTag<'a>;
fn deref(&self) -> &Self::Target {
&self.tag
}
}
impl<'a> BaseNbt<'a> {
pub fn write(&self, data: &mut Vec<u8>) {
@ -135,7 +227,7 @@ pub enum NbtTag<'a> {
IntArray(RawList<'a, i32>),
LongArray(RawList<'a, i64>),
}
impl<'a> NbtTag<'a> {
impl<'a, 'b> NbtTag<'a> {
/// Get the numerical ID of the tag type.
#[inline]
pub fn id(&self) -> u8 {
@ -155,8 +247,10 @@ impl<'a> NbtTag<'a> {
}
}
/// # Safety
/// The given TagAllocator must be valid for the lifetime of all the tags in this NBT.
#[inline(always)]
fn read_with_type(
unsafe fn read_with_type(
data: &mut Cursor<&'a [u8]>,
alloc: &TagAllocator<'a>,
tag_type: u8,
@ -202,12 +296,16 @@ impl<'a> NbtTag<'a> {
}
}
pub fn read(data: &mut Cursor<&'a [u8]>, alloc: &TagAllocator<'a>) -> Result<Self, Error> {
/// # Safety
/// The given TagAllocator must be valid for the lifetime of all the tags in this NBT.
unsafe fn read(data: &mut Cursor<&'a [u8]>, alloc: &TagAllocator<'a>) -> Result<Self, Error> {
let tag_type = data.read_u8().map_err(|_| Error::UnexpectedEof)?;
Self::read_with_type(data, alloc, tag_type, 0, 0)
}
pub fn read_optional(
/// # Safety
/// The given TagAllocator must be valid for the lifetime of all the tags in this NBT.
unsafe fn read_optional(
data: &mut Cursor<&'a [u8]>,
alloc: &TagAllocator<'a>,
) -> Result<Option<Self>, Error> {
@ -476,4 +574,24 @@ mod tests {
let res = Nbt::read(&mut Cursor::new(&data));
assert_eq!(res, Err(Error::UnexpectedEof));
}
#[test]
fn read_complexplayer_with_given_alloc() {
let src = include_bytes!("../../tests/complex_player.dat").to_vec();
let mut src_slice = src.as_slice();
let mut decoded_src_decoder = GzDecoder::new(&mut src_slice);
let mut decoded_src = Vec::new();
decoded_src_decoder.read_to_end(&mut decoded_src).unwrap();
let mut decoded_src_as_tag = Vec::new();
decoded_src_as_tag.push(COMPOUND_ID);
decoded_src_as_tag.extend_from_slice(&decoded_src);
decoded_src_as_tag.push(END_ID);
let nbt = super::read_tag(&mut Cursor::new(&decoded_src_as_tag)).unwrap();
let nbt = nbt.compound().unwrap().compound("").unwrap();
assert_eq!(nbt.float("foodExhaustionLevel").unwrap() as u32, 2);
assert_eq!(nbt.list("Rotation").unwrap().floats().unwrap().len(), 2);
}
}

17
simdnbt/src/borrow/tag_alloc.rs
View file

@ -29,6 +29,8 @@ use super::{NbtCompound, NbtList, NbtTag};
// this value appears to have the best results on my pc when testing with complex_player.dat
const MIN_ALLOC_SIZE: usize = 1024;
/// The data structure that contains all the parsed NBT tags. This must stay in scope for as long
/// as the borrowed NBT exists.
#[derive(Default)]
pub struct TagAllocator<'a>(UnsafeCell<TagAllocatorImpl<'a>>);
impl<'a> TagAllocator<'a> {
@ -235,21 +237,6 @@ impl<T> ContiguousTagsAllocator<T> {
self.alloc.len = self.size;
}
#[inline]
pub fn extend_from_slice(&mut self, slice: &[T]) {
while self.alloc.len + slice.len() > self.alloc.cap {
self.grow();
}
// copy the slice
unsafe {
let end = self.alloc.ptr.as_ptr().add(self.alloc.len);
std::ptr::copy_nonoverlapping(slice.as_ptr(), end, slice.len());
}
self.alloc.len += slice.len();
self.size += slice.len();
}
#[inline]
pub fn push(&mut self, value: T) {
// check if we need to reallocate

4
simdnbt/src/mutf8.rs
View file

@ -163,6 +163,10 @@ impl Borrow<Mutf8Str> for Mutf8String {
}
impl Mutf8String {
pub fn new() -> Self {
Self { vec: Vec::new() }
}
#[inline]
pub fn as_str(&self) -> &Mutf8Str {
Mutf8Str::from_slice(self.vec.as_slice())

6
simdnbt/src/owned/compound.rs
View file

@ -28,11 +28,11 @@ impl NbtCompound {
Self { values }
}
pub fn read(data: &mut Cursor<&[u8]>) -> Result<Self, Error> {
pub(crate) fn read(data: &mut Cursor<&[u8]>) -> Result<Self, Error> {
Self::read_with_depth(data, 0)
}
pub fn read_with_depth_and_capacity(
pub(crate) fn read_with_depth_and_capacity(
data: &mut Cursor<&[u8]>,
depth: usize,
capacity: usize,
@ -73,7 +73,7 @@ impl NbtCompound {
Ok(Self { values })
}
pub fn read_with_depth(data: &mut Cursor<&[u8]>, depth: usize) -> Result<Self, Error> {
pub(crate) fn read_with_depth(data: &mut Cursor<&[u8]>, depth: usize) -> Result<Self, Error> {
Self::read_with_depth_and_capacity(data, depth, 8)
}

2
simdnbt/src/owned/list.rs
View file

@ -37,7 +37,7 @@ pub enum NbtList {
LongArray(Vec<Vec<i64>>) = LONG_ARRAY_ID,
}
impl NbtList {
pub fn read(data: &mut Cursor<&[u8]>, depth: usize) -> Result<Self, Error> {
pub(crate) fn read(data: &mut Cursor<&[u8]>, depth: usize) -> Result<Self, Error> {
if depth > MAX_DEPTH {
return Err(Error::MaxDepthExceeded);
}

43
simdnbt/src/owned/mod.rs
View file

@ -1,4 +1,5 @@
//! The owned variant of NBT. This is useful if you're writing data from scratch or if you can't keep a reference to the original data.
//! The owned variant of NBT. This is useful if you're writing NBT or if you can't keep a reference
//! to the original data.
mod compound;
mod list;
@ -20,6 +21,36 @@ use crate::{
pub use self::{compound::NbtCompound, list::NbtList};
/// Read a normal root NBT compound. This is either empty or has a name and compound tag.
///
/// Returns `Ok(Nbt::None)` if there is no data.
pub fn read(data: &mut Cursor<&[u8]>) -> Result<Nbt, Error> {
Nbt::read(data)
}
/// Read a root NBT compound, but without reading the name. This is used in Minecraft when reading
/// NBT over the network.
///
/// This is similar to [`read_tag`], but returns an [`Nbt`] instead (guaranteeing it'll be either
/// empty or a compound).
pub fn read_unnamed(data: &mut Cursor<&[u8]>) -> Result<Nbt, Error> {
Nbt::read_unnamed(data)
}
/// Read a compound tag. This may have any number of items.
pub fn read_compound(data: &mut Cursor<&[u8]>) -> Result<NbtCompound, Error> {
NbtCompound::read(data)
}
/// Read an NBT tag, without reading its name. This may be any type of tag except for an end tag. If you need to be able to
/// handle end tags, use [`read_optional_tag`].
pub fn read_tag(data: &mut Cursor<&[u8]>) -> Result<NbtTag, Error> {
NbtTag::read(data)
}
/// Read any NBT tag, without reading its name. This may be any type of tag, including an end tag.
///
/// Returns `Ok(None)` if there is no data.
pub fn read_optional_tag(data: &mut Cursor<&[u8]>) -> Result<Option<NbtTag>, Error> {
Ok(NbtTag::read_optional(data)?)
}
/// A complete NBT container. This contains a name and a compound tag.
#[derive(Debug, Clone, PartialEq, Default)]
pub struct BaseNbt {
@ -39,8 +70,8 @@ impl Nbt {
Self::Some(BaseNbt { name, tag })
}
/// Reads NBT from the given data. Returns `Ok(None)` if there is no data.
pub fn read(data: &mut Cursor<&[u8]>) -> Result<Nbt, Error> {
/// Reads NBT from the given data. Returns `Ok(Nbt::None)` if there is no data.
fn read(data: &mut Cursor<&[u8]>) -> Result<Nbt, Error> {
let root_type = data.read_u8().map_err(|_| Error::UnexpectedEof)?;
if root_type == END_ID {
return Ok(Nbt::None);
@ -49,12 +80,12 @@ impl Nbt {
return Err(Error::InvalidRootType(root_type));
}
let name = read_string(data)?.to_owned();
let tag = NbtCompound::read_with_depth(data, 0)?;
let tag = NbtCompound::read(data)?;
Ok(Nbt::Some(BaseNbt { name, tag }))
}
pub fn read_unnamed(data: &mut Cursor<&[u8]>) -> Result<Nbt, Error> {
fn read_unnamed(data: &mut Cursor<&[u8]>) -> Result<Nbt, Error> {
let root_type = data.read_u8().map_err(|_| Error::UnexpectedEof)?;
if root_type == END_ID {
return Ok(Nbt::None);
@ -62,7 +93,7 @@ impl Nbt {
if root_type != COMPOUND_ID {
return Err(Error::InvalidRootType(root_type));
}
let tag = NbtCompound::read_with_depth(data, 0)?;
let tag = NbtCompound::read(data)?;
Ok(Nbt::Some(BaseNbt {
name: Mutf8String::from(""),