hannoy/
key.rs

1use std::borrow::Cow;
2use std::mem::size_of;
3
4use byteorder::{BigEndian, ByteOrder};
5use heed::BoxedError;
6
7use crate::{NodeId, NodeMode};
8
9/// This whole structure must fit in an u64 so we can tell LMDB to optimize its storage.
10/// The `index` is specified by the user and is used to differentiate between multiple hannoy indexes.
11/// The `mode` indicates what we're looking at.
12/// The `item` point to a specific node.
13/// If the mode is:
14///  - `Item`: we're looking at an `Item` node.
15///  - `Links`: we're looking at the `Links` bitmap of neighbours for a node
16///  - `Updated`: The list of items that has been updated since the last build of the database.
17///  - `Metadata`: There is only one item at `0` that contains the header required to read the index.
18#[derive(Debug, Copy, Clone)]
19pub struct Key {
20    /// The prefix specified by the user.
21    pub index: u16,
22    pub node: NodeId,
23}
24
25impl Key {
26    pub const fn new(index: u16, node: NodeId) -> Self {
27        Self { index, node }
28    }
29
30    pub const fn metadata(index: u16) -> Self {
31        Self::new(index, NodeId::metadata())
32    }
33
34    pub const fn version(index: u16) -> Self {
35        Self::new(index, NodeId::version())
36    }
37
38    pub const fn updated(index: u16, item: u32) -> Self {
39        Self::new(index, NodeId::updated(item))
40    }
41
42    pub const fn item(index: u16, item: u32) -> Self {
43        Self::new(index, NodeId::item(item))
44    }
45
46    pub const fn links(index: u16, item: u32, layer: u8) -> Self {
47        Self::new(index, NodeId::links(item, layer))
48    }
49}
50
51/// The heed codec used internally to encode/decoding the internal key type.
52pub enum KeyCodec {}
53
54impl<'a> heed::BytesEncode<'a> for KeyCodec {
55    type EItem = Key;
56
57    fn bytes_encode(item: &'a Self::EItem) -> Result<Cow<'a, [u8]>, BoxedError> {
58        let mut output = Vec::with_capacity(size_of::<u64>());
59        output.extend_from_slice(&item.index.to_be_bytes());
60        output.extend_from_slice(&(item.node.mode as u8).to_be_bytes());
61        output.extend_from_slice(&item.node.item.to_be_bytes());
62        output.extend_from_slice(&(item.node.layer).to_be_bytes());
63
64        Ok(Cow::Owned(output))
65    }
66}
67
68impl heed::BytesDecode<'_> for KeyCodec {
69    type DItem = Key;
70
71    fn bytes_decode(bytes: &[u8]) -> Result<Self::DItem, BoxedError> {
72        let prefix = BigEndian::read_u16(bytes);
73        let bytes = &bytes[size_of::<u16>()..];
74        let mode = bytes[0].try_into()?;
75        let bytes = &bytes[size_of::<u8>()..];
76        let item = BigEndian::read_u32(bytes);
77        let bytes = &bytes[size_of::<u32>()..];
78        let layer = bytes[0];
79
80        Ok(Key { index: prefix, node: NodeId { mode, item, layer } })
81    }
82}
83
84/// This is used to query part of a key.
85#[derive(Debug, Copy, Clone)]
86pub struct Prefix {
87    /// The index specified by the user.
88    index: u16,
89    // Indicate what the item represent.
90    mode: Option<NodeMode>,
91}
92
93impl Prefix {
94    pub const fn all(index: u16) -> Self {
95        Self { index, mode: None }
96    }
97
98    pub const fn item(index: u16) -> Self {
99        Self { index, mode: Some(NodeMode::Item) }
100    }
101
102    pub const fn links(index: u16) -> Self {
103        Self { index, mode: Some(NodeMode::Links) }
104    }
105
106    pub const fn updated(index: u16) -> Self {
107        Self { index, mode: Some(NodeMode::Updated) }
108    }
109}
110
111pub enum PrefixCodec {}
112
113impl<'a> heed::BytesEncode<'a> for PrefixCodec {
114    type EItem = Prefix;
115
116    fn bytes_encode(item: &'a Self::EItem) -> Result<Cow<'a, [u8]>, BoxedError> {
117        let mode_used = item.mode.is_some() as usize;
118        let mut output = Vec::with_capacity(size_of::<u16>() + mode_used);
119
120        output.extend_from_slice(&item.index.to_be_bytes());
121        if let Some(mode) = item.mode {
122            output.extend_from_slice(&(mode as u8).to_be_bytes());
123        }
124
125        Ok(Cow::Owned(output))
126    }
127}
128
129#[cfg(test)]
130mod test {
131    use heed::{BytesDecode, BytesEncode};
132
133    use super::*;
134
135    #[test]
136    fn check_size_of_types() {
137        let key = Key::metadata(0);
138        let encoded = KeyCodec::bytes_encode(&key).unwrap();
139        assert_eq!(encoded.len(), size_of::<u64>());
140    }
141
142    // TODO: fuzz this
143    #[test]
144    fn test_links_key() {
145        let key = Key::links(0, 1, 42);
146        let bytes = KeyCodec::bytes_encode(&key).unwrap();
147        let key2 = KeyCodec::bytes_decode(&bytes).unwrap();
148        assert_eq!(key.node.item, key2.node.item);
149        assert_eq!(key.node.layer, key2.node.layer);
150        assert_eq!(key.node.mode, key2.node.mode);
151    }
152
153    #[test]
154    fn test_item_key() {
155        let key = Key::item(0, 42);
156        let bytes = KeyCodec::bytes_encode(&key).unwrap();
157        let key2 = KeyCodec::bytes_decode(&bytes).unwrap();
158        assert_eq!(key.node.item, key2.node.item);
159        assert_eq!(key.node.layer, key2.node.layer);
160        assert_eq!(key.node.mode, key2.node.mode);
161    }
162}
hannoy/key.rs

hannoy/
key.rs
