1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

use crate::registry::{DatabaseStorePrefixes, SEPARATOR};
use smallvec::SmallVec;
use std::fmt::{Debug, Display};

#[derive(Clone)]
pub struct DbKey {
    path: SmallVec<[u8; 36]>, // Optimized for the common case of { prefix byte || HASH (32 bytes) }
    prefix_len: usize,
}

impl DbKey {
    pub fn new<TKey>(prefix: &[u8], key: TKey) -> Self
    where
        TKey: Clone + AsRef<[u8]>,
    {
        Self { path: prefix.iter().chain(key.as_ref().iter()).copied().collect(), prefix_len: prefix.len() }
    }

    pub fn new_with_bucket<TKey, TBucket>(prefix: &[u8], bucket: TBucket, key: TKey) -> Self
    where
        TKey: Clone + AsRef<[u8]>,
        TBucket: Copy + AsRef<[u8]>,
    {
        let mut db_key = Self::prefix_only(prefix);
        db_key.add_bucket(bucket);
        db_key.add_key(key);
        db_key
    }

    pub fn prefix_only(prefix: &[u8]) -> Self {
        Self::new(prefix, [])
    }

    /// add a bucket to the DBkey, this adds to the prefix length
    pub fn add_bucket<TBucket>(&mut self, bucket: TBucket)
    where
        TBucket: Copy + AsRef<[u8]>,
    {
        self.path.extend(bucket.as_ref().iter().copied());
        self.prefix_len += bucket.as_ref().len();
    }

    pub fn add_key<TKey>(&mut self, key: TKey)
    where
        TKey: Clone + AsRef<[u8]>,
    {
        self.path.extend(key.as_ref().iter().copied());
        self.prefix_len += key.as_ref().len();
    }

    pub fn prefix_len(&self) -> usize {
        self.prefix_len
    }
}

impl AsRef<[u8]> for DbKey {
    fn as_ref(&self) -> &[u8] {
        &self.path
    }
}

impl Display for DbKey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        use DatabaseStorePrefixes::*;
        let mut pos = 0;

        if self.prefix_len > 0 {
            if let Ok(prefix) = DatabaseStorePrefixes::try_from(self.path[0]) {
                prefix.fmt(f)?;
                f.write_str("/")?;
                pos += 1;
                if self.prefix_len > 1 {
                    match prefix {
                        Ghostdag
                        | GhostdagCompact
                        | RelationsParents
                        | RelationsChildren
                        | Reachability
                        | ReachabilityTreeChildren
                        | ReachabilityFutureCoveringSet => {
                            if self.path[1] != SEPARATOR {
                                // Expected to be a block level so we display as a number
                                Display::fmt(&self.path[1], f)?;
                                f.write_str("/")?;
                            }
                            pos += 1;
                        }
                        ReachabilityRelations => {
                            if let Ok(next_prefix) = DatabaseStorePrefixes::try_from(self.path[1]) {
                                next_prefix.fmt(f)?;
                                f.write_str("/")?;
                                pos += 1;
                            }
                        }
                        _ => {}
                    }
                }
            }
        }

        // We expect that the key part is usually more readable as hex
        f.write_str(&faster_hex::hex_string(&self.path[pos..]))
    }
}

impl Debug for DbKey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        Display::fmt(&self, f)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use kaspa_hashes::{Hash, HASH_SIZE};
    use DatabaseStorePrefixes::*;

    #[test]
    fn test_key_display() {
        let level = 37;
        let key1 = DbKey::new(&[ReachabilityRelations.into(), RelationsParents.into()], Hash::from_u64_word(34567890));
        let key2 = DbKey::new(&[Reachability.into(), Separator.into()], Hash::from_u64_word(345690));
        let key3 = DbKey::new(&[Reachability.into(), level], Hash::from_u64_word(345690));
        let key4 = DbKey::new(&[RelationsParents.into(), level], Hash::from_u64_word(345690));

        assert!(key1.to_string().starts_with(&format!("{:?}/{:?}/00", ReachabilityRelations, RelationsParents)));
        assert!(key2.to_string().starts_with(&format!("{:?}/00", Reachability)));
        assert!(key3.to_string().starts_with(&format!("{:?}/{}/00", Reachability, level)));
        assert!(key4.to_string().starts_with(&format!("{:?}/{}/00", RelationsParents, level)));

        let key5 = DbKey::new(b"human/readable", Hash::from_bytes([SEPARATOR; HASH_SIZE]));
        let key6 = DbKey::prefix_only(&[0xC0, 0xC1, 0xF5, 0xF6]);
        let key7 = DbKey::prefix_only(b"direct-prefix");

        // Make sure display can handle arbitrary strings
        let _ = key5.to_string();
        let _ = key6.to_string();
        let _ = key7.to_string();
    }
}