From fc4298751fc71506e52c536523c98cac54da7d11 Mon Sep 17 00:00:00 2001
From: Moritz Hoffmann <mh@materialize.com>
Date: Mon, 1 Jul 2024 10:53:16 -0400
Subject: [PATCH] Fix allocation behavior

Signed-off-by: Moritz Hoffmann <mh@materialize.com>
---
 src/ore/src/flatcontainer.rs |  2 ++
 src/ore/src/region.rs        | 14 +++++++++-----
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/src/ore/src/flatcontainer.rs b/src/ore/src/flatcontainer.rs
index f73d59851cc29..2a478b5d5fda2 100644
--- a/src/ore/src/flatcontainer.rs
+++ b/src/ore/src/flatcontainer.rs
@@ -483,6 +483,8 @@ mod lgalloc {
         #[inline]
         fn clear(&mut self) {
             self.slices.clear();
+            self.offsets.clear();
+            self.offsets.push(0);
         }
 
         #[inline]
diff --git a/src/ore/src/region.rs b/src/ore/src/region.rs
index 58e72f66ff12a..11aa51d526df7 100644
--- a/src/ore/src/region.rs
+++ b/src/ore/src/region.rs
@@ -554,14 +554,18 @@ mod vec {
             }
         }
 
+        const MIN_NON_ZERO_CAP: usize = if std::mem::size_of::<T>() == 1 {
+            8
+        } else if std::mem::size_of::<T>() <= 1024 {
+            4
+        } else {
+            1
+        };
+
         /// Grow the array to at least `new_len` elements. Reallocates the underlying storage.
         fn grow(&mut self, new_len: usize) {
             let new_capacity = std::cmp::max(self.capacity() * 2, new_len);
-            println!(
-                "Reallocating {} -> {}, requested {new_len}",
-                self.capacity(),
-                new_capacity
-            );
+            let new_capacity = std::cmp::max(new_capacity, Self::MIN_NON_ZERO_CAP);
             let mut new_vec = LgAllocVec::with_capacity(new_capacity);
 
             let src_ptr = self.elements.as_ptr();