From f47749f7a8de596de428db9ea0820d1412eaef4f Mon Sep 17 00:00:00 2001 From: mat Date: Tue, 14 May 2024 01:54:39 +0000 Subject: [PATCH] keep track of depths for compounds and lists separately for 10% speedup --- simdnbt/src/borrow/compound.rs | 30 ++++++++++---------- simdnbt/src/borrow/list.rs | 50 +++++++++++++++++++--------------- simdnbt/src/borrow/mod.rs | 19 +++++++++---- 3 files changed, 57 insertions(+), 42 deletions(-) diff --git a/simdnbt/src/borrow/compound.rs b/simdnbt/src/borrow/compound.rs index 11bbc69..ed93b70 100644 --- a/simdnbt/src/borrow/compound.rs +++ b/simdnbt/src/borrow/compound.rs @@ -20,19 +20,20 @@ pub struct NbtCompound<'a> { impl<'a> NbtCompound<'a> { pub fn read(data: &mut Cursor<&'a [u8]>, alloc: &TagAllocator<'a>) -> Result { - Self::read_with_depth(data, alloc, 0) + Self::read_with_depth(data, alloc, 0, 0) } pub fn read_with_depth( data: &mut Cursor<&'a [u8]>, alloc: &TagAllocator<'a>, - depth: usize, + compound_depth: usize, + list_depth: usize, ) -> Result { - if depth > MAX_DEPTH { + if compound_depth + list_depth > MAX_DEPTH { return Err(Error::MaxDepthExceeded); } - let mut tags = alloc.get().named.start(depth); + let mut tags = alloc.get().named.start(compound_depth); let mut tags_buffer = unsafe { MaybeUninit::<[MaybeUninit<(&Mutf8Str, NbtTag<'a>)>; 4]>::uninit().assume_init() @@ -43,7 +44,7 @@ impl<'a> NbtCompound<'a> { let tag_type = match data.read_u8() { Ok(tag_type) => tag_type, Err(_) => { - alloc.get().named.finish(tags, depth); + alloc.get().named.finish(tags, compound_depth); return Err(Error::UnexpectedEof); } }; @@ -54,19 +55,20 @@ impl<'a> NbtCompound<'a> { let tag_name = match read_string(data) { Ok(name) => name, Err(_) => { - alloc.get().named.finish(tags, depth); + alloc.get().named.finish(tags, compound_depth); // the only error read_string can return is UnexpectedEof, so this makes it // slightly faster return Err(Error::UnexpectedEof); } }; - let tag = match NbtTag::read_with_type(data, alloc, tag_type, depth) { - Ok(tag) => tag, - Err(e) => { - alloc.get().named.finish(tags, depth); - return Err(e); - } - }; + let tag = + match NbtTag::read_with_type(data, alloc, tag_type, compound_depth, list_depth) { + Ok(tag) => tag, + Err(e) => { + alloc.get().named.finish(tags, compound_depth); + return Err(e); + } + }; tags_buffer[tags_buffer_len] = MaybeUninit::new((tag_name, tag)); tags_buffer_len += 1; @@ -84,7 +86,7 @@ impl<'a> NbtCompound<'a> { tags.push(unsafe { tags_buffer.get_unchecked(i).assume_init_read() }); } - let values = alloc.get().named.finish(tags, depth); + let values = alloc.get().named.finish(tags, compound_depth); Ok(Self { values }) } diff --git a/simdnbt/src/borrow/list.rs b/simdnbt/src/borrow/list.rs index b06a5c9..9279b46 100644 --- a/simdnbt/src/borrow/list.rs +++ b/simdnbt/src/borrow/list.rs @@ -38,9 +38,10 @@ impl<'a> NbtList<'a> { pub fn read( data: &mut Cursor<&'a [u8]>, alloc: &TagAllocator<'a>, - depth: usize, + compound_depth: usize, + list_depth: usize, ) -> Result { - if depth > MAX_DEPTH { + if compound_depth + list_depth > MAX_DEPTH { return Err(Error::MaxDepthExceeded); } let tag_type = data.read_u8().map_err(|_| Error::UnexpectedEof)?; @@ -57,93 +58,98 @@ impl<'a> NbtList<'a> { DOUBLE_ID => NbtList::Double(RawList::new(read_with_u32_length(data, 8)?)), BYTE_ARRAY_ID => NbtList::ByteArray({ let length = read_u32(data)?; - let mut tags = alloc.get().unnamed_bytearray.start(depth); + let mut tags = alloc.get().unnamed_bytearray.start(list_depth); for _ in 0..length { let tag = match read_u8_array(data) { Ok(tag) => tag, Err(e) => { - alloc.get().unnamed_bytearray.finish(tags, depth); + alloc.get().unnamed_bytearray.finish(tags, list_depth); return Err(e); } }; tags.push(tag); } - alloc.get().unnamed_bytearray.finish(tags, depth) + alloc.get().unnamed_bytearray.finish(tags, list_depth) }), STRING_ID => NbtList::String({ let length = read_u32(data)?; - let mut tags = alloc.get().unnamed_string.start(depth); + let mut tags = alloc.get().unnamed_string.start(list_depth); for _ in 0..length { let tag = match read_string(data) { Ok(tag) => tag, Err(e) => { - alloc.get().unnamed_string.finish(tags, depth); + alloc.get().unnamed_string.finish(tags, list_depth); return Err(e); } }; tags.push(tag); } - alloc.get().unnamed_string.finish(tags, depth) + alloc.get().unnamed_string.finish(tags, list_depth) }), LIST_ID => NbtList::List({ let length = read_u32(data)?; - let mut tags = alloc.get().unnamed_list.start(depth); + let mut tags = alloc.get().unnamed_list.start(list_depth); for _ in 0..length { - let tag = match NbtList::read(data, alloc, depth + 1) { + let tag = match NbtList::read(data, alloc, compound_depth, list_depth + 1) { Ok(tag) => tag, Err(e) => { - alloc.get().unnamed_list.finish(tags, depth); + alloc.get().unnamed_list.finish(tags, list_depth); return Err(e); } }; tags.push(tag) } - alloc.get().unnamed_list.finish(tags, depth) + alloc.get().unnamed_list.finish(tags, list_depth) }), COMPOUND_ID => NbtList::Compound({ let length = read_u32(data)?; - let mut tags = alloc.get().unnamed_compound.start(depth); + let mut tags = alloc.get().unnamed_compound.start(list_depth); for _ in 0..length { - let tag = match NbtCompound::read_with_depth(data, alloc, depth + 1) { + let tag = match NbtCompound::read_with_depth( + data, + alloc, + compound_depth + 1, + list_depth, + ) { Ok(tag) => tag, Err(e) => { - alloc.get().unnamed_compound.finish(tags, depth); + alloc.get().unnamed_compound.finish(tags, list_depth); return Err(e); } }; tags.push(tag); } - alloc.get().unnamed_compound.finish(tags, depth) + alloc.get().unnamed_compound.finish(tags, list_depth) }), INT_ARRAY_ID => NbtList::IntArray({ let length = read_u32(data)?; - let mut tags = alloc.get().unnamed_intarray.start(depth); + let mut tags = alloc.get().unnamed_intarray.start(list_depth); for _ in 0..length { let tag = match read_int_array(data) { Ok(tag) => tag, Err(e) => { - alloc.get().unnamed_intarray.finish(tags, depth); + alloc.get().unnamed_intarray.finish(tags, list_depth); return Err(e); } }; tags.push(tag); } - alloc.get().unnamed_intarray.finish(tags, depth) + alloc.get().unnamed_intarray.finish(tags, list_depth) }), LONG_ARRAY_ID => NbtList::LongArray({ let length = read_u32(data)?; - let mut tags = alloc.get().unnamed_longarray.start(depth); + let mut tags = alloc.get().unnamed_longarray.start(list_depth); for _ in 0..length { let tag = match read_long_array(data) { Ok(tag) => tag, Err(e) => { - alloc.get().unnamed_longarray.finish(tags, depth); + alloc.get().unnamed_longarray.finish(tags, list_depth); return Err(e); } }; tags.push(tag); } - alloc.get().unnamed_longarray.finish(tags, depth) + alloc.get().unnamed_longarray.finish(tags, list_depth) }), _ => return Err(Error::UnknownTagId(tag_type)), }) diff --git a/simdnbt/src/borrow/mod.rs b/simdnbt/src/borrow/mod.rs index 67e2755..f8c0a84 100644 --- a/simdnbt/src/borrow/mod.rs +++ b/simdnbt/src/borrow/mod.rs @@ -50,7 +50,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)?; + let tag = NbtCompound::read_with_depth(data, &tag_alloc, 0, 0)?; Ok(Nbt::Some(BaseNbt { name, @@ -160,7 +160,8 @@ impl<'a> NbtTag<'a> { data: &mut Cursor<&'a [u8]>, alloc: &TagAllocator<'a>, tag_type: u8, - depth: usize, + compound_depth: usize, + list_depth: usize, ) -> Result { match tag_type { BYTE_ID => Ok(NbtTag::Byte( @@ -183,11 +184,17 @@ impl<'a> NbtTag<'a> { )), BYTE_ARRAY_ID => Ok(NbtTag::ByteArray(read_with_u32_length(data, 1)?)), STRING_ID => Ok(NbtTag::String(read_string(data)?)), - LIST_ID => Ok(NbtTag::List(NbtList::read(data, alloc, depth + 1)?)), + LIST_ID => Ok(NbtTag::List(NbtList::read( + data, + alloc, + compound_depth, + list_depth + 1, + )?)), COMPOUND_ID => Ok(NbtTag::Compound(NbtCompound::read_with_depth( data, alloc, - depth + 1, + compound_depth + 1, + list_depth, )?)), INT_ARRAY_ID => Ok(NbtTag::IntArray(read_int_array(data)?)), LONG_ARRAY_ID => Ok(NbtTag::LongArray(read_long_array(data)?)), @@ -197,7 +204,7 @@ impl<'a> NbtTag<'a> { pub fn read(data: &mut Cursor<&'a [u8]>, alloc: &TagAllocator<'a>) -> Result { let tag_type = data.read_u8().map_err(|_| Error::UnexpectedEof)?; - Self::read_with_type(data, alloc, tag_type, 0) + Self::read_with_type(data, alloc, tag_type, 0, 0) } pub fn read_optional( @@ -208,7 +215,7 @@ impl<'a> NbtTag<'a> { if tag_type == END_ID { return Ok(None); } - Ok(Some(Self::read_with_type(data, alloc, tag_type, 0)?)) + Ok(Some(Self::read_with_type(data, alloc, tag_type, 0, 0)?)) } pub fn byte(&self) -> Option {