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#[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#[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 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 pub const STORED_SIZE_BYTES: u64 = 64;
79
80 #[expect(clippy::cast_possible_truncation)]
82 pub const STORED_SIZE_USIZE: usize = Self::STORED_SIZE_BYTES as usize;
83
84 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 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 pub fn max_storable() -> Self {
115 Self::Account(Account::max_storable())
116 }
117
118 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 pub fn to_bytes(&self) -> Result<[u8; Self::STORED_SIZE_USIZE], KeyEncodeError> {
140 let mut buf = [0u8; Self::STORED_SIZE_USIZE];
141
142 buf[Self::TAG_OFFSET] = self.tag();
144 let payload = &mut buf[Self::PAYLOAD_OFFSET..=Self::PAYLOAD_SIZE];
145
146 #[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 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()), }
314 }
315}
316
317impl PartialOrd for Key {
318 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
319 Some(Ord::cmp(self, other))
320 }
321}