Base update capacity off updates instead of old capacity

This makes sure we never reallocate the update container. Without this, we can allocate a container that is too small to contain all updates that we might generate during merging. Signed-off-by: Moritz Hoffmann <antiguru@gmail.com>
antiguru · Jan 22, 2024 · 17e36c9 · 17e36c9
1 parent dc5ebef
commit 17e36c9
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 19 deletions.
diff --git a/src/trace/implementations/ord_neu.rs b/src/trace/implementations/ord_neu.rs
@@ -201,15 +201,15 @@ mod val_batch {
 
             let description = Description::new(batch1.lower().clone(), batch2.upper().clone(), since);
 
-            let batch1 = &batch1.storage;
-            let batch2 = &batch2.storage;
+            let storage1 = &batch1.storage;
+            let storage2 = &batch2.storage;
 
             let mut storage = OrdValStorage {
-                keys: L::KeyContainer::merge_capacity(&batch1.keys, &batch2.keys),
-                keys_offs: L::OffsetContainer::with_capacity(batch1.keys_offs.len() + batch2.keys_offs.len()),
-                vals: L::ValContainer::merge_capacity(&batch1.vals, &batch2.vals),
-                vals_offs: L::OffsetContainer::with_capacity(batch1.vals_offs.len() + batch2.vals_offs.len()),
-                updates: L::UpdContainer::merge_capacity(&batch1.updates, &batch2.updates),
+                keys: L::KeyContainer::merge_capacity(&storage1.keys, &storage2.keys),
+                keys_offs: L::OffsetContainer::with_capacity(storage1.keys_offs.len() + storage2.keys_offs.len()),
+                vals: L::ValContainer::merge_capacity(&storage1.vals, &storage2.vals),
+                vals_offs: L::OffsetContainer::with_capacity(storage1.vals_offs.len() + storage2.vals_offs.len()),
+                updates: L::UpdContainer::with_capacity(batch1.updates + batch2.updates),
             };
 
             // Mark explicit types because type inference fails to resolve it.
@@ -761,13 +761,13 @@ mod key_batch {
 
             let description = Description::new(batch1.lower().clone(), batch2.upper().clone(), since);
 
-            let batch1 = &batch1.storage;
-            let batch2 = &batch2.storage;
+            let storage1 = &batch1.storage;
+            let storage2 = &batch2.storage;
 
             let mut storage = OrdKeyStorage {
-                keys: L::KeyContainer::merge_capacity(&batch1.keys, &batch2.keys),
-                keys_offs: L::OffsetContainer::with_capacity(batch1.keys_offs.len() + batch2.keys_offs.len()),
-                updates: L::UpdContainer::merge_capacity(&batch1.updates, &batch2.updates),
+                keys: L::KeyContainer::merge_capacity(&storage1.keys, &storage2.keys),
+                keys_offs: L::OffsetContainer::with_capacity(storage1.keys_offs.len() + storage2.keys_offs.len()),
+                updates: L::UpdContainer::with_capacity(batch1.updates + batch2.updates),
             };
 
             let keys_offs: &mut L::OffsetContainer = &mut storage.keys_offs;

diff --git a/src/trace/implementations/rhh.rs b/src/trace/implementations/rhh.rs
@@ -334,15 +334,15 @@ mod val_batch {
             let max_cap = batch1.len() + batch2.len();
             let rhh_cap = 2 * max_cap;
 
-            let batch1 = &batch1.storage;
-            let batch2 = &batch2.storage;
+            let storage1 = &batch1.storage;
+            let storage2 = &batch2.storage;
 
             let mut storage = RhhValStorage {
-                keys: L::KeyContainer::merge_capacity(&batch1.keys, &batch2.keys),
-                keys_offs: L::OffsetContainer::with_capacity(batch1.keys_offs.len() + batch2.keys_offs.len()),
-                vals: L::ValContainer::merge_capacity(&batch1.vals, &batch2.vals),
-                vals_offs: L::OffsetContainer::with_capacity(batch1.vals_offs.len() + batch2.vals_offs.len()),
-                updates: L::UpdContainer::merge_capacity(&batch1.updates, &batch2.updates),
+                keys: L::KeyContainer::merge_capacity(&storage1.keys, &storage2.keys),
+                keys_offs: L::OffsetContainer::with_capacity(storage1.keys_offs.len() + storage2.keys_offs.len()),
+                vals: L::ValContainer::merge_capacity(&storage1.vals, &storage2.vals),
+                vals_offs: L::OffsetContainer::with_capacity(storage1.vals_offs.len() + storage2.vals_offs.len()),
+                updates: L::UpdContainer::with_capacity(batch1.updates + batch2.updates),
                 key_count: 0,
                 key_capacity: rhh_cap,
                 divisor: RhhValStorage::<L>::divisor_for_capacity(rhh_cap),