1#![expect(clippy::cast_possible_truncation)]
10
11use super::Value;
12use crate::{
13 error::InternalError,
14 traits::Repr,
15 types::{Account, Principal, Subaccount, Timestamp, Ulid},
16};
17use candid::CandidType;
18use serde::Deserialize;
19use std::cmp::Ordering;
20use thiserror::Error as ThisError;
21
22#[derive(Debug, ThisError)]
28pub enum StorageKeyEncodeError {
29 #[error("account owner principal exceeds max length: {len} bytes (limit {max})")]
30 AccountOwnerTooLarge { len: usize, max: usize },
31
32 #[error("value kind '{kind}' is not storage-key encodable")]
33 UnsupportedValueKind { kind: &'static str },
34
35 #[error("principal exceeds max length: {len} bytes (limit {max})")]
36 PrincipalTooLarge { len: usize, max: usize },
37}
38
39impl From<StorageKeyEncodeError> for InternalError {
40 fn from(err: StorageKeyEncodeError) -> Self {
41 Self::serialize_unsupported(err.to_string())
42 }
43}
44
45#[derive(Debug, ThisError)]
51pub enum StorageKeyDecodeError {
52 #[error("corrupted StorageKey: invalid size")]
53 InvalidSize,
54
55 #[error("corrupted StorageKey: invalid tag")]
56 InvalidTag,
57
58 #[error("corrupted StorageKey: invalid principal length")]
59 InvalidPrincipalLength,
60
61 #[error("corrupted StorageKey: non-zero {field} padding")]
62 NonZeroPadding { field: &'static str },
63
64 #[error("corrupted StorageKey: invalid account payload ({reason})")]
65 InvalidAccountPayload { reason: &'static str },
66}
67
68#[derive(CandidType, Clone, Copy, Debug, Deserialize, Eq, Hash, PartialEq)]
78pub enum StorageKey {
79 Account(Account),
80 Int(i64),
81 Principal(Principal),
82 Subaccount(Subaccount),
83 Timestamp(Timestamp),
84 Uint(u64),
85 Ulid(Ulid),
86 Unit,
87}
88
89macro_rules! value_is_storage_key_encodable_from_registry {
91 ( @args $value:expr; @entries $( ($scalar:ident, $family:expr, $value_pat:pat, is_numeric_value = $is_numeric:expr, supports_numeric_coercion = $supports_numeric_coercion:expr, supports_arithmetic = $supports_arithmetic:expr, supports_equality = $supports_equality:expr, supports_ordering = $supports_ordering:expr, is_keyable = $is_keyable:expr, is_storage_key_encodable = $is_storage_key_encodable:expr) ),* $(,)? ) => {
92 match $value {
93 $( $value_pat => $is_storage_key_encodable, )*
94 _ => false,
95 }
96 };
97}
98
99impl StorageKey {
100 pub(crate) const TAG_ACCOUNT: u8 = 0;
102 pub(crate) const TAG_INT: u8 = 1;
103 pub(crate) const TAG_PRINCIPAL: u8 = 2;
104 pub(crate) const TAG_SUBACCOUNT: u8 = 3;
105 pub(crate) const TAG_TIMESTAMP: u8 = 4;
106 pub(crate) const TAG_UINT: u8 = 5;
107 pub(crate) const TAG_ULID: u8 = 6;
108 pub(crate) const TAG_UNIT: u8 = 7;
109
110 pub const STORED_SIZE_BYTES: u64 = 64;
113 pub const STORED_SIZE_USIZE: usize = Self::STORED_SIZE_BYTES as usize;
114
115 const TAG_SIZE: usize = 1;
116 pub(crate) const TAG_OFFSET: usize = 0;
117
118 pub(crate) const PAYLOAD_OFFSET: usize = Self::TAG_SIZE;
119 const PAYLOAD_SIZE: usize = Self::STORED_SIZE_USIZE - Self::TAG_SIZE;
120
121 pub(crate) const INT_SIZE: usize = 8;
122 pub(crate) const UINT_SIZE: usize = 8;
123 pub(crate) const TIMESTAMP_SIZE: usize = 8;
124 pub(crate) const ULID_SIZE: usize = 16;
125 pub(crate) const SUBACCOUNT_SIZE: usize = 32;
126 const ACCOUNT_MAX_SIZE: usize = 62;
127
128 pub const fn try_from_value(value: &Value) -> Result<Self, StorageKeyEncodeError> {
129 let is_storage_key_encodable =
133 scalar_registry!(value_is_storage_key_encodable_from_registry, value);
134 if !is_storage_key_encodable {
135 return Err(StorageKeyEncodeError::UnsupportedValueKind {
136 kind: Self::value_kind_label(value),
137 });
138 }
139
140 match value {
141 Value::Account(v) => Ok(Self::Account(*v)),
142 Value::Int(v) => Ok(Self::Int(*v)),
143 Value::Principal(v) => Ok(Self::Principal(*v)),
144 Value::Subaccount(v) => Ok(Self::Subaccount(*v)),
145 Value::Timestamp(v) => Ok(Self::Timestamp(*v)),
146 Value::Uint(v) => Ok(Self::Uint(*v)),
147 Value::Ulid(v) => Ok(Self::Ulid(*v)),
148 Value::Unit => Ok(Self::Unit),
149
150 _ => Err(StorageKeyEncodeError::UnsupportedValueKind {
151 kind: Self::value_kind_label(value),
152 }),
153 }
154 }
155
156 const fn value_kind_label(value: &Value) -> &'static str {
157 match value {
158 Value::Account(_) => "Account",
159 Value::Blob(_) => "Blob",
160 Value::Bool(_) => "Bool",
161 Value::Date(_) => "Date",
162 Value::Decimal(_) => "Decimal",
163 Value::Duration(_) => "Duration",
164 Value::Enum(_) => "Enum",
165 Value::Float32(_) => "Float32",
166 Value::Float64(_) => "Float64",
167 Value::Int(_) => "Int",
168 Value::Int128(_) => "Int128",
169 Value::IntBig(_) => "IntBig",
170 Value::List(_) => "List",
171 Value::Map(_) => "Map",
172 Value::Null => "Null",
173 Value::Principal(_) => "Principal",
174 Value::Subaccount(_) => "Subaccount",
175 Value::Text(_) => "Text",
176 Value::Timestamp(_) => "Timestamp",
177 Value::Uint(_) => "Uint",
178 Value::Uint128(_) => "Uint128",
179 Value::UintBig(_) => "UintBig",
180 Value::Ulid(_) => "Ulid",
181 Value::Unit => "Unit",
182 }
183 }
184
185 const fn tag(&self) -> u8 {
186 match self {
187 Self::Account(_) => Self::TAG_ACCOUNT,
188 Self::Int(_) => Self::TAG_INT,
189 Self::Principal(_) => Self::TAG_PRINCIPAL,
190 Self::Subaccount(_) => Self::TAG_SUBACCOUNT,
191 Self::Timestamp(_) => Self::TAG_TIMESTAMP,
192 Self::Uint(_) => Self::TAG_UINT,
193 Self::Ulid(_) => Self::TAG_ULID,
194 Self::Unit => Self::TAG_UNIT,
195 }
196 }
197
198 #[must_use]
200 pub fn max_storable() -> Self {
201 Self::Account(Account::max_storable())
202 }
203
204 pub const MIN: Self = Self::Account(Account::from_parts(Principal::from_slice(&[]), None));
206
207 #[must_use]
208 pub const fn lower_bound() -> Self {
209 Self::MIN
210 }
211
212 #[must_use]
213 pub const fn upper_bound() -> Self {
214 Self::Unit
215 }
216
217 const fn variant_rank(&self) -> u8 {
218 self.tag()
219 }
220
221 const fn from_account_encode_error(
222 err: crate::types::AccountEncodeError,
223 ) -> StorageKeyEncodeError {
224 match err {
225 crate::types::AccountEncodeError::OwnerEncode(inner) => {
226 Self::from_principal_encode_error(inner)
227 }
228 crate::types::AccountEncodeError::OwnerTooLarge { len, max } => {
229 StorageKeyEncodeError::AccountOwnerTooLarge { len, max }
230 }
231 }
232 }
233
234 const fn from_principal_encode_error(
235 err: crate::types::PrincipalEncodeError,
236 ) -> StorageKeyEncodeError {
237 match err {
238 crate::types::PrincipalEncodeError::TooLarge { len, max } => {
239 StorageKeyEncodeError::PrincipalTooLarge { len, max }
240 }
241 }
242 }
243
244 pub fn to_bytes(self) -> Result<[u8; Self::STORED_SIZE_USIZE], StorageKeyEncodeError> {
246 let mut buf = [0u8; Self::STORED_SIZE_USIZE];
248 buf[Self::TAG_OFFSET] = self.tag();
249 let payload = &mut buf[Self::PAYLOAD_OFFSET..=Self::PAYLOAD_SIZE];
250
251 match self {
253 Self::Account(v) => {
254 let bytes = v
255 .to_stored_bytes()
256 .map_err(Self::from_account_encode_error)?;
257 payload[..Self::ACCOUNT_MAX_SIZE].copy_from_slice(&bytes);
258 }
259 Self::Int(v) => {
260 let biased = v.cast_unsigned() ^ (1u64 << 63);
261 payload[..Self::INT_SIZE].copy_from_slice(&biased.to_be_bytes());
262 }
263 Self::Uint(v) => payload[..Self::UINT_SIZE].copy_from_slice(&v.to_be_bytes()),
264 Self::Timestamp(v) => {
265 payload[..Self::TIMESTAMP_SIZE].copy_from_slice(&v.repr().to_be_bytes());
266 }
267 Self::Principal(v) => {
268 let bytes = v
269 .stored_bytes()
270 .map_err(Self::from_principal_encode_error)?;
271 let len = bytes.len();
272 payload[0] =
273 u8::try_from(len).map_err(|_| StorageKeyEncodeError::PrincipalTooLarge {
274 len,
275 max: Principal::MAX_LENGTH_IN_BYTES as usize,
276 })?;
277 payload[1..=len].copy_from_slice(bytes);
278 }
279 Self::Subaccount(v) => payload[..Self::SUBACCOUNT_SIZE].copy_from_slice(&v.to_array()),
280 Self::Ulid(v) => payload[..Self::ULID_SIZE].copy_from_slice(&v.to_bytes()),
281 Self::Unit => {}
282 }
283
284 Ok(buf)
285 }
286
287 pub fn try_from_bytes(bytes: &[u8]) -> Result<Self, StorageKeyDecodeError> {
288 let bytes: &[u8; Self::STORED_SIZE_USIZE] = bytes
289 .try_into()
290 .map_err(|_| StorageKeyDecodeError::InvalidSize)?;
291
292 Self::try_from_stored_bytes(bytes)
293 }
294
295 pub(crate) fn try_from_stored_bytes(
297 bytes: &[u8; Self::STORED_SIZE_USIZE],
298 ) -> Result<Self, StorageKeyDecodeError> {
299 let tag = bytes[Self::TAG_OFFSET];
300 let payload = &bytes[Self::PAYLOAD_OFFSET..=Self::PAYLOAD_SIZE];
301
302 let ensure_zero_padding = |used: usize, ctx: &'static str| {
303 if payload[used..].iter().all(|&b| b == 0) {
304 Ok(())
305 } else {
306 Err(StorageKeyDecodeError::NonZeroPadding { field: ctx })
307 }
308 };
309
310 match tag {
312 Self::TAG_ACCOUNT => {
313 let end = Account::STORED_SIZE as usize;
314 ensure_zero_padding(end, "account")?;
315 Ok(Self::Account(
316 Account::try_from_bytes(&payload[..end]).map_err(|reason| {
317 StorageKeyDecodeError::InvalidAccountPayload { reason }
318 })?,
319 ))
320 }
321 Self::TAG_INT => {
322 let mut buf = [0u8; Self::INT_SIZE];
323 buf.copy_from_slice(&payload[..Self::INT_SIZE]);
324 ensure_zero_padding(Self::INT_SIZE, "int")?;
325 Ok(Self::Int(
326 (u64::from_be_bytes(buf) ^ (1u64 << 63)).cast_signed(),
327 ))
328 }
329 Self::TAG_PRINCIPAL => {
330 let len = payload[0] as usize;
331 if len > Principal::MAX_LENGTH_IN_BYTES as usize {
332 return Err(StorageKeyDecodeError::InvalidPrincipalLength);
333 }
334 ensure_zero_padding(1 + len, "principal")?;
335 Ok(Self::Principal(Principal::from_slice(&payload[1..=len])))
336 }
337 Self::TAG_SUBACCOUNT => {
338 ensure_zero_padding(Self::SUBACCOUNT_SIZE, "subaccount")?;
339 let mut buf = [0u8; Self::SUBACCOUNT_SIZE];
340 buf.copy_from_slice(&payload[..Self::SUBACCOUNT_SIZE]);
341 Ok(Self::Subaccount(Subaccount::from_array(buf)))
342 }
343 Self::TAG_TIMESTAMP => {
344 ensure_zero_padding(Self::TIMESTAMP_SIZE, "timestamp")?;
345 let mut buf = [0u8; Self::TIMESTAMP_SIZE];
346 buf.copy_from_slice(&payload[..Self::TIMESTAMP_SIZE]);
347 Ok(Self::Timestamp(Timestamp::from_repr(i64::from_be_bytes(
348 buf,
349 ))))
350 }
351 Self::TAG_UINT => {
352 ensure_zero_padding(Self::UINT_SIZE, "uint")?;
353 let mut buf = [0u8; Self::UINT_SIZE];
354 buf.copy_from_slice(&payload[..Self::UINT_SIZE]);
355 Ok(Self::Uint(u64::from_be_bytes(buf)))
356 }
357 Self::TAG_ULID => {
358 ensure_zero_padding(Self::ULID_SIZE, "ulid")?;
359 let mut buf = [0u8; Self::ULID_SIZE];
360 buf.copy_from_slice(&payload[..Self::ULID_SIZE]);
361 Ok(Self::Ulid(Ulid::from_bytes(buf)))
362 }
363 Self::TAG_UNIT => {
364 ensure_zero_padding(0, "unit")?;
365 Ok(Self::Unit)
366 }
367 _ => Err(StorageKeyDecodeError::InvalidTag),
368 }
369 }
370
371 #[must_use]
376 pub const fn as_value(&self) -> Value {
377 match self {
378 Self::Account(v) => Value::Account(*v),
379 Self::Int(v) => Value::Int(*v),
380 Self::Principal(v) => Value::Principal(*v),
381 Self::Subaccount(v) => Value::Subaccount(*v),
382 Self::Timestamp(v) => Value::Timestamp(*v),
383 Self::Uint(v) => Value::Uint(*v),
384 Self::Ulid(v) => Value::Ulid(*v),
385 Self::Unit => Value::Unit,
386 }
387 }
388}
389
390impl Ord for StorageKey {
391 fn cmp(&self, other: &Self) -> Ordering {
392 match (self, other) {
393 (Self::Account(a), Self::Account(b)) => a.cmp(b),
394 (Self::Int(a), Self::Int(b)) => a.cmp(b),
395 (Self::Principal(a), Self::Principal(b)) => a.cmp(b),
396 (Self::Uint(a), Self::Uint(b)) => a.cmp(b),
397 (Self::Ulid(a), Self::Ulid(b)) => a.cmp(b),
398 (Self::Subaccount(a), Self::Subaccount(b)) => a.cmp(b),
399 (Self::Timestamp(a), Self::Timestamp(b)) => a.cmp(b),
400 _ => self.variant_rank().cmp(&other.variant_rank()),
401 }
402 }
403}
404
405impl PartialOrd for StorageKey {
406 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
407 Some(self.cmp(other))
408 }
409}
410
411impl TryFrom<&[u8]> for StorageKey {
412 type Error = StorageKeyDecodeError;
413 fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
414 Self::try_from_bytes(bytes)
415 }
416}
417
418#[cfg(test)]
423mod tests {
424 use super::{StorageKey, StorageKeyDecodeError, StorageKeyEncodeError};
425 use crate::{
426 types::{
427 Account, Date, Decimal, Duration, Float32, Float64, Int, Int128, Nat, Nat128,
428 Principal, Subaccount, Timestamp, Ulid,
429 },
430 value::{Value, ValueEnum},
431 };
432
433 macro_rules! sample_value_for_scalar {
434 (Account) => {
435 Value::Account(Account::dummy(7))
436 };
437 (Blob) => {
438 Value::Blob(vec![1u8, 2u8, 3u8])
439 };
440 (Bool) => {
441 Value::Bool(true)
442 };
443 (Date) => {
444 Value::Date(Date::new(2024, 1, 2))
445 };
446 (Decimal) => {
447 Value::Decimal(Decimal::new(123, 2))
448 };
449 (Duration) => {
450 Value::Duration(Duration::from_secs(1))
451 };
452 (Enum) => {
453 Value::Enum(ValueEnum::loose("example"))
454 };
455 (Float32) => {
456 Value::Float32(Float32::try_new(1.25).expect("Float32 sample should be finite"))
457 };
458 (Float64) => {
459 Value::Float64(Float64::try_new(2.5).expect("Float64 sample should be finite"))
460 };
461 (Int) => {
462 Value::Int(-7)
463 };
464 (Int128) => {
465 Value::Int128(Int128::from(123i128))
466 };
467 (IntBig) => {
468 Value::IntBig(Int::from(99i32))
469 };
470 (Principal) => {
471 Value::Principal(Principal::from_slice(&[1u8, 2u8, 3u8]))
472 };
473 (Subaccount) => {
474 Value::Subaccount(Subaccount::new([1u8; 32]))
475 };
476 (Text) => {
477 Value::Text("example".to_string())
478 };
479 (Timestamp) => {
480 Value::Timestamp(Timestamp::from_secs(1))
481 };
482 (Uint) => {
483 Value::Uint(7)
484 };
485 (Uint128) => {
486 Value::Uint128(Nat128::from(9u128))
487 };
488 (UintBig) => {
489 Value::UintBig(Nat::from(11u64))
490 };
491 (Ulid) => {
492 Value::Ulid(Ulid::from_u128(42))
493 };
494 (Unit) => {
495 Value::Unit
496 };
497 }
498
499 fn registry_storage_encodable_cases() -> Vec<(Value, bool)> {
500 macro_rules! collect_cases {
501 ( @entries $( ($scalar:ident, $family:expr, $value_pat:pat, is_numeric_value = $is_numeric:expr, supports_numeric_coercion = $supports_numeric_coercion:expr, supports_arithmetic = $supports_arithmetic:expr, supports_equality = $supports_equality:expr, supports_ordering = $supports_ordering:expr, is_keyable = $is_keyable:expr, is_storage_key_encodable = $is_storage_key_encodable:expr) ),* $(,)? ) => {
502 vec![ $( (sample_value_for_scalar!($scalar), $is_storage_key_encodable) ),* ]
503 };
504 ( @args $($ignore:tt)*; @entries $( ($scalar:ident, $family:expr, $value_pat:pat, is_numeric_value = $is_numeric:expr, supports_numeric_coercion = $supports_numeric_coercion:expr, supports_arithmetic = $supports_arithmetic:expr, supports_equality = $supports_equality:expr, supports_ordering = $supports_ordering:expr, is_keyable = $is_keyable:expr, is_storage_key_encodable = $is_storage_key_encodable:expr) ),* $(,)? ) => {
505 vec![ $( (sample_value_for_scalar!($scalar), $is_storage_key_encodable) ),* ]
506 };
507 }
508
509 scalar_registry!(collect_cases)
510 }
511
512 #[test]
513 fn storage_key_try_from_value_matches_registry_flag() {
514 for (value, expected_encodable) in registry_storage_encodable_cases() {
515 assert_eq!(
516 StorageKey::try_from_value(&value).is_ok(),
517 expected_encodable,
518 "value: {value:?}"
519 );
520 }
521 }
522
523 #[test]
524 fn storage_key_known_encodability_contracts() {
525 assert!(StorageKey::try_from_value(&Value::Unit).is_ok());
526 assert!(StorageKey::try_from_value(&Value::Decimal(Decimal::new(1, 0))).is_err());
527 assert!(StorageKey::try_from_value(&Value::Text("x".to_string())).is_err());
528 assert!(StorageKey::try_from_value(&Value::Account(Account::dummy(1))).is_ok());
529 }
530
531 #[test]
532 fn storage_key_unsupported_values_report_kind() {
533 let decimal_err = StorageKey::try_from_value(&Value::Decimal(Decimal::new(1, 0)))
534 .expect_err("Decimal is not storage-key encodable");
535 assert!(matches!(
536 decimal_err,
537 StorageKeyEncodeError::UnsupportedValueKind { kind } if kind == "Decimal"
538 ));
539
540 let text_err = StorageKey::try_from_value(&Value::Text("x".to_string()))
541 .expect_err("Text is not storage-key encodable");
542 assert!(matches!(
543 text_err,
544 StorageKeyEncodeError::UnsupportedValueKind { kind } if kind == "Text"
545 ));
546 }
547
548 #[test]
549 fn storage_keys_sort_deterministically_across_mixed_variants() {
550 let mut keys = vec![
551 StorageKey::try_from_value(&Value::Unit).expect("Unit is encodable"),
552 StorageKey::try_from_value(&Value::Ulid(Ulid::from_u128(2)))
553 .expect("Ulid is encodable"),
554 StorageKey::try_from_value(&Value::Uint(2)).expect("Uint is encodable"),
555 StorageKey::try_from_value(&Value::Timestamp(Timestamp::from_secs(2)))
556 .expect("Timestamp is encodable"),
557 StorageKey::try_from_value(&Value::Subaccount(Subaccount::new([3u8; 32])))
558 .expect("Subaccount is encodable"),
559 StorageKey::try_from_value(&Value::Principal(Principal::from_slice(&[9u8])))
560 .expect("Principal is encodable"),
561 StorageKey::try_from_value(&Value::Int(-1)).expect("Int is encodable"),
562 StorageKey::try_from_value(&Value::Account(Account::dummy(3)))
563 .expect("Account is encodable"),
564 ];
565
566 keys.sort();
567
568 let expected = vec![
569 StorageKey::Account(Account::dummy(3)),
570 StorageKey::Int(-1),
571 StorageKey::Principal(Principal::from_slice(&[9u8])),
572 StorageKey::Subaccount(Subaccount::new([3u8; 32])),
573 StorageKey::Timestamp(Timestamp::from_secs(2)),
574 StorageKey::Uint(2),
575 StorageKey::Ulid(Ulid::from_u128(2)),
576 StorageKey::Unit,
577 ];
578
579 assert_eq!(keys, expected);
580 }
581
582 #[test]
583 fn storage_key_decode_rejects_invalid_size_as_structured_error() {
584 let err =
585 StorageKey::try_from_bytes(&[]).expect_err("decode should reject invalid key size");
586 assert!(matches!(err, StorageKeyDecodeError::InvalidSize));
587 }
588
589 #[test]
590 fn storage_key_decode_rejects_invalid_tag_as_structured_error() {
591 let mut bytes = [0u8; StorageKey::STORED_SIZE_USIZE];
592 bytes[StorageKey::TAG_OFFSET] = 0xFF;
593
594 let err = StorageKey::try_from_bytes(&bytes).expect_err("decode should reject invalid tag");
595 assert!(matches!(err, StorageKeyDecodeError::InvalidTag));
596 }
597
598 #[test]
599 fn storage_key_decode_rejects_non_zero_padding_with_segment_context() {
600 let mut bytes = [0u8; StorageKey::STORED_SIZE_USIZE];
601 bytes[StorageKey::TAG_OFFSET] = StorageKey::TAG_UNIT;
602 bytes[StorageKey::PAYLOAD_OFFSET] = 1;
603
604 let err = StorageKey::try_from_bytes(&bytes)
605 .expect_err("decode should reject non-zero padding for unit payload");
606 assert!(matches!(
607 err,
608 StorageKeyDecodeError::NonZeroPadding { field } if field == "unit"
609 ));
610 }
611}