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