Skip to main content

reifydb_codec/key/
buf.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use crate::key::encoded::{EncodedKey, INLINE_CAP};
5
6pub enum KeyBuf {
7	Inline {
8		len: u8,
9		buf: [u8; INLINE_CAP],
10	},
11	Spill(Vec<u8>),
12}
13
14impl KeyBuf {
15	pub fn new() -> Self {
16		KeyBuf::Inline {
17			len: 0,
18			buf: [0u8; INLINE_CAP],
19		}
20	}
21
22	pub fn with_capacity(capacity: usize) -> Self {
23		if capacity <= INLINE_CAP {
24			Self::new()
25		} else {
26			KeyBuf::Spill(Vec::with_capacity(capacity))
27		}
28	}
29
30	pub fn as_slice(&self) -> &[u8] {
31		match self {
32			KeyBuf::Inline {
33				len,
34				buf,
35			} => &buf[..*len as usize],
36			KeyBuf::Spill(v) => v.as_slice(),
37		}
38	}
39
40	pub fn len(&self) -> usize {
41		match self {
42			KeyBuf::Inline {
43				len,
44				..
45			} => *len as usize,
46			KeyBuf::Spill(v) => v.len(),
47		}
48	}
49
50	pub fn is_empty(&self) -> bool {
51		self.len() == 0
52	}
53
54	pub fn push(&mut self, byte: u8) {
55		match self {
56			KeyBuf::Inline {
57				len,
58				buf,
59			} => {
60				let cur = *len as usize;
61				if cur < INLINE_CAP {
62					buf[cur] = byte;
63					*len += 1;
64					return;
65				}
66				let mut vec = Vec::with_capacity(cur + 1);
67				vec.extend_from_slice(&buf[..cur]);
68				vec.push(byte);
69				*self = KeyBuf::Spill(vec);
70			}
71			KeyBuf::Spill(v) => v.push(byte),
72		}
73	}
74
75	pub fn extend_from_slice(&mut self, slice: &[u8]) {
76		match self {
77			KeyBuf::Inline {
78				len,
79				buf,
80			} => {
81				let cur = *len as usize;
82				let total = cur + slice.len();
83				if total <= INLINE_CAP {
84					buf[cur..total].copy_from_slice(slice);
85					*len = total as u8;
86					return;
87				}
88				let mut vec = Vec::with_capacity(total);
89				vec.extend_from_slice(&buf[..cur]);
90				vec.extend_from_slice(slice);
91				*self = KeyBuf::Spill(vec);
92			}
93			KeyBuf::Spill(v) => v.extend_from_slice(slice),
94		}
95	}
96
97	pub fn finish(self) -> EncodedKey {
98		match self {
99			KeyBuf::Inline {
100				len,
101				buf,
102			} => EncodedKey::Inline {
103				len,
104				buf,
105			},
106			KeyBuf::Spill(v) => EncodedKey::new(v),
107		}
108	}
109}
110
111impl Default for KeyBuf {
112	fn default() -> Self {
113		Self::new()
114	}
115}