Skip to main content

icydb_core/db/store/
key.rs

1use crate::{
2    db::identity::EntityName,
3    error::{ErrorClass, ErrorOrigin, InternalError},
4    key::{Key, KeyEncodeError},
5    traits::Storable,
6};
7use canic_cdk::structures::storable::Bound;
8use std::{
9    borrow::Cow,
10    fmt::{self, Display},
11};
12use thiserror::Error as ThisError;
13
14///
15/// DataKey
16///
17
18#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
19pub struct DataKey {
20    entity: EntityName,
21    key: Key,
22}
23
24///
25/// DataKeyEncodeError
26///
27/// Errors returned when encoding a data key for persistence.
28///
29
30#[derive(Debug, ThisError)]
31pub enum DataKeyEncodeError {
32    #[error("data key encoding failed for {key}: {source}")]
33    KeyEncoding {
34        key: DataKey,
35        source: KeyEncodeError,
36    },
37}
38
39impl From<DataKeyEncodeError> for InternalError {
40    fn from(err: DataKeyEncodeError) -> Self {
41        Self::new(
42            ErrorClass::Unsupported,
43            ErrorOrigin::Serialize,
44            err.to_string(),
45        )
46    }
47}
48
49impl DataKey {
50    #[allow(clippy::cast_possible_truncation)]
51    pub const STORED_SIZE: u32 = EntityName::STORED_SIZE + Key::STORED_SIZE as u32;
52
53    #[must_use]
54    /// Build a data key for the given entity type and primary key.
55    pub fn new<E: crate::traits::EntityKind>(key: impl Into<Key>) -> Self {
56        Self {
57            entity: EntityName::from_static_unchecked(E::ENTITY_NAME),
58            key: key.into(),
59        }
60    }
61
62    #[must_use]
63    pub const fn lower_bound<E: crate::traits::EntityKind>() -> Self {
64        Self {
65            entity: EntityName::from_static_unchecked(E::ENTITY_NAME),
66            key: Key::MIN,
67        }
68    }
69
70    #[must_use]
71    pub const fn upper_bound<E: crate::traits::EntityKind>() -> Self {
72        Self {
73            entity: EntityName::from_static_unchecked(E::ENTITY_NAME),
74            key: Key::upper_bound(),
75        }
76    }
77
78    /// Return the primary key component of this data key.
79    #[must_use]
80    pub const fn key(&self) -> Key {
81        self.key
82    }
83
84    /// Entity name (stable, compile-time constant per entity type).
85    #[must_use]
86    pub const fn entity_name(&self) -> &EntityName {
87        &self.entity
88    }
89
90    /// Compute the on-disk size used by a single data entry from its value length.
91    /// Includes the bounded `DataKey` size and the value bytes.
92    #[must_use]
93    pub const fn entry_size_bytes(value_len: u64) -> u64 {
94        Self::STORED_SIZE as u64 + value_len
95    }
96
97    #[must_use]
98    /// Max sentinel key for sizing.
99    pub fn max_storable() -> Self {
100        Self {
101            entity: EntityName::max_storable(),
102            key: Key::max_storable(),
103        }
104    }
105
106    /// Encode this data key into the fixed-size on-disk representation.
107    pub fn to_raw(&self) -> Result<RawDataKey, InternalError> {
108        let mut buf = [0u8; Self::STORED_SIZE as usize];
109
110        buf[0] = self.entity.len;
111        let entity_end = EntityName::STORED_SIZE_USIZE;
112        buf[1..entity_end].copy_from_slice(&self.entity.bytes);
113
114        let key_bytes = self
115            .key
116            .to_bytes()
117            .map_err(|err| DataKeyEncodeError::KeyEncoding {
118                key: self.clone(),
119                source: err,
120            })?;
121        let key_offset = EntityName::STORED_SIZE_USIZE;
122        buf[key_offset..key_offset + Key::STORED_SIZE].copy_from_slice(&key_bytes);
123
124        Ok(RawDataKey(buf))
125    }
126
127    pub fn try_from_raw(raw: &RawDataKey) -> Result<Self, &'static str> {
128        let bytes = &raw.0;
129
130        let mut offset = 0;
131        let entity = EntityName::from_bytes(&bytes[offset..offset + EntityName::STORED_SIZE_USIZE])
132            .map_err(|_| "corrupted DataKey: invalid EntityName bytes")?;
133        offset += EntityName::STORED_SIZE_USIZE;
134
135        let key = Key::try_from_bytes(&bytes[offset..offset + Key::STORED_SIZE])
136            .map_err(|_| "corrupted DataKey: invalid Key bytes")?;
137
138        Ok(Self { entity, key })
139    }
140}
141
142impl Display for DataKey {
143    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
144        write!(f, "#{} ({})", self.entity, self.key)
145    }
146}
147
148impl From<DataKey> for Key {
149    fn from(key: DataKey) -> Self {
150        key.key()
151    }
152}
153
154///
155/// RawDataKey
156///
157
158#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
159pub struct RawDataKey([u8; DataKey::STORED_SIZE as usize]);
160
161impl RawDataKey {
162    #[must_use]
163    pub const fn as_bytes(&self) -> &[u8; DataKey::STORED_SIZE as usize] {
164        &self.0
165    }
166}
167
168impl Storable for RawDataKey {
169    fn to_bytes(&self) -> Cow<'_, [u8]> {
170        Cow::Borrowed(&self.0)
171    }
172
173    fn from_bytes(bytes: Cow<'_, [u8]>) -> Self {
174        let mut out = [0u8; DataKey::STORED_SIZE as usize];
175        if bytes.len() == out.len() {
176            out.copy_from_slice(bytes.as_ref());
177        }
178        Self(out)
179    }
180
181    fn into_bytes(self) -> Vec<u8> {
182        self.0.to_vec()
183    }
184
185    const BOUND: Bound = Bound::Bounded {
186        max_size: DataKey::STORED_SIZE,
187        is_fixed_size: true,
188    };
189}
190
191///
192/// TESTS
193///
194
195#[cfg(test)]
196mod tests {
197    use super::*;
198    use std::borrow::Cow;
199
200    #[test]
201    fn data_key_is_exactly_fixed_size() {
202        let data_key = DataKey::max_storable();
203        let size = data_key.to_raw().expect("data key encode").as_bytes().len();
204
205        assert_eq!(
206            size,
207            DataKey::STORED_SIZE as usize,
208            "DataKey must serialize to exactly STORED_SIZE bytes"
209        );
210    }
211
212    #[test]
213    fn data_key_ordering_matches_bytes() {
214        let keys = vec![
215            DataKey {
216                entity: EntityName::from_static_unchecked("a"),
217                key: Key::Int(0),
218            },
219            DataKey {
220                entity: EntityName::from_static_unchecked("aa"),
221                key: Key::Int(0),
222            },
223            DataKey {
224                entity: EntityName::from_static_unchecked("b"),
225                key: Key::Int(0),
226            },
227            DataKey {
228                entity: EntityName::from_static_unchecked("a"),
229                key: Key::Uint(1),
230            },
231        ];
232
233        let mut sorted_by_ord = keys.clone();
234        sorted_by_ord.sort();
235
236        let mut sorted_by_bytes = keys;
237        sorted_by_bytes.sort_by(|a, b| {
238            a.to_raw()
239                .expect("data key encode")
240                .as_bytes()
241                .cmp(b.to_raw().expect("data key encode").as_bytes())
242        });
243
244        assert_eq!(
245            sorted_by_ord, sorted_by_bytes,
246            "DataKey Ord and byte ordering diverged"
247        );
248    }
249
250    #[test]
251    fn data_key_rejects_undersized_bytes() {
252        let buf = vec![0u8; DataKey::STORED_SIZE as usize - 1];
253        let raw = RawDataKey::from_bytes(Cow::Borrowed(&buf));
254        let err = DataKey::try_from_raw(&raw).unwrap_err();
255        assert!(
256            err.contains("corrupted"),
257            "expected corruption error, got: {err}"
258        );
259    }
260
261    #[test]
262    fn data_key_rejects_oversized_bytes() {
263        let buf = vec![0u8; DataKey::STORED_SIZE as usize + 1];
264        let raw = RawDataKey::from_bytes(Cow::Borrowed(&buf));
265        let err = DataKey::try_from_raw(&raw).unwrap_err();
266        assert!(
267            err.contains("corrupted"),
268            "expected corruption error, got: {err}"
269        );
270    }
271
272    #[test]
273    fn data_key_rejects_invalid_entity_len() {
274        let mut raw = DataKey::max_storable().to_raw().expect("data key encode");
275        raw.0[0] = 0;
276        assert!(DataKey::try_from_raw(&raw).is_err());
277    }
278
279    #[test]
280    fn data_key_rejects_non_ascii_entity_bytes() {
281        let data_key = DataKey {
282            entity: EntityName::from_static_unchecked("a"),
283            key: Key::Int(1),
284        };
285        let mut raw = data_key.to_raw().expect("data key encode");
286        raw.0[1] = 0xFF;
287        assert!(DataKey::try_from_raw(&raw).is_err());
288    }
289
290    #[test]
291    fn data_key_rejects_entity_padding() {
292        let data_key = DataKey {
293            entity: EntityName::from_static_unchecked("user"),
294            key: Key::Int(1),
295        };
296        let mut raw = data_key.to_raw().expect("data key encode");
297        let padding_offset = 1 + data_key.entity.len();
298        raw.0[padding_offset] = b'x';
299        assert!(DataKey::try_from_raw(&raw).is_err());
300    }
301
302    #[test]
303    #[allow(clippy::cast_possible_truncation)]
304    fn data_key_decode_fuzz_roundtrip_is_canonical() {
305        const RUNS: u64 = 1_000;
306
307        let mut seed = 0xDEAD_BEEF_u64;
308        for _ in 0..RUNS {
309            let mut bytes = [0u8; DataKey::STORED_SIZE as usize];
310            for b in &mut bytes {
311                seed = seed.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1);
312                *b = (seed >> 24) as u8;
313            }
314
315            let raw = RawDataKey(bytes);
316            if let Ok(decoded) = DataKey::try_from_raw(&raw) {
317                let re = decoded.to_raw().expect("data key encode");
318                assert_eq!(
319                    raw.as_bytes(),
320                    re.as_bytes(),
321                    "decoded DataKey must be canonical"
322                );
323            }
324        }
325    }
326}