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