From e1a1ce52bc7d40a8c07b5a2635f26afae8ef49d3 Mon Sep 17 00:00:00 2001
From: marvin-j97 <marvin.janke.97@gmail.com>
Date: Fri, 20 Dec 2024 14:49:39 +0100
Subject: [PATCH] perf: skip heap alloc in Memtable::get

---
 Cargo.toml          |   2 +-
 src/key.rs          | 100 ++++++++++++++++++++++----------------------
 src/memtable/mod.rs |   4 +-
 3 files changed, 52 insertions(+), 54 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 57cda897..b29604aa 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -26,7 +26,7 @@ bytes = ["value-log/bytes"]
 
 [dependencies]
 byteorder = "1.5.0"
-crossbeam-skiplist = "0.1.3"
+crossbeam-skiplist = { git = "https://github.com/crossbeam-rs/crossbeam", rev = "45425b032b75d40c8f79be2133eb7d33aaa1d4e4", package = "crossbeam-skiplist" }
 double-ended-peekable = "0.1.0"
 enum_dispatch = "0.3.13"
 guardian = "1.1.0"
diff --git a/src/key.rs b/src/key.rs
index 45ed596f..2027cb62 100644
--- a/src/key.rs
+++ b/src/key.rs
@@ -7,6 +7,7 @@ use crate::{
     SeqNo, UserKey, ValueType,
 };
 use byteorder::{ReadBytesExt, WriteBytesExt};
+use crossbeam_skiplist::equivalent::{Comparable, Equivalent};
 use std::{
     cmp::Reverse,
     io::{Read, Write},
@@ -105,54 +106,51 @@ impl Ord for InternalKey {
     }
 }
 
-// TODO: wait for new crossbeam-skiplist
-// TODO: https://github.com/crossbeam-rs/crossbeam/pull/1162
-//
-// impl Equivalent<InternalKeyRef<'_>> for InternalKey {
-//     fn equivalent(&self, other: &InternalKeyRef<'_>) -> bool {
-//         self.user_key == other.user_key && self.seqno == other.seqno
-//     }
-// }
-
-// impl Comparable<InternalKeyRef<'_>> for InternalKey {
-//     fn compare(&self, other: &InternalKeyRef<'_>) -> std::cmp::Ordering {
-//         (&*self.user_key, Reverse(self.seqno)).cmp(&(other.user_key, Reverse(other.seqno)))
-//     }
-// }
-
-// Temporary internal key without heap allocation
-// #[derive(Debug, Eq)]
-// pub struct InternalKeyRef<'a> {
-//     pub user_key: &'a [u8],
-//     pub seqno: SeqNo,
-//     pub value_type: ValueType,
-// }
-
-// impl<'a> InternalKeyRef<'a> {
-//     // Constructor for InternalKeyRef
-//     pub fn new(user_key: &'a [u8], seqno: u64, value_type: ValueType) -> Self {
-//         InternalKeyRef {
-//             user_key,
-//             seqno,
-//             value_type,
-//         }
-//     }
-// }
-
-// impl<'a> PartialEq for InternalKeyRef<'a> {
-//     fn eq(&self, other: &Self) -> bool {
-//         self.user_key == other.user_key && self.seqno == other.seqno
-//     }
-// }
-
-// impl<'a> PartialOrd for InternalKeyRef<'a> {
-//     fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
-//         Some(self.cmp(other))
-//     }
-// }
-
-// impl<'a> Ord for InternalKeyRef<'a> {
-//     fn cmp(&self, other: &Self) -> std::cmp::Ordering {
-//         (&self.user_key, Reverse(self.seqno)).cmp(&(&other.user_key, Reverse(other.seqno)))
-//     }
-// }
+impl Equivalent<InternalKeyRef<'_>> for InternalKey {
+    fn equivalent(&self, other: &InternalKeyRef<'_>) -> bool {
+        self.user_key == other.user_key && self.seqno == other.seqno
+    }
+}
+
+impl Comparable<InternalKeyRef<'_>> for InternalKey {
+    fn compare(&self, other: &InternalKeyRef<'_>) -> std::cmp::Ordering {
+        (&*self.user_key, Reverse(self.seqno)).cmp(&(other.user_key, Reverse(other.seqno)))
+    }
+}
+
+/// Temporary internal key without heap allocation
+#[derive(Debug, Eq)]
+pub struct InternalKeyRef<'a> {
+    pub user_key: &'a [u8],
+    pub seqno: SeqNo,
+    pub value_type: ValueType,
+}
+
+impl<'a> InternalKeyRef<'a> {
+    // Constructor for InternalKeyRef
+    pub fn new(user_key: &'a [u8], seqno: u64, value_type: ValueType) -> Self {
+        InternalKeyRef {
+            user_key,
+            seqno,
+            value_type,
+        }
+    }
+}
+
+impl<'a> PartialEq for InternalKeyRef<'a> {
+    fn eq(&self, other: &Self) -> bool {
+        self.user_key == other.user_key && self.seqno == other.seqno
+    }
+}
+
+impl<'a> PartialOrd for InternalKeyRef<'a> {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl<'a> Ord for InternalKeyRef<'a> {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        (&self.user_key, Reverse(self.seqno)).cmp(&(&other.user_key, Reverse(other.seqno)))
+    }
+}
diff --git a/src/memtable/mod.rs b/src/memtable/mod.rs
index a5b390f1..b96ea0a9 100644
--- a/src/memtable/mod.rs
+++ b/src/memtable/mod.rs
@@ -2,7 +2,7 @@
 // This source code is licensed under both the Apache 2.0 and MIT License
 // (found in the LICENSE-* files in the repository)
 
-use crate::key::InternalKey;
+use crate::key::{InternalKey, InternalKeyRef};
 use crate::segment::block::ItemSize;
 use crate::value::{InternalValue, SeqNo, UserValue, ValueType};
 use crossbeam_skiplist::SkipMap;
@@ -84,7 +84,7 @@ impl Memtable {
         // abcdef -> 6
         // abcdef -> 5
         //
-        let lower_bound = InternalKey::new(
+        let lower_bound = InternalKeyRef::new(
             key,
             match seqno {
                 Some(seqno) => seqno - 1,