From b0147d96ea1903ffb5136adecb55b8736167866d Mon Sep 17 00:00:00 2001
From: John Spray <john@neon.tech>
Date: Wed, 17 Apr 2024 15:15:47 +0100
Subject: [PATCH] utils: use SmallVec in VecMap

---
 Cargo.lock                |  1 +
 libs/utils/Cargo.toml     |  1 +
 libs/utils/src/vec_map.rs | 21 ++++++++++++++-------
 3 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 2505d4d3ed..3bae9c1516 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -6850,6 +6850,7 @@ dependencies = [
  "serde_path_to_error",
  "serde_with",
  "signal-hook",
+ "smallvec",
  "strum",
  "strum_macros",
  "thiserror",
diff --git a/libs/utils/Cargo.toml b/libs/utils/Cargo.toml
index 261ca2cc1a..412e9fe829 100644
--- a/libs/utils/Cargo.toml
+++ b/libs/utils/Cargo.toml
@@ -36,6 +36,7 @@ routerify.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 signal-hook.workspace = true
+smallvec.workspace = true
 thiserror.workspace = true
 tokio.workspace = true
 tokio-tar.workspace = true
diff --git a/libs/utils/src/vec_map.rs b/libs/utils/src/vec_map.rs
index 8b9f8e167f..35595fab77 100644
--- a/libs/utils/src/vec_map.rs
+++ b/libs/utils/src/vec_map.rs
@@ -1,11 +1,15 @@
 use std::{alloc::Layout, cmp::Ordering, ops::RangeBounds};
 
+use smallvec::SmallVec;
+
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum VecMapOrdering {
     Greater,
     GreaterOrEqual,
 }
 
+const INLINE_ELEMENTS: usize = 1;
+
 /// Ordered map datastructure implemented in a Vec.
 /// Append only - can only add keys that are larger than the
 /// current max key.
@@ -13,7 +17,7 @@ pub enum VecMapOrdering {
 /// during `VecMap` construction.
 #[derive(Clone, Debug)]
 pub struct VecMap<K, V> {
-    data: Vec<(K, V)>,
+    data: SmallVec<[(K, V); INLINE_ELEMENTS]>,
     ordering: VecMapOrdering,
 }
 
@@ -37,7 +41,7 @@ pub enum VecMapError {
 impl<K: Ord, V> VecMap<K, V> {
     pub fn new(ordering: VecMapOrdering) -> Self {
         Self {
-            data: Vec::new(),
+            data: Default::default(),
             ordering,
         }
     }
@@ -48,7 +52,7 @@ impl<K: Ord, V> VecMap<K, V> {
 
     pub fn with_capacity(capacity: usize, ordering: VecMapOrdering) -> Self {
         Self {
-            data: Vec::with_capacity(capacity),
+            data: SmallVec::with_capacity(capacity),
             ordering,
         }
     }
@@ -144,11 +148,11 @@ impl<K: Ord, V> VecMap<K, V> {
 
         (
             VecMap {
-                data: self.data[..split_idx].to_vec(),
+                data: SmallVec::from(&self.data[..split_idx]),
                 ordering: self.ordering,
             },
             VecMap {
-                data: self.data[split_idx..].to_vec(),
+                data: SmallVec::from(&self.data[split_idx..]),
                 ordering: self.ordering,
             },
         )
@@ -195,7 +199,10 @@ impl<K: Ord, V> VecMap<K, V> {
     /// Instrument an operation on the underlying [`Vec`].
     /// Will panic if the operation decreases capacity.
     /// Returns the increase in memory usage caused by the op.
-    fn instrument_vec_op(&mut self, op: impl FnOnce(&mut Vec<(K, V)>)) -> usize {
+    fn instrument_vec_op(
+        &mut self,
+        op: impl FnOnce(&mut SmallVec<[(K, V); INLINE_ELEMENTS]>),
+    ) -> usize {
         let old_cap = self.data.capacity();
         op(&mut self.data);
         let new_cap = self.data.capacity();
@@ -235,7 +242,7 @@ impl<K: Ord, V> VecMap<K, V> {
 
 impl<K: Ord, V> IntoIterator for VecMap<K, V> {
     type Item = (K, V);
-    type IntoIter = std::vec::IntoIter<(K, V)>;
+    type IntoIter = smallvec::IntoIter<[(K, V); INLINE_ELEMENTS]>;
 
     fn into_iter(self) -> Self::IntoIter {
         self.data.into_iter()