From bbadfb87c4762613bbbfe022147c02489ffd33d5 Mon Sep 17 00:00:00 2001
From: Will Greenberg <ifnspifn@gmail.com>
Date: Mon, 20 Jan 2025 12:13:52 -0800
Subject: [PATCH] Unity optimization + bugfix (#746)

* unity: faster handling of byte arrays

* unity: fix error when reloading a scene
---
 rust/src/unity/types/binary.rs   | 41 +++++++++++++++++++++++++++-----
 src/Common/Unity/AssetManager.ts |  6 +++--
 2 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/rust/src/unity/types/binary.rs b/rust/src/unity/types/binary.rs
index f88d80479..096c0db9b 100644
--- a/rust/src/unity/types/binary.rs
+++ b/rust/src/unity/types/binary.rs
@@ -135,7 +135,7 @@ pub struct Mesh {
     pub keep_vertices: u8,
     pub keep_indices: u8,
     pub index_format: IndexFormat,
-    pub index_buffer: UnityArray<u8>,
+    pub index_buffer: ByteArray,
     #[deku(count = "(4 - deku::byte_offset % 4) % 4")] _alignment2: Vec<u8>,
     #[deku(ctx = "version")]
     pub vertex_data: VertexData,
@@ -143,9 +143,9 @@ pub struct Mesh {
     pub compressed_mesh: CompressedMesh,
     pub local_aabb: AABB,
     pub mesh_usage_flags: i32,
-    pub baked_convex_collision_mesh: UnityArray<u8>,
+    pub baked_convex_collision_mesh: ByteArray,
     #[deku(count = "(4 - deku::byte_offset % 4) % 4")] _alignment4: Vec<u8>,
-    pub baked_triangle_collision_mesh: UnityArray<u8>,
+    pub baked_triangle_collision_mesh: ByteArray,
     #[deku(count = "(4 - deku::byte_offset % 4) % 4")] _alignment5: Vec<u8>,
     pub mesh_metrics: [f32; 2],
     #[deku(ctx = "version")]
@@ -211,10 +211,39 @@ pub struct SubMesh {
 pub struct VertexData {
     pub vertex_count: u32,
     pub channels: UnityArray<ChannelInfo>,
-    pub data: UnityArray<u8>,
+    pub data: ByteArray,
     #[deku(count = "(4 - deku::byte_offset % 4) % 4")] _alignment: Vec<u8>,
 }
 
+#[derive(Default, Debug, Clone)]
+pub struct ByteArray {
+    pub data: Vec<u8>,
+}
+
+impl<'a> DekuRead<'a> for ByteArray {
+    fn read(
+        input: &'a deku::bitvec::BitSlice<u8, deku::bitvec::Msb0>,
+        ctx: (),
+    ) -> Result<(&'a deku::bitvec::BitSlice<u8, deku::bitvec::Msb0>, Self), DekuError>
+    where Self: Sized {
+        let (rest, count) = i32::read(input, ctx)?;
+        let (data_bits, rest) = rest.split_at(count as usize * 8);
+        let bytes = data_bits.domain().region().unwrap().1;
+        Ok((
+            rest,
+            Self {
+                data: bytes.to_vec(),
+            },
+        ))
+    }
+}
+
+impl From<ByteArray> for Vec<u8> {
+    fn from(value: ByteArray) -> Self {
+        value.data
+    }
+}
+
 #[derive(DekuRead, Clone, Debug)]
 pub struct CompressedMesh {
     pub vertices: Packedf32Vec,
@@ -323,9 +352,9 @@ pub struct Texture2D {
     pub lightmap_format: i32,
     pub color_space: ColorSpace,
     #[deku(cond = "version >= UnityVersion::V2020_3_16f1")]
-    pub platform_blob: UnityArray<u8>,
+    pub platform_blob: ByteArray,
     #[deku(count = "(4 - deku::byte_offset % 4) % 4")] _alignment2: Vec<u8>,
-    pub data: UnityArray<u8>,
+    pub data: ByteArray,
     #[deku(count = "(4 - deku::byte_offset % 4) % 4")] _alignment3: Vec<u8>,
     #[deku(ctx = "version")]
     pub streaming_info: StreamingInfo,
diff --git a/src/Common/Unity/AssetManager.ts b/src/Common/Unity/AssetManager.ts
index 28bfb02e2..ba294c000 100644
--- a/src/Common/Unity/AssetManager.ts
+++ b/src/Common/Unity/AssetManager.ts
@@ -156,8 +156,10 @@ export class AssetFile {
         this.waitForHeaderPromise = null;
     }
 
-    public waitForHeader(): Promise<void> {
-        return assertExists(this.waitForHeaderPromise);
+    public async waitForHeader() {
+        if (this.waitForHeaderPromise !== null) {
+            await this.waitForHeaderPromise;
+        }
     }
 
     private async initFullInternal(dataFetcher: DataFetcher): Promise<void> {