Skip to main content

icydb_core/db/store/
key.rs

1#![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///
16/// DataKeyEncodeError
17/// (serialize boundary)
18///
19
20#[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///
40/// KeyDecodeError
41/// (decode / corruption boundary)
42///
43
44#[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///
57/// DataKeyDecodeError
58/// (decode / corruption boundary)
59///
60
61#[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///
71/// DataKey
72///
73
74#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
75pub struct DataKey {
76    entity: EntityName,
77    key: Key,
78}
79
80impl DataKey {
81    /// Fixed on-disk size in bytes (stable, protocol-level)
82    pub const STORED_SIZE_BYTES: u64 = EntityName::STORED_SIZE_BYTES + Key::STORED_SIZE_BYTES;
83
84    /// Fixed in-memory size (for buffers and arrays only)
85    pub const STORED_SIZE_USIZE: usize = Self::STORED_SIZE_BYTES as usize;
86
87    // ------------------------------------------------------------------
88    // Constructors
89    // ------------------------------------------------------------------
90
91    /// Construct using compile-time entity metadata.
92    ///
93    /// This cannot fail in practice: `ENTITY_NAME` is generated and validated
94    /// at compile time. If it is ever invalid, that is a programmer error.
95    #[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        // SAFETY: ENTITY_NAME is generated code and guaranteed valid.
122        // A failure here indicates a schema/codegen bug, not runtime input.
123        EntityName::try_from_str(E::ENTITY_NAME).unwrap()
124    }
125
126    // ------------------------------------------------------------------
127    // Accessors
128    // ------------------------------------------------------------------
129
130    #[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    /// Compute on-disk entry size from value length.
141    #[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    // ------------------------------------------------------------------
155    // Encoding / decoding
156    // ------------------------------------------------------------------
157
158    /// Encode into fixed-size on-disk representation.
159    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///
204/// RawDataKey
205///
206
207#[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///
241/// TESTS
242///
243
244#[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}