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#[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 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 pub const STORED_SIZE: usize = 64;
40
41 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 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 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 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 #[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 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()), }
254 }
255}
256
257impl PartialOrd for Key {
258 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
259 Some(Ord::cmp(self, other))
260 }
261}