icydb_core/key/
core.rs

1use crate::types::{Account, Principal, Subaccount, Timestamp, Ulid};
2use candid::CandidType;
3use derive_more::Display;
4use serde::{Deserialize, Serialize};
5use std::cmp::Ordering;
6
7///
8/// Key
9///
10/// Treating IndexKey as the atomic, normalized unit of the keyspace
11/// Backing primary keys and secondary indexes with the same value representation
12/// Planning to enforce Copy semantics (i.e., fast, clean, safe)
13///
14
15#[derive(CandidType, Clone, Copy, Debug, Deserialize, Display, Eq, Hash, PartialEq, Serialize)]
16pub enum Key {
17    Account(Account),
18    Int(i64),
19    Principal(Principal),
20    Subaccount(Subaccount),
21    Timestamp(Timestamp),
22    Uint(u64),
23    Ulid(Ulid),
24    Unit,
25}
26
27impl Key {
28    // ── Variant tags (do not reorder) ─────────────────
29    pub(crate) const TAG_ACCOUNT: u8 = 0;
30    pub(crate) const TAG_INT: u8 = 1;
31    pub(crate) const TAG_PRINCIPAL: u8 = 2;
32    pub(crate) const TAG_SUBACCOUNT: u8 = 3;
33    pub(crate) const TAG_TIMESTAMP: u8 = 4;
34    pub(crate) const TAG_UINT: u8 = 5;
35    pub(crate) const TAG_ULID: u8 = 6;
36    pub(crate) const TAG_UNIT: u8 = 7;
37
38    /// Fixed serialized size (do not change without migration)
39    pub const STORED_SIZE: usize = 64;
40
41    // ── Layout ─────────────────────────────────────
42    const TAG_SIZE: usize = 1;
43    pub(crate) const TAG_OFFSET: usize = 0;
44
45    pub(crate) const PAYLOAD_OFFSET: usize = Self::TAG_SIZE;
46    const PAYLOAD_SIZE: usize = Self::STORED_SIZE - Self::TAG_SIZE;
47
48    // ── Payload sizes ──────────────────────────────
49    pub(crate) const INT_SIZE: usize = 8;
50    pub(crate) const UINT_SIZE: usize = 8;
51    pub(crate) const TIMESTAMP_SIZE: usize = 8;
52    pub(crate) const ULID_SIZE: usize = 16;
53    pub(crate) const SUBACCOUNT_SIZE: usize = 32;
54    const ACCOUNT_MAX_SIZE: usize = 62;
55
56    const fn tag(&self) -> u8 {
57        match self {
58            Self::Account(_) => Self::TAG_ACCOUNT,
59            Self::Int(_) => Self::TAG_INT,
60            Self::Principal(_) => Self::TAG_PRINCIPAL,
61            Self::Subaccount(_) => Self::TAG_SUBACCOUNT,
62            Self::Timestamp(_) => Self::TAG_TIMESTAMP,
63            Self::Uint(_) => Self::TAG_UINT,
64            Self::Ulid(_) => Self::TAG_ULID,
65            Self::Unit => Self::TAG_UNIT,
66        }
67    }
68
69    #[must_use]
70    /// Sentinel key representing the maximum storable account value.
71    pub fn max_storable() -> Self {
72        Self::Account(Account::max_storable())
73    }
74
75    #[must_use]
76    pub const fn lower_bound() -> Self {
77        Self::Int(i64::MIN)
78    }
79
80    #[must_use]
81    pub const fn upper_bound() -> Self {
82        Self::Unit
83    }
84
85    const fn variant_rank(&self) -> u8 {
86        self.tag()
87    }
88
89    #[must_use]
90    pub fn to_bytes(&self) -> [u8; Self::STORED_SIZE] {
91        let mut buf = [0u8; Self::STORED_SIZE];
92
93        // ── Tag ─────────────────────────────────────
94        buf[Self::TAG_OFFSET] = self.tag();
95        let payload = &mut buf[Self::PAYLOAD_OFFSET..];
96
97        debug_assert_eq!(payload.len(), Self::PAYLOAD_SIZE);
98
99        // ── Payload ─────────────────────────────────
100        #[allow(clippy::cast_possible_truncation)]
101        match self {
102            Self::Account(v) => {
103                let bytes = v.to_bytes();
104                debug_assert_eq!(bytes.len(), Self::ACCOUNT_MAX_SIZE);
105                payload[..bytes.len()].copy_from_slice(&bytes);
106            }
107
108            Self::Int(v) => {
109                // Flip sign bit to preserve ordering in lexicographic bytes.
110                let biased = (*v).cast_unsigned() ^ (1u64 << 63);
111                payload[..Self::INT_SIZE].copy_from_slice(&biased.to_be_bytes());
112            }
113
114            Self::Uint(v) => {
115                payload[..Self::UINT_SIZE].copy_from_slice(&v.to_be_bytes());
116            }
117
118            Self::Timestamp(v) => {
119                payload[..Self::TIMESTAMP_SIZE].copy_from_slice(&v.get().to_be_bytes());
120            }
121
122            Self::Principal(v) => {
123                let bytes = v.to_bytes();
124                let len = bytes.len();
125                assert!(
126                    (1..=Principal::MAX_LENGTH_IN_BYTES as usize).contains(&len),
127                    "invalid Key principal length"
128                );
129                payload[0] = len as u8;
130                if len > 0 {
131                    payload[1..=len].copy_from_slice(&bytes[..len]);
132                }
133            }
134
135            Self::Subaccount(v) => {
136                let bytes = v.to_array();
137                debug_assert_eq!(bytes.len(), Self::SUBACCOUNT_SIZE);
138                payload[..Self::SUBACCOUNT_SIZE].copy_from_slice(&bytes);
139            }
140
141            Self::Ulid(v) => {
142                payload[..Self::ULID_SIZE].copy_from_slice(&v.to_bytes());
143            }
144
145            Self::Unit => {}
146        }
147
148        buf
149    }
150
151    pub fn try_from_bytes(bytes: &[u8]) -> Result<Self, &'static str> {
152        if bytes.len() != Self::STORED_SIZE {
153            return Err("corrupted Key: invalid size");
154        }
155
156        let tag = bytes[Self::TAG_OFFSET];
157        let payload = &bytes[Self::PAYLOAD_OFFSET..];
158
159        let ensure_zero_padding = |used: usize, context: &str| {
160            if payload[used..].iter().all(|&b| b == 0) {
161                Ok(())
162            } else {
163                Err(match context {
164                    "account" => "corrupted Key: non-zero account padding",
165                    "int" => "corrupted Key: non-zero int padding",
166                    "principal" => "corrupted Key: non-zero principal padding",
167                    "subaccount" => "corrupted Key: non-zero subaccount padding",
168                    "timestamp" => "corrupted Key: non-zero timestamp padding",
169                    "uint" => "corrupted Key: non-zero uint padding",
170                    "ulid" => "corrupted Key: non-zero ulid padding",
171                    "unit" => "corrupted Key: non-zero unit padding",
172                    _ => "corrupted Key: non-zero padding",
173                })
174            }
175        };
176
177        match tag {
178            Self::TAG_ACCOUNT => {
179                let end = Account::STORED_SIZE as usize;
180                ensure_zero_padding(end, "account")?;
181                let account = Account::try_from_bytes(&payload[..end])?;
182                Ok(Self::Account(account))
183            }
184
185            Self::TAG_INT => {
186                let mut buf = [0u8; Self::INT_SIZE];
187                buf.copy_from_slice(&payload[..Self::INT_SIZE]);
188                let biased = u64::from_be_bytes(buf);
189                ensure_zero_padding(Self::INT_SIZE, "int")?;
190                Ok(Self::Int((biased ^ (1u64 << 63)).cast_signed()))
191            }
192
193            Self::TAG_PRINCIPAL => {
194                let len = payload[0] as usize;
195                if !(1..=Principal::MAX_LENGTH_IN_BYTES as usize).contains(&len) {
196                    return Err("corrupted Key: invalid principal length");
197                }
198                let end = 1 + len;
199                ensure_zero_padding(end, "principal")?;
200                Ok(Self::Principal(Principal::from_slice(&payload[1..end])))
201            }
202
203            Self::TAG_SUBACCOUNT => {
204                let mut buf = [0u8; Self::SUBACCOUNT_SIZE];
205                buf.copy_from_slice(&payload[..Self::SUBACCOUNT_SIZE]);
206                ensure_zero_padding(Self::SUBACCOUNT_SIZE, "subaccount")?;
207                Ok(Self::Subaccount(Subaccount::from_array(buf)))
208            }
209
210            Self::TAG_TIMESTAMP => {
211                let mut buf = [0u8; Self::TIMESTAMP_SIZE];
212                buf.copy_from_slice(&payload[..Self::TIMESTAMP_SIZE]);
213                ensure_zero_padding(Self::TIMESTAMP_SIZE, "timestamp")?;
214                Ok(Self::Timestamp(Timestamp::from(u64::from_be_bytes(buf))))
215            }
216
217            Self::TAG_UINT => {
218                let mut buf = [0u8; Self::UINT_SIZE];
219                buf.copy_from_slice(&payload[..Self::UINT_SIZE]);
220                ensure_zero_padding(Self::UINT_SIZE, "uint")?;
221                Ok(Self::Uint(u64::from_be_bytes(buf)))
222            }
223
224            Self::TAG_ULID => {
225                let mut buf = [0u8; Self::ULID_SIZE];
226                buf.copy_from_slice(&payload[..Self::ULID_SIZE]);
227                ensure_zero_padding(Self::ULID_SIZE, "ulid")?;
228                Ok(Self::Ulid(Ulid::from_bytes(buf)))
229            }
230
231            Self::TAG_UNIT => {
232                ensure_zero_padding(0, "unit")?;
233                Ok(Self::Unit)
234            }
235
236            _ => Err("corrupted Key: invalid tag"),
237        }
238    }
239}
240
241impl Ord for Key {
242    fn cmp(&self, other: &Self) -> Ordering {
243        match (self, other) {
244            (Self::Account(a), Self::Account(b)) => Ord::cmp(a, b),
245            (Self::Int(a), Self::Int(b)) => Ord::cmp(a, b),
246            (Self::Principal(a), Self::Principal(b)) => Ord::cmp(a, b),
247            (Self::Uint(a), Self::Uint(b)) => Ord::cmp(a, b),
248            (Self::Ulid(a), Self::Ulid(b)) => Ord::cmp(a, b),
249            (Self::Subaccount(a), Self::Subaccount(b)) => Ord::cmp(a, b),
250            (Self::Timestamp(a), Self::Timestamp(b)) => Ord::cmp(a, b),
251
252            _ => Ord::cmp(&self.variant_rank(), &other.variant_rank()), // fallback for cross-type comparison
253        }
254    }
255}
256
257impl PartialOrd for Key {
258    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
259        Some(Ord::cmp(self, other))
260    }
261}