use smallvec::SmallVec;
pub const TXN_KEYS_INLINE_CAPACITY: usize = 16;
#[derive(Debug, Default, Clone)]
pub struct TxnKeysBuffer {
bytes: Vec<u8>,
offsets: SmallVec<[u32; TXN_KEYS_INLINE_CAPACITY]>,
}
impl TxnKeysBuffer {
#[inline]
pub fn new() -> Self {
Self {
bytes: Vec::new(),
offsets: SmallVec::new(),
}
}
#[inline]
pub fn with_capacity(bytes_capacity: usize, keys_capacity: usize) -> Self {
Self {
bytes: Vec::with_capacity(bytes_capacity),
offsets: SmallVec::with_capacity(keys_capacity),
}
}
#[inline]
pub fn is_empty(&self) -> bool {
self.offsets.is_empty()
}
#[inline]
pub fn len(&self) -> usize {
self.offsets.len()
}
#[inline]
pub fn capacity(&self) -> (usize, usize) {
(self.bytes.capacity(), self.offsets.capacity())
}
#[inline]
pub fn iter(&self) -> TxnKeysIter<'_> {
TxnKeysIter {
bytes: &self.bytes,
offsets: &self.offsets,
idx: 0,
prev: 0,
}
}
#[inline]
pub fn get(&self, index: usize) -> Option<&[u8]> {
let &end = self.offsets.get(index)?;
let start = if index == 0 {
0
} else {
unsafe { *self.offsets.get_unchecked(index - 1) as usize }
};
Some(unsafe { self.bytes.get_unchecked(start..end as usize) })
}
pub fn push(&mut self, key: &[u8]) {
if self.iter().any(|k| k == key) {
return;
}
self.bytes.extend_from_slice(key);
self.offsets.push(self.bytes.len() as u32);
}
#[inline]
pub fn to_smallvec(&self) -> SmallVec<[&[u8]; TXN_KEYS_INLINE_CAPACITY]> {
self.iter().collect()
}
#[inline]
pub fn clear(&mut self) {
self.bytes.clear();
self.offsets.clear();
}
}
pub struct TxnKeysIter<'a> {
bytes: &'a [u8],
offsets: &'a [u32],
idx: usize,
prev: usize,
}
impl<'a> Iterator for TxnKeysIter<'a> {
type Item = &'a [u8];
#[inline]
fn next(&mut self) -> Option<Self::Item> {
let &end = self.offsets.get(self.idx)?;
self.idx += 1;
let end = end as usize;
let start = self.prev;
self.prev = end;
Some(unsafe { self.bytes.get_unchecked(start..end) })
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let rem = self.offsets.len() - self.idx;
(rem, Some(rem))
}
}
impl ExactSizeIterator for TxnKeysIter<'_> {}
impl<'a> IntoIterator for &'a TxnKeysBuffer {
type Item = &'a [u8];
type IntoIter = TxnKeysIter<'a>;
#[inline]
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_txn_keys_buffer_push_and_dedup() {
let mut buf = TxnKeysBuffer::new();
assert!(buf.is_empty());
assert_eq!(buf.len(), 0);
buf.push(b"key1");
buf.push(b"key2");
buf.push(b"key1"); buf.push(b"key3");
assert_eq!(buf.len(), 3);
assert!(!buf.is_empty());
assert_eq!(buf.get(0), Some(&b"key1"[..]));
assert_eq!(buf.get(1), Some(&b"key2"[..]));
assert_eq!(buf.get(2), Some(&b"key3"[..]));
assert_eq!(buf.get(3), None);
let collected: Vec<&[u8]> = buf.iter().collect();
assert_eq!(collected, vec![&b"key1"[..], &b"key2"[..], &b"key3"[..]]);
}
#[test]
fn test_txn_keys_buffer_clear_reuse_capacity() {
let mut buf = TxnKeysBuffer::new();
buf.push(b"first_key_long_enough_to_allocate");
let (byte_cap, offset_cap) = buf.capacity();
buf.clear();
assert!(buf.is_empty());
assert_eq!(buf.len(), 0);
assert_eq!(buf.capacity(), (byte_cap, offset_cap));
buf.push(b"second_key");
assert_eq!(buf.len(), 1);
assert_eq!(buf.get(0), Some(&b"second_key"[..]));
}
#[test]
fn test_txn_keys_buffer_inline_capacity_16() {
let mut buf = TxnKeysBuffer::new();
for i in 0..16 {
let key = format!("k{i}");
buf.push(key.as_bytes());
}
assert_eq!(buf.len(), 16);
assert!(!buf.offsets.spilled(), "16 键以内应保持内联,0 堆分配");
buf.push(b"overflow_17");
assert_eq!(buf.len(), 17);
assert!(buf.offsets.spilled(), "超过 16 键优雅溢出至堆");
}
}