Skip to main content

icydb_core/key/
mod.rs

1mod convert;
2#[cfg(test)]
3mod tests;
4
5use crate::{
6    error::{ErrorClass, ErrorOrigin, InternalError},
7    types::{
8        Account, AccountEncodeError, Principal, PrincipalEncodeError, Subaccount, Timestamp, Ulid,
9    },
10};
11use candid::CandidType;
12use derive_more::Display;
13use serde::{Deserialize, Serialize};
14use std::cmp::Ordering;
15use thiserror::Error as ThisError;
16
17///
18/// Key
19///
20/// Treating IndexKey as the atomic, normalized unit of the keyspace
21/// Backing primary keys and secondary indexes with the same value representation
22/// Planning to enforce Copy semantics (i.e., fast, clean, safe)
23///
24
25#[derive(CandidType, Clone, Copy, Debug, Deserialize, Display, Eq, Hash, PartialEq, Serialize)]
26pub enum Key {
27    Account(Account),
28    Int(i64),
29    Principal(Principal),
30    Subaccount(Subaccount),
31    Timestamp(Timestamp),
32    Uint(u64),
33    Ulid(Ulid),
34    Unit,
35}
36
37///
38/// KeyEncodeError
39///
40/// Errors returned when encoding a key for persistence.
41///
42
43#[derive(Debug, ThisError)]
44pub enum KeyEncodeError {
45    #[error("account encoding failed: {0}")]
46    Account(#[from] AccountEncodeError),
47
48    #[error("account payload length mismatch: {len} bytes (expected {expected})")]
49    AccountLengthMismatch { len: usize, expected: usize },
50
51    #[error("principal encoding failed: {0}")]
52    Principal(#[from] PrincipalEncodeError),
53}
54
55impl From<KeyEncodeError> for InternalError {
56    fn from(err: KeyEncodeError) -> Self {
57        Self::new(
58            ErrorClass::Unsupported,
59            ErrorOrigin::Serialize,
60            err.to_string(),
61        )
62    }
63}
64
65impl Key {
66    // ── Variant tags (do not reorder) ─────────────────
67    pub(crate) const TAG_ACCOUNT: u8 = 0;
68    pub(crate) const TAG_INT: u8 = 1;
69    pub(crate) const TAG_PRINCIPAL: u8 = 2;
70    pub(crate) const TAG_SUBACCOUNT: u8 = 3;
71    pub(crate) const TAG_TIMESTAMP: u8 = 4;
72    pub(crate) const TAG_UINT: u8 = 5;
73    pub(crate) const TAG_ULID: u8 = 6;
74    pub(crate) const TAG_UNIT: u8 = 7;
75
76    /// Fixed serialized size in bytes (stable, protocol-level)
77    /// DO NOT CHANGE without migration.
78    pub const STORED_SIZE_BYTES: u64 = 64;
79
80    /// Fixed in-memory size (for buffers and indexing only)
81    #[expect(clippy::cast_possible_truncation)]
82    pub const STORED_SIZE_USIZE: usize = Self::STORED_SIZE_BYTES as usize;
83
84    // ── Layout ─────────────────────────────────────
85    const TAG_SIZE: usize = 1;
86    pub(crate) const TAG_OFFSET: usize = 0;
87
88    pub(crate) const PAYLOAD_OFFSET: usize = Self::TAG_SIZE;
89    const PAYLOAD_SIZE: usize = Self::STORED_SIZE_USIZE - Self::TAG_SIZE;
90
91    // ── Payload sizes ──────────────────────────────
92    pub(crate) const INT_SIZE: usize = 8;
93    pub(crate) const UINT_SIZE: usize = 8;
94    pub(crate) const TIMESTAMP_SIZE: usize = 8;
95    pub(crate) const ULID_SIZE: usize = 16;
96    pub(crate) const SUBACCOUNT_SIZE: usize = 32;
97    const ACCOUNT_MAX_SIZE: usize = 62;
98
99    const fn tag(&self) -> u8 {
100        match self {
101            Self::Account(_) => Self::TAG_ACCOUNT,
102            Self::Int(_) => Self::TAG_INT,
103            Self::Principal(_) => Self::TAG_PRINCIPAL,
104            Self::Subaccount(_) => Self::TAG_SUBACCOUNT,
105            Self::Timestamp(_) => Self::TAG_TIMESTAMP,
106            Self::Uint(_) => Self::TAG_UINT,
107            Self::Ulid(_) => Self::TAG_ULID,
108            Self::Unit => Self::TAG_UNIT,
109        }
110    }
111
112    #[must_use]
113    /// Sentinel key representing the maximum storable account value.
114    pub fn max_storable() -> Self {
115        Self::Account(Account::max_storable())
116    }
117
118    /// Global minimum key for scan bounds and key range construction.
119    pub const MIN: Self = Self::Account(Account {
120        owner: Principal::from_slice(&[]),
121        subaccount: None,
122    });
123
124    #[must_use]
125    pub const fn lower_bound() -> Self {
126        Self::MIN
127    }
128
129    #[must_use]
130    pub const fn upper_bound() -> Self {
131        Self::Unit
132    }
133
134    const fn variant_rank(&self) -> u8 {
135        self.tag()
136    }
137
138    /// Encode this key into its fixed-size storage representation.
139    pub fn to_bytes(&self) -> Result<[u8; Self::STORED_SIZE_USIZE], KeyEncodeError> {
140        let mut buf = [0u8; Self::STORED_SIZE_USIZE];
141
142        // ── Tag ─────────────────────────────────────
143        buf[Self::TAG_OFFSET] = self.tag();
144        let payload = &mut buf[Self::PAYLOAD_OFFSET..=Self::PAYLOAD_SIZE];
145
146        // ── Payload ─────────────────────────────────
147        #[allow(clippy::cast_possible_truncation)]
148        match self {
149            Self::Account(v) => {
150                let bytes = v.to_bytes()?;
151                if bytes.len() != Self::ACCOUNT_MAX_SIZE {
152                    return Err(KeyEncodeError::AccountLengthMismatch {
153                        len: bytes.len(),
154                        expected: Self::ACCOUNT_MAX_SIZE,
155                    });
156                }
157                payload[..bytes.len()].copy_from_slice(&bytes);
158            }
159
160            Self::Int(v) => {
161                // Flip sign bit to preserve ordering in lexicographic bytes.
162                let biased = (*v).cast_unsigned() ^ (1u64 << 63);
163                payload[..Self::INT_SIZE].copy_from_slice(&biased.to_be_bytes());
164            }
165
166            Self::Uint(v) => {
167                payload[..Self::UINT_SIZE].copy_from_slice(&v.to_be_bytes());
168            }
169
170            Self::Timestamp(v) => {
171                payload[..Self::TIMESTAMP_SIZE].copy_from_slice(&v.get().to_be_bytes());
172            }
173
174            Self::Principal(v) => {
175                let bytes = v.to_bytes()?;
176                let len = bytes.len();
177                payload[0] = u8::try_from(len).map_err(|_| {
178                    KeyEncodeError::Principal(PrincipalEncodeError::TooLarge {
179                        len,
180                        max: Principal::MAX_LENGTH_IN_BYTES as usize,
181                    })
182                })?;
183                if len > 0 {
184                    payload[1..=len].copy_from_slice(&bytes[..len]);
185                }
186            }
187
188            Self::Subaccount(v) => {
189                let bytes = v.to_array();
190                payload[..Self::SUBACCOUNT_SIZE].copy_from_slice(&bytes);
191            }
192
193            Self::Ulid(v) => {
194                payload[..Self::ULID_SIZE].copy_from_slice(&v.to_bytes());
195            }
196
197            Self::Unit => {}
198        }
199
200        Ok(buf)
201    }
202
203    pub fn try_from_bytes(bytes: &[u8]) -> Result<Self, &'static str> {
204        if bytes.len() != Self::STORED_SIZE_USIZE {
205            return Err("corrupted Key: invalid size");
206        }
207
208        let tag = bytes[Self::TAG_OFFSET];
209        let payload = &bytes[Self::PAYLOAD_OFFSET..=Self::PAYLOAD_SIZE];
210
211        let ensure_zero_padding = |used: usize, context: &str| {
212            if payload[used..].iter().all(|&b| b == 0) {
213                Ok(())
214            } else {
215                Err(match context {
216                    "account" => "corrupted Key: non-zero account padding",
217                    "int" => "corrupted Key: non-zero int padding",
218                    "principal" => "corrupted Key: non-zero principal padding",
219                    "subaccount" => "corrupted Key: non-zero subaccount padding",
220                    "timestamp" => "corrupted Key: non-zero timestamp padding",
221                    "uint" => "corrupted Key: non-zero uint padding",
222                    "ulid" => "corrupted Key: non-zero ulid padding",
223                    "unit" => "corrupted Key: non-zero unit padding",
224                    _ => "corrupted Key: non-zero padding",
225                })
226            }
227        };
228
229        match tag {
230            Self::TAG_ACCOUNT => {
231                let end = Account::STORED_SIZE as usize;
232                ensure_zero_padding(end, "account")?;
233                let account = Account::try_from_bytes(&payload[..end])?;
234                Ok(Self::Account(account))
235            }
236
237            Self::TAG_INT => {
238                let mut buf = [0u8; Self::INT_SIZE];
239                buf.copy_from_slice(&payload[..Self::INT_SIZE]);
240                let biased = u64::from_be_bytes(buf);
241                ensure_zero_padding(Self::INT_SIZE, "int")?;
242                Ok(Self::Int((biased ^ (1u64 << 63)).cast_signed()))
243            }
244
245            Self::TAG_PRINCIPAL => {
246                let len = payload[0] as usize;
247                if len > Principal::MAX_LENGTH_IN_BYTES as usize {
248                    return Err("corrupted Key: invalid principal length");
249                }
250                let end = 1 + len;
251                ensure_zero_padding(end, "principal")?;
252                Ok(Self::Principal(Principal::from_slice(&payload[1..end])))
253            }
254
255            Self::TAG_SUBACCOUNT => {
256                let mut buf = [0u8; Self::SUBACCOUNT_SIZE];
257                buf.copy_from_slice(&payload[..Self::SUBACCOUNT_SIZE]);
258                ensure_zero_padding(Self::SUBACCOUNT_SIZE, "subaccount")?;
259                Ok(Self::Subaccount(Subaccount::from_array(buf)))
260            }
261
262            Self::TAG_TIMESTAMP => {
263                let mut buf = [0u8; Self::TIMESTAMP_SIZE];
264                buf.copy_from_slice(&payload[..Self::TIMESTAMP_SIZE]);
265                ensure_zero_padding(Self::TIMESTAMP_SIZE, "timestamp")?;
266                Ok(Self::Timestamp(Timestamp::from(u64::from_be_bytes(buf))))
267            }
268
269            Self::TAG_UINT => {
270                let mut buf = [0u8; Self::UINT_SIZE];
271                buf.copy_from_slice(&payload[..Self::UINT_SIZE]);
272                ensure_zero_padding(Self::UINT_SIZE, "uint")?;
273                Ok(Self::Uint(u64::from_be_bytes(buf)))
274            }
275
276            Self::TAG_ULID => {
277                let mut buf = [0u8; Self::ULID_SIZE];
278                buf.copy_from_slice(&payload[..Self::ULID_SIZE]);
279                ensure_zero_padding(Self::ULID_SIZE, "ulid")?;
280                Ok(Self::Ulid(Ulid::from_bytes(buf)))
281            }
282
283            Self::TAG_UNIT => {
284                ensure_zero_padding(0, "unit")?;
285                Ok(Self::Unit)
286            }
287
288            _ => Err("corrupted Key: invalid tag"),
289        }
290    }
291}
292
293impl TryFrom<&[u8]> for Key {
294    type Error = &'static str;
295
296    fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
297        Self::try_from_bytes(bytes)
298    }
299}
300
301impl Ord for Key {
302    fn cmp(&self, other: &Self) -> Ordering {
303        match (self, other) {
304            (Self::Account(a), Self::Account(b)) => Ord::cmp(a, b),
305            (Self::Int(a), Self::Int(b)) => Ord::cmp(a, b),
306            (Self::Principal(a), Self::Principal(b)) => Ord::cmp(a, b),
307            (Self::Uint(a), Self::Uint(b)) => Ord::cmp(a, b),
308            (Self::Ulid(a), Self::Ulid(b)) => Ord::cmp(a, b),
309            (Self::Subaccount(a), Self::Subaccount(b)) => Ord::cmp(a, b),
310            (Self::Timestamp(a), Self::Timestamp(b)) => Ord::cmp(a, b),
311
312            _ => Ord::cmp(&self.variant_rank(), &other.variant_rank()), // fallback for cross-type comparison
313        }
314    }
315}
316
317impl PartialOrd for Key {
318    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
319        Some(Ord::cmp(self, other))
320    }
321}