icydb_core/db/store/
key.rs1#![expect(clippy::cast_possible_truncation)]
2use crate::{
3 db::identity::{EntityName, IdentityDecodeError},
4 error::{ErrorClass, ErrorOrigin, InternalError},
5 key::{Key, KeyEncodeError},
6 traits::{EntityKind, Storable},
7};
8use canic_cdk::structures::storable::Bound;
9use std::{
10 borrow::Cow,
11 fmt::{self, Display},
12};
13use thiserror::Error as ThisError;
14
15#[derive(Debug, ThisError)]
21pub enum DataKeyEncodeError {
22 #[error("data key encoding failed for {key}: {source}")]
23 KeyEncoding {
24 key: DataKey,
25 source: KeyEncodeError,
26 },
27}
28
29impl From<DataKeyEncodeError> for InternalError {
30 fn from(err: DataKeyEncodeError) -> Self {
31 Self::new(
32 ErrorClass::Unsupported,
33 ErrorOrigin::Serialize,
34 err.to_string(),
35 )
36 }
37}
38
39#[derive(Debug, ThisError)]
45pub enum KeyDecodeError {
46 #[error("invalid primary key encoding")]
47 InvalidEncoding,
48}
49
50impl From<&'static str> for KeyDecodeError {
51 fn from(_: &'static str) -> Self {
52 Self::InvalidEncoding
53 }
54}
55
56#[derive(Debug, ThisError)]
62pub enum DataKeyDecodeError {
63 #[error("invalid entity name")]
64 Entity(#[from] IdentityDecodeError),
65
66 #[error("invalid primary key")]
67 Key(#[from] KeyDecodeError),
68}
69
70#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
75pub struct DataKey {
76 entity: EntityName,
77 key: Key,
78}
79
80impl DataKey {
81 pub const STORED_SIZE_BYTES: u64 = EntityName::STORED_SIZE_BYTES + Key::STORED_SIZE_BYTES;
83
84 pub const STORED_SIZE_USIZE: usize = Self::STORED_SIZE_BYTES as usize;
86
87 #[must_use]
96 pub fn new<E: EntityKind>(key: impl Into<Key>) -> Self {
97 Self {
98 entity: Self::entity_for::<E>(),
99 key: key.into(),
100 }
101 }
102
103 #[must_use]
104 pub fn lower_bound<E: EntityKind>() -> Self {
105 Self {
106 entity: Self::entity_for::<E>(),
107 key: Key::MIN,
108 }
109 }
110
111 #[must_use]
112 pub fn upper_bound<E: EntityKind>() -> Self {
113 Self {
114 entity: Self::entity_for::<E>(),
115 key: Key::upper_bound(),
116 }
117 }
118
119 #[inline]
120 fn entity_for<E: EntityKind>() -> EntityName {
121 EntityName::try_from_str(E::ENTITY_NAME).unwrap()
124 }
125
126 #[must_use]
131 pub const fn key(&self) -> Key {
132 self.key
133 }
134
135 #[must_use]
136 pub const fn entity_name(&self) -> &EntityName {
137 &self.entity
138 }
139
140 #[must_use]
142 pub const fn entry_size_bytes(value_len: u64) -> u64 {
143 Self::STORED_SIZE_BYTES + value_len
144 }
145
146 #[must_use]
147 pub fn max_storable() -> Self {
148 Self {
149 entity: EntityName::max_storable(),
150 key: Key::max_storable(),
151 }
152 }
153
154 pub fn to_raw(&self) -> Result<RawDataKey, InternalError> {
160 let mut buf = [0u8; Self::STORED_SIZE_USIZE];
161
162 let entity_bytes = self.entity.to_bytes();
163 buf[..EntityName::STORED_SIZE_USIZE].copy_from_slice(&entity_bytes);
164
165 let key_bytes = self
166 .key
167 .to_bytes()
168 .map_err(|err| DataKeyEncodeError::KeyEncoding {
169 key: self.clone(),
170 source: err,
171 })?;
172
173 let key_offset = EntityName::STORED_SIZE_USIZE;
174 buf[key_offset..key_offset + Key::STORED_SIZE_USIZE].copy_from_slice(&key_bytes);
175
176 Ok(RawDataKey(buf))
177 }
178
179 pub fn try_from_raw(raw: &RawDataKey) -> Result<Self, DataKeyDecodeError> {
180 let bytes = &raw.0;
181
182 let entity = EntityName::from_bytes(&bytes[..EntityName::STORED_SIZE_USIZE])?;
183
184 let key = Key::try_from_bytes(&bytes[EntityName::STORED_SIZE_USIZE..])
185 .map_err(KeyDecodeError::from)?;
186
187 Ok(Self { entity, key })
188 }
189}
190
191impl Display for DataKey {
192 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
193 write!(f, "#{} ({})", self.entity, self.key)
194 }
195}
196
197impl From<DataKey> for Key {
198 fn from(key: DataKey) -> Self {
199 key.key()
200 }
201}
202
203#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
208pub struct RawDataKey([u8; DataKey::STORED_SIZE_USIZE]);
209
210impl RawDataKey {
211 #[must_use]
212 pub const fn as_bytes(&self) -> &[u8; DataKey::STORED_SIZE_USIZE] {
213 &self.0
214 }
215}
216
217impl Storable for RawDataKey {
218 fn to_bytes(&self) -> Cow<'_, [u8]> {
219 Cow::Borrowed(&self.0)
220 }
221
222 fn from_bytes(bytes: Cow<'_, [u8]>) -> Self {
223 let mut out = [0u8; DataKey::STORED_SIZE_USIZE];
224 if bytes.len() == out.len() {
225 out.copy_from_slice(bytes.as_ref());
226 }
227 Self(out)
228 }
229
230 fn into_bytes(self) -> Vec<u8> {
231 self.0.to_vec()
232 }
233
234 const BOUND: Bound = Bound::Bounded {
235 max_size: DataKey::STORED_SIZE_BYTES as u32,
236 is_fixed_size: true,
237 };
238}
239
240#[cfg(test)]
245mod tests {
246 use super::*;
247 use std::borrow::Cow;
248
249 #[test]
250 fn data_key_is_exactly_fixed_size() {
251 let data_key = DataKey::max_storable();
252 let size = data_key.to_raw().unwrap().as_bytes().len();
253 assert_eq!(size, DataKey::STORED_SIZE_USIZE);
254 }
255
256 #[test]
257 fn data_key_ordering_matches_bytes() {
258 let keys = vec![
259 DataKey {
260 entity: EntityName::try_from_str("a").unwrap(),
261 key: Key::Int(0),
262 },
263 DataKey {
264 entity: EntityName::try_from_str("aa").unwrap(),
265 key: Key::Int(0),
266 },
267 DataKey {
268 entity: EntityName::try_from_str("b").unwrap(),
269 key: Key::Int(0),
270 },
271 DataKey {
272 entity: EntityName::try_from_str("a").unwrap(),
273 key: Key::Uint(1),
274 },
275 ];
276
277 let mut by_ord = keys.clone();
278 by_ord.sort();
279
280 let mut by_bytes = keys;
281 by_bytes.sort_by(|a, b| {
282 a.to_raw()
283 .unwrap()
284 .as_bytes()
285 .cmp(b.to_raw().unwrap().as_bytes())
286 });
287
288 assert_eq!(by_ord, by_bytes);
289 }
290
291 #[test]
292 fn data_key_rejects_corrupt_entity() {
293 let mut raw = DataKey::max_storable().to_raw().unwrap();
294 raw.0[0] = 0;
295 assert!(DataKey::try_from_raw(&raw).is_err());
296 }
297
298 #[test]
299 fn data_key_rejects_corrupt_key() {
300 let mut raw = DataKey::max_storable().to_raw().unwrap();
301 let off = EntityName::STORED_SIZE_USIZE;
302 raw.0[off] = 0xFF;
303 assert!(DataKey::try_from_raw(&raw).is_err());
304 }
305
306 #[test]
307 #[allow(clippy::cast_possible_truncation)]
308 fn data_key_fuzz_roundtrip_is_canonical() {
309 let mut seed = 0xDEAD_BEEF_u64;
310
311 for _ in 0..1_000 {
312 let mut bytes = [0u8; DataKey::STORED_SIZE_USIZE];
313 for b in &mut bytes {
314 seed = seed.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1);
315 *b = (seed >> 24) as u8;
316 }
317
318 let raw = RawDataKey(bytes);
319 if let Ok(decoded) = DataKey::try_from_raw(&raw) {
320 let re = decoded.to_raw().unwrap();
321 assert_eq!(raw.as_bytes(), re.as_bytes());
322 }
323 }
324 }
325
326 #[test]
327 fn raw_data_key_storable_roundtrip() {
328 let key = DataKey::max_storable().to_raw().unwrap();
329 let bytes = key.to_bytes();
330 let decoded = RawDataKey::from_bytes(Cow::Borrowed(&bytes));
331 assert_eq!(key, decoded);
332 }
333}