From 70e70fd5947e120dc31d0cbdbeb4d9ec170209e3 Mon Sep 17 00:00:00 2001
From: Kenneth Loeffler <kenloef@gmail.com>
Date: Wed, 30 Oct 2024 20:21:53 +0000
Subject: [PATCH] Reserve space for additional instances in binary deserializer
 (#465)

---
 rbx_binary/src/deserializer/state.rs | 4 +++-
 rbx_dom_weak/CHANGELOG.md            | 2 ++
 rbx_dom_weak/src/dom.rs              | 6 ++++++
 3 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/rbx_binary/src/deserializer/state.rs b/rbx_binary/src/deserializer/state.rs
index 4ca5f2ffd..172287f34 100644
--- a/rbx_binary/src/deserializer/state.rs
+++ b/rbx_binary/src/deserializer/state.rs
@@ -216,13 +216,15 @@ impl<'db, R: Read> DeserializerState<'db, R> {
         deserializer: &'db Deserializer<'db>,
         mut input: R,
     ) -> Result<Self, InnerError> {
-        let tree = WeakDom::new(InstanceBuilder::new("DataModel"));
+        let mut tree = WeakDom::new(InstanceBuilder::new("DataModel"));
 
         let header = FileHeader::decode(&mut input)?;
 
         let type_infos = HashMap::with_capacity(header.num_types as usize);
         let instances_by_ref = HashMap::with_capacity(1 + header.num_instances as usize);
 
+        tree.reserve(header.num_instances as usize);
+
         Ok(DeserializerState {
             deserializer,
             input,
diff --git a/rbx_dom_weak/CHANGELOG.md b/rbx_dom_weak/CHANGELOG.md
index f089c48af..f1a05354a 100644
--- a/rbx_dom_weak/CHANGELOG.md
+++ b/rbx_dom_weak/CHANGELOG.md
@@ -94,7 +94,9 @@ where
 * Added `UstrMapExt`, a helper trait providing convenience methods `UstrMap::new` and `UstrMap::with_capacity`.
 * Added re-exports for `ustr` (a convenience function for creating `Ustr`s), `Ustr`, `UstrMap`, and `UstrSet`.
 * Added `InstanceBuilder::with_property_capacity`, which can preallocate an `InstanceBuilder`'s property table. [#464]
+* Added `WeakDom::reserve`, which can preallocate additional space for instances in the `WeakDom`. [#465]
 
+[#465]: https://github.com/rojo-rbx/rbx-dom/pull/465
 [#464]: https://github.com/rojo-rbx/rbx-dom/pull/464
 
 ## 2.9.0 (2024-08-22)
diff --git a/rbx_dom_weak/src/dom.rs b/rbx_dom_weak/src/dom.rs
index 81c942bb3..a9e4fa423 100644
--- a/rbx_dom_weak/src/dom.rs
+++ b/rbx_dom_weak/src/dom.rs
@@ -32,6 +32,12 @@ impl WeakDom {
         dom
     }
 
+    /// Reserve at least enough space for `additional` number of instances in
+    /// the WeakDom.
+    pub fn reserve(&mut self, additional: usize) {
+        self.instances.reserve(additional);
+    }
+
     /// Consumes the WeakDom, returning its underlying root ref and backing
     /// storage. This method is useful when tree-preserving operations are too
     /// slow.