Skip to main content

kaspa_database/
key.rs

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]>, // Optimized for the common case of { prefix byte || HASH (32 bytes) }
8    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    /// Convenience constructor: `[prefix || key || suffix]`. The suffix bytes
31    /// sit *after* the logical key and are not counted toward `prefix_len` —
32    /// see [`DbKey::add_suffix`] for the rationale.
33    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    /// add a bucket to the DBkey, this adds to the prefix length
48    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    /// Append bytes after the logical key without changing `prefix_len`.
65    ///
66    /// Unlike [`DbKey::add_bucket`] and [`DbKey::add_key`] — both of which sit
67    /// conceptually in front of the row key and therefore grow `prefix_len` —
68    /// a suffix is part of the row key itself (for example a version byte
69    /// appended to a versioned store's keys). `prefix_len` must continue to
70    /// reflect only the store/bucket path so that iterator helpers that strip
71    /// the prefix see the correct boundary.
72    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                            // Expected to be a block level so we display as a number
114                            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        // We expect that the key part is usually more readable as hex
132        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        // Make sure display can handle arbitrary strings
166        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        // `new_with_suffix` produces [prefix || key || suffix] and reports the
174        // store-prefix length (without the suffix) via `prefix_len()`.
175        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        // `add_suffix` on an existing key appends bytes and keeps prefix_len untouched.
183        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}