1use crate::registry::{DatabaseStorePrefixes, SEPARATOR};
2use smallvec::SmallVec;
3use std::fmt::{Debug, Display};
4
5#[derive(Clone)]
6pub struct DbKey {
7 path: SmallVec<[u8; 36]>, prefix_len: usize,
9}
10
11impl DbKey {
12 pub fn new<TKey>(prefix: &[u8], key: TKey) -> Self
13 where
14 TKey: Clone + AsRef<[u8]>,
15 {
16 Self { path: prefix.iter().chain(key.as_ref().iter()).copied().collect(), prefix_len: prefix.len() }
17 }
18
19 pub fn new_with_bucket<TKey, TBucket>(prefix: &[u8], bucket: TBucket, key: TKey) -> Self
20 where
21 TKey: Clone + AsRef<[u8]>,
22 TBucket: Copy + AsRef<[u8]>,
23 {
24 let mut db_key = Self::prefix_only(prefix);
25 db_key.add_bucket(bucket);
26 db_key.add_key(key);
27 db_key
28 }
29
30 pub fn new_with_suffix<TKey, TSuffix>(prefix: &[u8], key: TKey, suffix: TSuffix) -> Self
34 where
35 TKey: Clone + AsRef<[u8]>,
36 TSuffix: AsRef<[u8]>,
37 {
38 let mut db_key = Self::new(prefix, key);
39 db_key.add_suffix(suffix);
40 db_key
41 }
42
43 pub fn prefix_only(prefix: &[u8]) -> Self {
44 Self::new(prefix, [])
45 }
46
47 pub fn add_bucket<TBucket>(&mut self, bucket: TBucket)
49 where
50 TBucket: Copy + AsRef<[u8]>,
51 {
52 self.path.extend(bucket.as_ref().iter().copied());
53 self.prefix_len += bucket.as_ref().len();
54 }
55
56 pub fn add_key<TKey>(&mut self, key: TKey)
57 where
58 TKey: Clone + AsRef<[u8]>,
59 {
60 self.path.extend(key.as_ref().iter().copied());
61 self.prefix_len += key.as_ref().len();
62 }
63
64 pub fn add_suffix<TSuffix>(&mut self, suffix: TSuffix)
73 where
74 TSuffix: AsRef<[u8]>,
75 {
76 self.path.extend(suffix.as_ref().iter().copied());
77 }
78
79 pub fn prefix_len(&self) -> usize {
80 self.prefix_len
81 }
82}
83
84impl AsRef<[u8]> for DbKey {
85 fn as_ref(&self) -> &[u8] {
86 &self.path
87 }
88}
89
90impl Display for DbKey {
91 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
92 use DatabaseStorePrefixes::*;
93 let mut pos = 0;
94
95 if self.prefix_len > 0
96 && let Ok(prefix) = DatabaseStorePrefixes::try_from(self.path[0])
97 {
98 prefix.fmt(f)?;
99 f.write_str("/")?;
100 pos += 1;
101 if self.prefix_len > 1 {
102 match prefix {
103 Ghostdag
104 | GhostdagCompact
105 | TempGhostdag
106 | TempGhostdagCompact
107 | RelationsParents
108 | RelationsChildren
109 | Reachability
110 | ReachabilityTreeChildren
111 | ReachabilityFutureCoveringSet => {
112 if self.path[1] != SEPARATOR {
113 Display::fmt(&self.path[1], f)?;
115 f.write_str("/")?;
116 }
117 pos += 1;
118 }
119 ReachabilityRelations => {
120 if let Ok(next_prefix) = DatabaseStorePrefixes::try_from(self.path[1]) {
121 next_prefix.fmt(f)?;
122 f.write_str("/")?;
123 pos += 1;
124 }
125 }
126 _ => {}
127 }
128 }
129 }
130
131 f.write_str(&faster_hex::hex_string(&self.path[pos..]))
133 }
134}
135
136impl Debug for DbKey {
137 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
138 Display::fmt(&self, f)
139 }
140}
141
142#[cfg(test)]
143mod tests {
144 use super::*;
145 use DatabaseStorePrefixes::*;
146 use kaspa_hashes::{HASH_SIZE, Hash};
147
148 #[test]
149 fn test_key_display() {
150 let level = 37;
151 let key1 = DbKey::new(&[ReachabilityRelations.into(), RelationsParents.into()], Hash::from_u64_word(34567890));
152 let key2 = DbKey::new(&[Reachability.into(), Separator.into()], Hash::from_u64_word(345690));
153 let key3 = DbKey::new(&[Reachability.into(), level], Hash::from_u64_word(345690));
154 let key4 = DbKey::new(&[RelationsParents.into(), level], Hash::from_u64_word(345690));
155
156 assert!(key1.to_string().starts_with(&format!("{:?}/{:?}/00", ReachabilityRelations, RelationsParents)));
157 assert!(key2.to_string().starts_with(&format!("{:?}/00", Reachability)));
158 assert!(key3.to_string().starts_with(&format!("{:?}/{}/00", Reachability, level)));
159 assert!(key4.to_string().starts_with(&format!("{:?}/{}/00", RelationsParents, level)));
160
161 let key5 = DbKey::new(b"human/readable", Hash::from_bytes([SEPARATOR; HASH_SIZE]));
162 let key6 = DbKey::prefix_only(&[0xC0, 0xC1, 0xF5, 0xF6]);
163 let key7 = DbKey::prefix_only(b"direct-prefix");
164
165 let _ = key5.to_string();
167 let _ = key6.to_string();
168 let _ = key7.to_string();
169 }
170
171 #[test]
172 fn test_add_suffix_and_new_with_suffix() {
173 let prefix: [u8; 1] = [42];
176 let key = [0x01u8; 4];
177 let suffix = [0xFFu8];
178 let db_key = DbKey::new_with_suffix(&prefix, key, suffix);
179 assert_eq!(db_key.as_ref(), &[42, 0x01, 0x01, 0x01, 0x01, 0xFF]);
180 assert_eq!(db_key.prefix_len(), prefix.len(), "suffix must not grow prefix_len");
181
182 let mut extended = DbKey::new(&prefix, key);
184 let prefix_len_before = extended.prefix_len();
185 extended.add_suffix([0x77u8, 0x88u8]);
186 assert_eq!(extended.as_ref(), &[42, 0x01, 0x01, 0x01, 0x01, 0x77, 0x88]);
187 assert_eq!(extended.prefix_len(), prefix_len_before);
188 }
189}