From 43b21f2a9b43279c2069b851b52d3386b8412d5b Mon Sep 17 00:00:00 2001
From: Moritz Hoffmann <antiguru@gmail.com>
Date: Fri, 18 Aug 2023 12:37:59 +0200
Subject: [PATCH] Shrink containers if their capacity isn't needed

This eagerly releases memory after retracting large amount of data, without
having to wait for another operation to happen.

Signed-off-by: Moritz Hoffmann <antiguru@gmail.com>
---
 src/trace/implementations/ord.rs |  8 ++++++++
 src/trace/layers/mod.rs          | 10 ++++++++++
 2 files changed, 18 insertions(+)

diff --git a/src/trace/implementations/ord.rs b/src/trace/implementations/ord.rs
index 5c995e048..a8368bd88 100644
--- a/src/trace/implementations/ord.rs
+++ b/src/trace/implementations/ord.rs
@@ -196,6 +196,7 @@ where
             }
         }
         layer.vals.vals.vals.truncate(write_position);
+        layer.vals.vals.vals.maybe_shrink();
         layer.vals.offs[layer.vals.keys.len()] = O::try_from(write_position).unwrap();
 
         // 3. Remove values with empty histories. In addition, we need to update offsets
@@ -221,6 +222,7 @@ where
         });
         debug_assert_eq!(write_position, layer.vals.keys.len());
         layer.vals.offs.truncate(write_position + 1);
+        layer.vals.offs.maybe_shrink();
         layer.offs[layer.keys.len()] = O::try_from(write_position).unwrap();
 
         // 4. Remove empty keys.
@@ -238,6 +240,9 @@ where
         });
         debug_assert_eq!(write_position, layer.keys.len());
         layer.offs.truncate(layer.keys.len()+1);
+        layer.offs.maybe_shrink();
+        layer.vals.keys.maybe_shrink();
+        layer.keys.maybe_shrink();
     }
 }
 
@@ -568,6 +573,7 @@ where
             }
         }
         layer.vals.vals.truncate(write_position);
+        layer.vals.vals.maybe_shrink();
         layer.offs[layer.keys.len()] = O::try_from(write_position).unwrap();
 
         // 4. Remove empty keys.
@@ -585,6 +591,8 @@ where
         });
         debug_assert_eq!(write_position, layer.keys.len());
         layer.offs.truncate(layer.keys.len()+1);
+        layer.offs.maybe_shrink();
+        layer.keys.maybe_shrink();
     }
 }
 
diff --git a/src/trace/layers/mod.rs b/src/trace/layers/mod.rs
index e0d26088a..cbc7d509d 100644
--- a/src/trace/layers/mod.rs
+++ b/src/trace/layers/mod.rs
@@ -125,6 +125,8 @@ pub trait BatchContainer: Default {
     fn reserve(&mut self, additional: usize);
     /// Creates a new container with sufficient capacity.
     fn merge_capacity(cont1: &Self, cont2: &Self) -> Self;
+    /// Shrink the container if needed
+    fn maybe_shrink(&mut self);
 }
 
 impl<T: Clone> BatchContainer for Vec<T> {
@@ -147,6 +149,11 @@ impl<T: Clone> BatchContainer for Vec<T> {
     fn merge_capacity(cont1: &Self, cont2: &Self) -> Self {
         Vec::with_capacity(cont1.len() + cont2.len())
     }
+    fn maybe_shrink(&mut self) {
+        if self.len() < self.capacity() / 4 {
+            self.shrink_to_fit();
+        }
+    }
 }
 
 impl<T: Columnation> BatchContainer for TimelyStack<T> {
@@ -173,6 +180,9 @@ impl<T: Columnation> BatchContainer for TimelyStack<T> {
         new.reserve_regions(std::iter::once(cont1).chain(std::iter::once(cont2)));
         new
     }
+    fn maybe_shrink(&mut self) {
+        // Cannot shrink TimelyStack yet.
+    }
 }