From 8d8ecbbdc9fcd0b2c640a52e2a71e1a3150d983c Mon Sep 17 00:00:00 2001
From: Moritz Hoffmann <antiguru@gmail.com>
Date: Sat, 25 Nov 2023 13:17:49 -0500
Subject: [PATCH] Avoid reallocations in first `keep`, bound stash size

A follow-up to #425. Retain two elements instead of three in the stash.
Avoid doubling reallocations in the first `keep` object.

Signed-off-by: Moritz Hoffmann <antiguru@gmail.com>
---
 src/trace/implementations/merge_batcher_col.rs | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/trace/implementations/merge_batcher_col.rs b/src/trace/implementations/merge_batcher_col.rs
index 7aaf80340..5d4fad18a 100644
--- a/src/trace/implementations/merge_batcher_col.rs
+++ b/src/trace/implementations/merge_batcher_col.rs
@@ -105,7 +105,11 @@ where
             for datum @ ((_key, _val), time, _diff) in &buffer[..] {
                 if upper.less_equal(time) {
                     self.frontier.insert(time.clone());
-                    if !keep.is_empty() && keep.len() == keep.capacity() {
+                    if keep.is_empty() {
+                        if keep.capacity() != MergeSorterColumnation::<(U::Key, U::Val), U::Time, U::Diff>::buffer_size() {
+                            keep = self.sorter.empty();
+                        }
+                    } else if keep.len() == keep.capacity() {
                         kept.push(keep);
                         keep = self.sorter.empty();
                     }
@@ -228,7 +232,7 @@ impl<D: Ord+Clone+Columnation+'static, T: Ord+Clone+Columnation+'static, R: Semi
 
     /// Insert an empty buffer into the stash. Panics if the buffer is not empty.
     fn recycle(&mut self, mut buffer: TimelyStack<(D, T, R)>) {
-        if buffer.capacity() == Self::buffer_size() && self.stash.len() <= 2 {
+        if buffer.capacity() == Self::buffer_size() && self.stash.len() < 2 {
             buffer.clear();
             self.stash.push(buffer);
         }