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(crate) fn as_mut_slice(&mut self) -> &mut [u8] {
41		match self {
42			KeyBuf::Inline {
43				len,
44				buf,
45			} => &mut buf[..*len as usize],
46			KeyBuf::Spill(v) => v.as_mut_slice(),
47		}
48	}
49
50	pub fn len(&self) -> usize {
51		match self {
52			KeyBuf::Inline {
53				len,
54				..
55			} => *len as usize,
56			KeyBuf::Spill(v) => v.len(),
57		}
58	}
59
60	pub fn is_empty(&self) -> bool {
61		self.len() == 0
62	}
63
64	pub fn push(&mut self, byte: u8) {
65		match self {
66			KeyBuf::Inline {
67				len,
68				buf,
69			} => {
70				let cur = *len as usize;
71				if cur < INLINE_CAP {
72					buf[cur] = byte;
73					*len += 1;
74					return;
75				}
76				let mut vec = Vec::with_capacity(cur + 1);
77				vec.extend_from_slice(&buf[..cur]);
78				vec.push(byte);
79				*self = KeyBuf::Spill(vec);
80			}
81			KeyBuf::Spill(v) => v.push(byte),
82		}
83	}
84
85	pub fn extend_from_slice(&mut self, slice: &[u8]) {
86		match self {
87			KeyBuf::Inline {
88				len,
89				buf,
90			} => {
91				let cur = *len as usize;
92				let total = cur + slice.len();
93				if total <= INLINE_CAP {
94					buf[cur..total].copy_from_slice(slice);
95					*len = total as u8;
96					return;
97				}
98				let mut vec = Vec::with_capacity(total);
99				vec.extend_from_slice(&buf[..cur]);
100				vec.extend_from_slice(slice);
101				*self = KeyBuf::Spill(vec);
102			}
103			KeyBuf::Spill(v) => v.extend_from_slice(slice),
104		}
105	}
106
107	pub fn finish(self) -> EncodedKey {
108		match self {
109			KeyBuf::Inline {
110				len,
111				buf,
112			} => EncodedKey::Inline {
113				len,
114				buf,
115			},
116			KeyBuf::Spill(v) => EncodedKey::new(v),
117		}
118	}
119}
120
121impl Default for KeyBuf {
122	fn default() -> Self {
123		Self::new()
124	}
125}